Просмотр исходного кода

test(model): switch pricing tests to global defaults

master
fengsilin 2 недель назад
Родитель
Сommit
415174a676
2 измененных файлов: 23 добавлений и 217 удалений
  1. +0
    -121
      model/channel_pricing_test.go
  2. +23
    -96
      model/pricing_test.go

+ 0
- 121
model/channel_pricing_test.go Просмотреть файл

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

+ 23
- 96
model/pricing_test.go Просмотреть файл

@@ -4,119 +4,46 @@ import (
"testing"

"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

const testPricingModel = "test-pricing-model-apply"
const testPricingCallModel = "test-pricing-call-model"

func setupPricingTest(t *testing.T) {
func resetPricingRatios(t *testing.T) {
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
require.NoError(t, db.AutoMigrate(&ChannelPricing{}))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`))
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`))
require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`))
require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`))
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
}

// 全局默认定价
func TestApplyGlobalDefault_PerTokenModel(t *testing.T) {
resetPricingRatios(t)
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`))
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`))
require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`))
require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`))

t.Cleanup(func() {
DB = origDB
sqlDB.Close()
ratio_setting.UpdateModelRatioByJSONString(`{}`)
ratio_setting.UpdateCompletionRatioByJSONString(`{}`)
ratio_setting.UpdateCacheRatioByJSONString(`{}`)
ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)
ratio_setting.UpdateModelPriceByJSONString(`{}`)
})
}

func TestApplyBestChannelPricing_NoChannelPricing(t *testing.T) {
setupPricingTest(t)

p := &Pricing{}
applyBestChannelPricing(p, nil, testPricingModel)

require.Equal(t, 0, p.QuotaType, "应使用全局按量计费")
require.Equal(t, 10.0, p.ModelRatio, "应使用全局 model_ratio")
require.Equal(t, 3.0, p.CompletionRatio, "应使用全局 completion_ratio")
require.Equal(t, 0.5, p.CacheRatio, "应使用全局 cache_ratio")
require.Equal(t, 0.75, p.CacheCreationRatio, "应使用全局 cache_creation_ratio")
}

func TestApplyBestChannelPricing_PerTokenCheapest(t *testing.T) {
setupPricingTest(t)

// 渠道A: model_ratio=8(更便宜)
require.NoError(t, (&ChannelPricing{
ModelName: testPricingModel, ChannelId: 1,
QuotaType: QuotaTypeByTokens, ModelRatio: 8, CompletionRatio: 2,
CacheRatio: 0.3, CacheCreationRatio: 0.6,
}).Insert())
// 渠道B: model_ratio=12(更贵)
require.NoError(t, (&ChannelPricing{
ModelName: testPricingModel, ChannelId: 2,
QuotaType: QuotaTypeByTokens, ModelRatio: 12, CompletionRatio: 4,
CacheRatio: 0.8, CacheCreationRatio: 1.0,
}).Insert())

cps := []ChannelPricing{
{ModelRatio: 8, CompletionRatio: 2, CacheRatio: 0.3, CacheCreationRatio: 0.6},
{ModelRatio: 12, CompletionRatio: 4, CacheRatio: 0.8, CacheCreationRatio: 1.0},
}

p := &Pricing{}
applyBestChannelPricing(p, cps, testPricingModel)

require.Equal(t, 0, p.QuotaType)
require.Equal(t, 8.0, p.ModelRatio, "应选最便宜的渠道A")
require.Equal(t, 2.0, p.CompletionRatio)
require.Equal(t, 0.3, p.CacheRatio)
require.Equal(t, 0.6, p.CacheCreationRatio)
}

func TestApplyBestChannelPricing_ExtendedRatioZeroFallback(t *testing.T) {
setupPricingTest(t)

// 渠道定价中扩展比率为 0,应回退到全局值
require.NoError(t, (&ChannelPricing{
ModelName: testPricingModel, ChannelId: 1,
QuotaType: QuotaTypeByTokens, ModelRatio: 5, CompletionRatio: 1,
CacheRatio: 0, CacheCreationRatio: 0, // 0 = 未设置
}).Insert())

cps := []ChannelPricing{
{ModelRatio: 5, CompletionRatio: 1, CacheRatio: 0, CacheCreationRatio: 0},
}

p := &Pricing{}
applyBestChannelPricing(p, cps, testPricingModel)
applyGlobalDefault(p, testPricingModel)

require.Equal(t, 0.5, p.CacheRatio, "cache_ratio=0 应回退全局 0.5")
require.Equal(t, 0.75, p.CacheCreationRatio, "cache_creation_ratio=0 应回退全局 0.75")
require.Equal(t, QuotaTypeByTokens, p.QuotaType)
require.Equal(t, 10.0, p.ModelRatio)
require.Equal(t, 3.0, p.CompletionRatio)
require.Equal(t, 0.5, p.CacheRatio)
require.Equal(t, 0.75, p.CacheCreationRatio)
}

func TestApplyBestChannelPricing_PerCallCheapest(t *testing.T) {
setupPricingTest(t)
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingModel+`":0.8}`))

// 只有按次计费的渠道
cps := []ChannelPricing{
{QuotaType: QuotaTypeByCall, ModelPrice: 0.3},
{QuotaType: QuotaTypeByCall, ModelPrice: 0.5},
}
func TestApplyGlobalDefault_PerCallModel(t *testing.T) {
resetPricingRatios(t)
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`))

p := &Pricing{}
applyBestChannelPricing(p, cps, testPricingModel)
applyGlobalDefault(p, testPricingCallModel)

require.Equal(t, 1, p.QuotaType)
require.Equal(t, 0.3, p.ModelPrice, "应选最便宜的按次渠道")
require.Equal(t, QuotaTypeByCall, p.QuotaType)
require.Equal(t, 0.3, p.ModelPrice)
require.Equal(t, 0.0, p.ModelRatio)
}

Загрузка…
Отмена
Сохранить