Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

188 строки
5.2 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 setupChannelSelectDB(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, err := db.DB()
  15. require.NoError(t, err)
  16. sqlDB.SetMaxOpenConns(1)
  17. origDB := DB
  18. origUsingSQLite := common.UsingSQLite
  19. origRedisEnabled := common.RedisEnabled
  20. origMemoryCache := common.MemoryCacheEnabled
  21. origGroupCol := commonGroupCol
  22. origTrueVal := commonTrueVal
  23. DB = db
  24. common.UsingSQLite = true
  25. common.RedisEnabled = false
  26. common.MemoryCacheEnabled = false
  27. // Set column quoting for SQLite
  28. commonGroupCol = "`group`"
  29. commonTrueVal = "1"
  30. require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{}))
  31. t.Cleanup(func() {
  32. DB = origDB
  33. common.UsingSQLite = origUsingSQLite
  34. common.RedisEnabled = origRedisEnabled
  35. common.MemoryCacheEnabled = origMemoryCache
  36. commonGroupCol = origGroupCol
  37. commonTrueVal = origTrueVal
  38. sqlDB.Close()
  39. })
  40. return db
  41. }
  42. func int64Ptr(v int64) *int64 { return &v }
  43. func intPtr(v int) *int { return &v }
  44. func uintPtr(v uint) *uint { return &v }
  45. func createChannelForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, priority int64) {
  46. t.Helper()
  47. ch := Channel{
  48. Id: id,
  49. Type: 1,
  50. Key: "test-key-" + name,
  51. Status: common.ChannelStatusEnabled,
  52. Name: name,
  53. Group: group,
  54. Models: models,
  55. Priority: int64Ptr(priority),
  56. Weight: uintPtr(10),
  57. AutoBan: intPtr(1),
  58. CreatedTime: 1,
  59. }
  60. require.NoError(t, db.Create(&ch).Error)
  61. }
  62. func createAbilityForTest(t *testing.T, db *gorm.DB, group, modelName string, channelId int, priority int64, enabled bool) {
  63. t.Helper()
  64. ability := Ability{
  65. Group: group,
  66. Model: modelName,
  67. ChannelId: channelId,
  68. Enabled: enabled,
  69. Priority: int64Ptr(priority),
  70. Weight: 10,
  71. }
  72. require.NoError(t, db.Create(&ability).Error)
  73. }
  74. // --- GetChannel (DB path) tests ---
  75. func TestGetChannel_BasicSelection(t *testing.T) {
  76. db := setupChannelSelectDB(t)
  77. createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
  78. createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
  79. ch, err := GetChannel("default", "gpt-4", 0)
  80. require.NoError(t, err)
  81. assert.NotNil(t, ch)
  82. assert.Equal(t, 1, ch.Id)
  83. }
  84. func TestGetChannel_NoChannelForModel(t *testing.T) {
  85. db := setupChannelSelectDB(t)
  86. createChannelForTest(t, db, 1, "ch1", "default", "gpt-3.5", 10)
  87. createAbilityForTest(t, db, "default", "gpt-3.5", 1, 10, true)
  88. ch, err := GetChannel("default", "gpt-4", 0)
  89. assert.NoError(t, err)
  90. assert.Nil(t, ch)
  91. }
  92. func TestGetChannel_PrioritySelection(t *testing.T) {
  93. db := setupChannelSelectDB(t)
  94. // Two channels with different priorities
  95. createChannelForTest(t, db, 1, "ch-low", "default", "gpt-4", 0)
  96. createChannelForTest(t, db, 2, "ch-high", "default", "gpt-4", 10)
  97. createAbilityForTest(t, db, "default", "gpt-4", 1, 0, true)
  98. createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true)
  99. // Retry=0 should get highest priority (10)
  100. ch, err := GetChannel("default", "gpt-4", 0)
  101. require.NoError(t, err)
  102. assert.NotNil(t, ch)
  103. assert.Equal(t, 2, ch.Id) // ch-high has priority 10
  104. // Retry=1 should get next priority (0)
  105. ch, err = GetChannel("default", "gpt-4", 1)
  106. require.NoError(t, err)
  107. assert.NotNil(t, ch)
  108. assert.Equal(t, 1, ch.Id) // ch-low has priority 0
  109. }
  110. func TestGetChannel_DifferentGroups(t *testing.T) {
  111. db := setupChannelSelectDB(t)
  112. // Channel in "default" group
  113. createChannelForTest(t, db, 1, "ch-default", "default", "gpt-4", 10)
  114. createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
  115. // Channel in "vip" group
  116. createChannelForTest(t, db, 2, "ch-vip", "vip", "gpt-4", 10)
  117. createAbilityForTest(t, db, "vip", "gpt-4", 2, 10, true)
  118. // Query default group
  119. ch, err := GetChannel("default", "gpt-4", 0)
  120. require.NoError(t, err)
  121. assert.Equal(t, 1, ch.Id)
  122. // Query vip group
  123. ch, err = GetChannel("vip", "gpt-4", 0)
  124. require.NoError(t, err)
  125. assert.Equal(t, 2, ch.Id)
  126. }
  127. func TestGetChannel_DisabledAbility(t *testing.T) {
  128. db := setupChannelSelectDB(t)
  129. createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
  130. createAbilityForTest(t, db, "default", "gpt-4", 1, 10, false) // disabled
  131. ch, err := GetChannel("default", "gpt-4", 0)
  132. assert.NoError(t, err)
  133. assert.Nil(t, ch) // no enabled ability
  134. }
  135. func TestGetChannel_MultipleChannelsWeighted(t *testing.T) {
  136. db := setupChannelSelectDB(t)
  137. createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
  138. createChannelForTest(t, db, 2, "ch2", "default", "gpt-4", 10)
  139. createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
  140. createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true)
  141. // Run multiple times to verify both channels can be selected
  142. found := make(map[int]bool)
  143. for i := 0; i < 50; i++ {
  144. ch, err := GetChannel("default", "gpt-4", 0)
  145. require.NoError(t, err)
  146. require.NotNil(t, ch)
  147. found[ch.Id] = true
  148. }
  149. // With equal weights, both should appear at least once in 50 tries
  150. assert.True(t, found[1], "channel 1 should be selected at least once")
  151. assert.True(t, found[2], "channel 2 should be selected at least once")
  152. }