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 setupChannelPricingDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) sqlDB, _ := db.DB() sqlDB.SetMaxOpenConns(1) origDB := DB DB = db common.UsingSQLite = true common.RedisEnabled = false require.NoError(t, db.AutoMigrate(&ChannelPricing{})) t.Cleanup(func() { DB = origDB sqlDB.Close() }) return db } func TestCacheWriteThrough_Insert(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ ModelName: "test-insert-model", ChannelId: 9001, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, CacheRatio: 0.8, ImageRatio: 1.2, } require.NoError(t, cp.Insert()) defer cp.Delete() found, ok := GetEffectivePricing("test-insert-model", 9001) assert.True(t, ok) assert.Equal(t, 0.8, found.CacheRatio) assert.Equal(t, 1.2, found.ImageRatio) } func TestCacheWriteThrough_Update(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ ModelName: "test-update-model", ChannelId: 9002, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, } require.NoError(t, cp.Insert()) defer cp.Delete() cp.CacheRatio = 0.9 cp.AudioRatio = 1.5 require.NoError(t, cp.Update()) found, ok := GetEffectivePricing("test-update-model", 9002) assert.True(t, ok) assert.Equal(t, 0.9, found.CacheRatio) assert.Equal(t, 1.5, found.AudioRatio) } func TestCacheWriteThrough_Delete(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ ModelName: "test-delete-model", ChannelId: 9003, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, } require.NoError(t, cp.Insert()) require.NoError(t, cp.Delete()) found, ok := GetEffectivePricing("test-delete-model", 9003) assert.False(t, ok) assert.Nil(t, found) } func TestExtendedFields_DefaultZero(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ ModelName: "test-default-model", ChannelId: 9004, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, } require.NoError(t, cp.Insert()) defer cp.Delete() found, ok := GetEffectivePricing("test-default-model", 9004) assert.True(t, ok) assert.Equal(t, 0.0, found.CacheRatio) assert.Equal(t, 0.0, found.CacheCreationRatio) assert.Equal(t, 0.0, found.ImageRatio) assert.Equal(t, 0.0, found.AudioRatio) assert.Equal(t, 0.0, found.AudioCompletionRatio) } func TestGetEffectivePricing_NotFound(t *testing.T) { setupChannelPricingDB(t) found, ok := GetEffectivePricing("nonexistent-model-xyz", 99999) assert.False(t, ok) assert.Nil(t, found) } // === 默认通道测试 === func TestDefaultChannel_SetAndGet(t *testing.T) { setupChannelPricingDB(t) // 创建两条定价记录 cp1 := &ChannelPricing{ModelName: "default-test-model", ChannelId: 100, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0} cp2 := &ChannelPricing{ModelName: "default-test-model", ChannelId: 200, QuotaType: QuotaTypeByTokens, ModelRatio: 2.0} require.NoError(t, cp1.Insert()) require.NoError(t, cp2.Insert()) t.Cleanup(func() { cp1.Delete(); cp2.Delete() }) // 初始没有默认 _, ok := GetDefaultChannelId("default-test-model") assert.False(t, ok) // 设置通道 100 为默认 require.NoError(t, SetDefaultChannel("default-test-model", 100)) id, ok := GetDefaultChannelId("default-test-model") assert.True(t, ok) assert.Equal(t, 100, id) // 切换默认到通道 200 require.NoError(t, SetDefaultChannel("default-test-model", 200)) id, ok = GetDefaultChannelId("default-test-model") assert.True(t, ok) assert.Equal(t, 200, id) // 验证旧的默认标记被清除 found, _ := GetEffectivePricing("default-test-model", 100) assert.False(t, found.IsDefault) found, _ = GetEffectivePricing("default-test-model", 200) assert.True(t, found.IsDefault) } func TestDefaultChannel_Clear(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ModelName: "clear-test-model", ChannelId: 300, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0} require.NoError(t, cp.Insert()) t.Cleanup(func() { cp.Delete() }) require.NoError(t, SetDefaultChannel("clear-test-model", 300)) _, ok := GetDefaultChannelId("clear-test-model") assert.True(t, ok) require.NoError(t, ClearDefaultChannel("clear-test-model")) _, ok = GetDefaultChannelId("clear-test-model") assert.False(t, ok) } func TestDefaultChannel_DeleteClearsCache(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ModelName: "delete-test-model", ChannelId: 400, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, IsDefault: true} require.NoError(t, cp.Insert()) id, ok := GetDefaultChannelId("delete-test-model") assert.True(t, ok) assert.Equal(t, 400, id) // 删除记录后缓存应清除 require.NoError(t, cp.Delete()) _, ok = GetDefaultChannelId("delete-test-model") assert.False(t, ok) } func TestDefaultChannel_LoadCache(t *testing.T) { db := setupChannelPricingDB(t) // 直接插入数据(绕过缓存) db.Create(&ChannelPricing{ModelName: "load-model", ChannelId: 500, ModelRatio: 1.0, IsDefault: true}) db.Create(&ChannelPricing{ModelName: "load-model", ChannelId: 501, ModelRatio: 2.0, IsDefault: false}) db.Create(&ChannelPricing{ModelName: "other-model", ChannelId: 502, ModelRatio: 1.0, IsDefault: true}) // 全量加载缓存 LoadChannelPricingCache() id, ok := GetDefaultChannelId("load-model") assert.True(t, ok) assert.Equal(t, 500, id) id, ok = GetDefaultChannelId("other-model") assert.True(t, ok) assert.Equal(t, 502, id) // 无默认的模型 _, ok = GetDefaultChannelId("nonexistent") assert.False(t, ok) } func TestDefaultChannel_InsertWithDefault(t *testing.T) { setupChannelPricingDB(t) cp := &ChannelPricing{ModelName: "insert-default-model", ChannelId: 600, ModelRatio: 1.0, IsDefault: true} require.NoError(t, cp.Insert()) t.Cleanup(func() { cp.Delete() }) id, ok := GetDefaultChannelId("insert-default-model") assert.True(t, ok) assert.Equal(t, 600, id) } func TestDefaultChannel_UpdateWithDefault(t *testing.T) { setupChannelPricingDB(t) cp1 := &ChannelPricing{ModelName: "update-default-model", ChannelId: 700, ModelRatio: 1.0, IsDefault: true} cp2 := &ChannelPricing{ModelName: "update-default-model", ChannelId: 701, ModelRatio: 2.0} require.NoError(t, cp1.Insert()) require.NoError(t, cp2.Insert()) t.Cleanup(func() { cp1.Delete(); cp2.Delete() }) // cp1 是默认,通过 Update 把 cp2 设为默认 cp2.IsDefault = true require.NoError(t, cp2.Update()) id, ok := GetDefaultChannelId("update-default-model") assert.True(t, ok) assert.Equal(t, 701, id) }