|
- package model
-
- import (
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/constant"
- "github.com/QuantumNous/new-api/setting/ratio_setting"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- const testPricingModel = "test-pricing-global-model"
-
- func setupPricingTest(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
- origUsingSQLite := common.UsingSQLite
- DB = db
- common.UsingSQLite = true
-
- require.NoError(t, db.AutoMigrate(&Model{}, &Vendor{}, &Channel{}, &Ability{}))
-
- t.Cleanup(func() {
- DB = origDB
- common.UsingSQLite = origUsingSQLite
- sqlDB.Close()
- 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 TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) {
- setupPricingTest(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}`))
-
- require.NoError(t, (&Model{
- ModelName: testPricingModel,
- Status: 1,
- }).Insert())
- require.NoError(t, (&Vendor{
- Id: 1,
- Name: "test-vendor",
- SortOrder: 1,
- }).Insert())
- require.NoError(t, (&Channel{
- Id: 1,
- Type: constant.ChannelTypeOpenAI,
- Key: "k",
- Name: "channel-a",
- Status: common.ChannelStatusEnabled,
- Models: testPricingModel,
- Group: "default",
- }).Insert())
-
- RefreshPricing()
- pricing := GetPricingByModel(testPricingModel)
-
- require.NotNil(t, pricing)
- require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
- require.Equal(t, 10.0, pricing.ModelRatio)
- require.Equal(t, 3.0, pricing.CompletionRatio)
- require.Equal(t, 0.5, pricing.CacheRatio)
- require.Equal(t, 0.75, pricing.CacheCreationRatio)
- }
|