package model import ( "testing" "time" "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 testPricingApplyModel = "test-pricing-model-apply" const testPricingCallModel = "test-pricing-call-model" const testPricingGlobalModel = "test-pricing-global-model" func resetPricingRatios(t *testing.T) { t.Helper() 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 setupPricingTest(t *testing.T) { t.Helper() resetPricingRatios(t) db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) origDB := DB origUsingSQLite := common.UsingSQLite origPricingMap := pricingMap origPricingByModel := pricingByModel origVendorsList := vendorsList origSupportedEndpointMap := supportedEndpointMap origModelSupportEndpointTypes := modelSupportEndpointTypes origLastGetPricingTime := lastGetPricingTime DB = db common.UsingSQLite = true pricingMap = nil pricingByModel = nil vendorsList = nil supportedEndpointMap = nil modelSupportEndpointTypes = make(map[string][]constant.EndpointType) lastGetPricingTime = time.Time{} require.NoError(t, db.AutoMigrate(&Model{}, &Vendor{}, &Channel{}, &Ability{})) t.Cleanup(func() { DB = origDB common.UsingSQLite = origUsingSQLite pricingMap = origPricingMap pricingByModel = origPricingByModel vendorsList = origVendorsList supportedEndpointMap = origSupportedEndpointMap modelSupportEndpointTypes = origModelSupportEndpointTypes lastGetPricingTime = origLastGetPricingTime require.NoError(t, sqlDB.Close()) resetPricingRatios(t) }) } func TestApplyGlobalDefault_PerTokenModel(t *testing.T) { resetPricingRatios(t) t.Cleanup(func() { resetPricingRatios(t) }) require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingApplyModel+`":10}`)) require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingApplyModel+`":3}`)) require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.5}`)) require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.75}`)) p := &Pricing{} applyGlobalDefault(p, testPricingApplyModel) 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 TestApplyGlobalDefault_PerCallModel(t *testing.T) { resetPricingRatios(t) t.Cleanup(func() { resetPricingRatios(t) }) require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`)) p := &Pricing{} applyGlobalDefault(p, testPricingCallModel) require.Equal(t, QuotaTypeByCall, p.QuotaType) require.Equal(t, 0.3, p.ModelPrice) require.Equal(t, 0.0, p.ModelRatio) } func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) { setupPricingTest(t) require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingGlobalModel+`":10}`)) require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingGlobalModel+`":3}`)) require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.5}`)) require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.75}`)) require.NoError(t, (&Model{ ModelName: testPricingGlobalModel, 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: testPricingGlobalModel, Group: "default", }).Insert()) RefreshPricing() pricing := GetPricingByModel(testPricingGlobalModel) 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) } func TestParseModelEndpointTypes_SupportsObjectAndArray(t *testing.T) { objectEndpoints := parseModelEndpointTypes(`{ "anthropic": {"path": "/v1/messages", "method": "POST"}, "openai": "/v1/chat/completions" }`) require.ElementsMatch(t, []string{"anthropic", "openai"}, objectEndpoints) arrayEndpoints := parseModelEndpointTypes(`["anthropic", "openai", "anthropic"]`) require.Equal(t, []string{"anthropic", "openai"}, arrayEndpoints) } func TestParseModelEndpointTypes_IgnoresInvalidJSON(t *testing.T) { require.Nil(t, parseModelEndpointTypes(`123`)) require.Nil(t, parseModelEndpointTypes(`{"anthropic": true}`)) require.Nil(t, parseModelEndpointTypes(``)) info, ok := common.GetDefaultEndpointInfo("anthropic") require.True(t, ok) require.Equal(t, "/v1/messages", info.Path) }