|
- package model
-
- import (
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupChannelSelectDB(t *testing.T) *gorm.DB {
- t.Helper()
-
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
-
- sqlDB, err := db.DB()
- require.NoError(t, err)
- sqlDB.SetMaxOpenConns(1)
-
- origDB := DB
- origUsingSQLite := common.UsingSQLite
- origRedisEnabled := common.RedisEnabled
- origMemoryCache := common.MemoryCacheEnabled
- origGroupCol := commonGroupCol
- origTrueVal := commonTrueVal
-
- DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
- common.MemoryCacheEnabled = false
- // Set column quoting for SQLite
- commonGroupCol = "`group`"
- commonTrueVal = "1"
-
- require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{}))
-
- t.Cleanup(func() {
- DB = origDB
- common.UsingSQLite = origUsingSQLite
- common.RedisEnabled = origRedisEnabled
- common.MemoryCacheEnabled = origMemoryCache
- commonGroupCol = origGroupCol
- commonTrueVal = origTrueVal
- sqlDB.Close()
- })
-
- return db
- }
-
- func int64Ptr(v int64) *int64 { return &v }
- func intPtr(v int) *int { return &v }
- func uintPtr(v uint) *uint { return &v }
-
- func createChannelForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, priority int64) {
- t.Helper()
- ch := Channel{
- Id: id,
- Type: 1,
- Key: "test-key-" + name,
- Status: common.ChannelStatusEnabled,
- Name: name,
- Group: group,
- Models: models,
- Priority: int64Ptr(priority),
- Weight: uintPtr(10),
- AutoBan: intPtr(1),
- CreatedTime: 1,
- }
- require.NoError(t, db.Create(&ch).Error)
- }
-
- func createAbilityForTest(t *testing.T, db *gorm.DB, group, modelName string, channelId int, priority int64, enabled bool) {
- t.Helper()
- ability := Ability{
- Group: group,
- Model: modelName,
- ChannelId: channelId,
- Enabled: enabled,
- Priority: int64Ptr(priority),
- Weight: 10,
- }
- require.NoError(t, db.Create(&ability).Error)
- }
-
- // --- GetChannel (DB path) tests ---
-
- func TestGetChannel_BasicSelection(t *testing.T) {
- db := setupChannelSelectDB(t)
-
- createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
- createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
-
- ch, err := GetChannel("default", "gpt-4", 0)
- require.NoError(t, err)
- assert.NotNil(t, ch)
- assert.Equal(t, 1, ch.Id)
- }
-
- func TestGetChannel_NoChannelForModel(t *testing.T) {
- db := setupChannelSelectDB(t)
-
- createChannelForTest(t, db, 1, "ch1", "default", "gpt-3.5", 10)
- createAbilityForTest(t, db, "default", "gpt-3.5", 1, 10, true)
-
- ch, err := GetChannel("default", "gpt-4", 0)
- assert.NoError(t, err)
- assert.Nil(t, ch)
- }
-
- func TestGetChannel_PrioritySelection(t *testing.T) {
- db := setupChannelSelectDB(t)
-
- // Two channels with different priorities
- createChannelForTest(t, db, 1, "ch-low", "default", "gpt-4", 0)
- createChannelForTest(t, db, 2, "ch-high", "default", "gpt-4", 10)
- createAbilityForTest(t, db, "default", "gpt-4", 1, 0, true)
- createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true)
-
- // Retry=0 should get highest priority (10)
- ch, err := GetChannel("default", "gpt-4", 0)
- require.NoError(t, err)
- assert.NotNil(t, ch)
- assert.Equal(t, 2, ch.Id) // ch-high has priority 10
-
- // Retry=1 should get next priority (0)
- ch, err = GetChannel("default", "gpt-4", 1)
- require.NoError(t, err)
- assert.NotNil(t, ch)
- assert.Equal(t, 1, ch.Id) // ch-low has priority 0
- }
-
- func TestGetChannel_DifferentGroups(t *testing.T) {
- db := setupChannelSelectDB(t)
-
- // Channel in "default" group
- createChannelForTest(t, db, 1, "ch-default", "default", "gpt-4", 10)
- createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
-
- // Channel in "vip" group
- createChannelForTest(t, db, 2, "ch-vip", "vip", "gpt-4", 10)
- createAbilityForTest(t, db, "vip", "gpt-4", 2, 10, true)
-
- // Query default group
- ch, err := GetChannel("default", "gpt-4", 0)
- require.NoError(t, err)
- assert.Equal(t, 1, ch.Id)
-
- // Query vip group
- ch, err = GetChannel("vip", "gpt-4", 0)
- require.NoError(t, err)
- assert.Equal(t, 2, ch.Id)
- }
-
- func TestGetChannel_DisabledAbility(t *testing.T) {
- db := setupChannelSelectDB(t)
-
- createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
- createAbilityForTest(t, db, "default", "gpt-4", 1, 10, false) // disabled
-
- ch, err := GetChannel("default", "gpt-4", 0)
- assert.NoError(t, err)
- assert.Nil(t, ch) // no enabled ability
- }
-
- func TestGetChannel_MultipleChannelsWeighted(t *testing.T) {
- db := setupChannelSelectDB(t)
-
- createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
- createChannelForTest(t, db, 2, "ch2", "default", "gpt-4", 10)
- createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
- createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true)
-
- // Run multiple times to verify both channels can be selected
- found := make(map[int]bool)
- for i := 0; i < 50; i++ {
- ch, err := GetChannel("default", "gpt-4", 0)
- require.NoError(t, err)
- require.NotNil(t, ch)
- found[ch.Id] = true
- }
- // With equal weights, both should appear at least once in 50 tries
- assert.True(t, found[1], "channel 1 should be selected at least once")
- assert.True(t, found[2], "channel 2 should be selected at least once")
- }
|