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") }