瀏覽代碼

test(channel-pricing): 缓存写穿 + 字段默认值单元测试

验证 ChannelPricing 的 Insert/Update/Delete 操作正确更新内存缓存,
以及未设置的扩展字段(CacheRatio 等)默认为零值。

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
master
fengsilin 3 週之前
父節點
當前提交
d5d4714908
共有 2 個檔案被更改,包括 127 行新增1 行删除
  1. +6
    -1
      model/channel_pricing.go
  2. +121
    -0
      model/channel_pricing_test.go

+ 6
- 1
model/channel_pricing.go 查看文件

@@ -201,7 +201,12 @@ type ChannelPricingWithChannel struct {
ModelRatio float64 `json:"model_ratio"`
CompletionRatio float64 `json:"completion_ratio"`
ModelPrice float64 `json:"model_price"`
HasCustomPricing bool `json:"has_custom_pricing"` // 是否有自定义定价
HasCustomPricing bool `json:"has_custom_pricing"` // 是否有自定义定价
CacheRatio float64 `json:"cache_ratio"`
CacheCreationRatio float64 `json:"cache_creation_ratio"`
ImageRatio float64 `json:"image_ratio"`
AudioRatio float64 `json:"audio_ratio"`
AudioCompletionRatio float64 `json:"audio_completion_ratio"`
}

// GetChannelPricingByModelWithChannelInfo 获取指定模型的渠道定价(带渠道信息)


+ 121
- 0
model/channel_pricing_test.go 查看文件

@@ -0,0 +1,121 @@
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)
}

Loading…
取消
儲存