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