Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

106 righe
2.6 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/common"
  5. "github.com/glebarez/sqlite"
  6. "github.com/stretchr/testify/assert"
  7. "github.com/stretchr/testify/require"
  8. "gorm.io/gorm"
  9. )
  10. func setupUserChannelRatioDB(t *testing.T) *gorm.DB {
  11. t.Helper()
  12. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  13. require.NoError(t, err)
  14. sqlDB, _ := db.DB()
  15. sqlDB.SetMaxOpenConns(1)
  16. origDB := DB
  17. DB = db
  18. common.UsingSQLite = true
  19. common.RedisEnabled = false
  20. require.NoError(t, db.AutoMigrate(&UserChannelRatio{}))
  21. t.Cleanup(func() {
  22. DB = origDB
  23. sqlDB.Close()
  24. })
  25. return db
  26. }
  27. func TestUserChannelRatioCRUD(t *testing.T) {
  28. setupUserChannelRatioDB(t)
  29. LoadUserChannelRatioCache()
  30. ucr := &UserChannelRatio{
  31. UserId: 1,
  32. ModelName: "gpt-4o",
  33. ChannelId: 5,
  34. Ratio: 0.8,
  35. }
  36. err := ucr.Insert()
  37. require.NoError(t, err)
  38. assert.True(t, ucr.Id > 0)
  39. ratio := GetUserChannelRatio(1, "gpt-4o", 5)
  40. assert.Equal(t, 0.8, ratio)
  41. ratio = GetUserChannelRatio(999, "nonexistent", 999)
  42. assert.Equal(t, 1.0, ratio)
  43. ucr.Ratio = 1.2
  44. err = ucr.Update()
  45. require.NoError(t, err)
  46. ratio = GetUserChannelRatio(1, "gpt-4o", 5)
  47. assert.Equal(t, 1.2, ratio)
  48. list, err := GetUserChannelRatiosByUserId(1)
  49. require.NoError(t, err)
  50. assert.Len(t, list, 1)
  51. assert.Equal(t, "gpt-4o", list[0].ModelName)
  52. err = DeleteUserChannelRatioById(ucr.Id)
  53. require.NoError(t, err)
  54. ratio = GetUserChannelRatio(1, "gpt-4o", 5)
  55. assert.Equal(t, 1.0, ratio)
  56. list, err = GetUserChannelRatiosByUserId(1)
  57. require.NoError(t, err)
  58. assert.Len(t, list, 0)
  59. }
  60. func TestUserChannelRatioUniqueConstraint(t *testing.T) {
  61. setupUserChannelRatioDB(t)
  62. LoadUserChannelRatioCache()
  63. ucr1 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 5, Ratio: 0.8}
  64. err := ucr1.Insert()
  65. require.NoError(t, err)
  66. ucr2 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 5, Ratio: 1.0}
  67. err = ucr2.Insert()
  68. assert.Error(t, err)
  69. ucr3 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 6, Ratio: 1.5}
  70. err = ucr3.Insert()
  71. assert.NoError(t, err)
  72. }
  73. func TestLoadUserChannelRatioCache(t *testing.T) {
  74. setupUserChannelRatioDB(t)
  75. ucr1 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 5, Ratio: 0.8}
  76. ucr2 := &UserChannelRatio{UserId: 2, ModelName: "claude-3", ChannelId: 10, Ratio: 1.5}
  77. require.NoError(t, ucr1.Insert())
  78. require.NoError(t, ucr2.Insert())
  79. LoadUserChannelRatioCache()
  80. assert.Equal(t, 0.8, GetUserChannelRatio(1, "gpt-4o", 5))
  81. assert.Equal(t, 1.5, GetUserChannelRatio(2, "claude-3", 10))
  82. assert.Equal(t, 1.0, GetUserChannelRatio(3, "nonexistent", 99))
  83. }