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