new-api/model/pricing_test.go

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/QuantumNous/new-api/types"
	"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.UpdateModelPricingRulesByJSONString(`{}`))
	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 TestGetPricingByModel_IncludesMatrixPricingConfig(t *testing.T) {
	setupPricingTest(t)

	const matrixModel = "doubao-seedance-2-0-260128"
	require.NoError(t, ratio_setting.SetPricingConfig(matrixModel, &types.PricingConfig{
		SchemaVersion:      types.PricingSchemaVersion,
		Scope:              types.PricingScopeModel,
		BillingUnit:        types.BillingUnitPer1MTokens,
		PreconsumeStrategy: types.PreconsumeStrategyMinimum,
		Dimensions: []types.PricingDimension{
			{Key: "resolution", Source: "request.resolution", Type: "string"},
			{Key: "video_input", Source: "derived.video_input", Type: "boolean"},
		},
		Table: []types.PricingRow{
			{"resolution": "480p", "video_input": true, "price": 28.0, "source": types.PricingRowSourceManual},
			{"resolution": "480p", "video_input": false, "price": 46.0, "source": types.PricingRowSourceManual},
		},
		Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject},
	}))

	require.NoError(t, (&Model{
		ModelName: matrixModel,
		Status:    1,
	}).Insert())
	require.NoError(t, (&Channel{
		Id:     1,
		Type:   constant.ChannelTypeDoubaoVideoCompatibleAiping,
		Key:    "k",
		Name:   "channel-a",
		Status: common.ChannelStatusEnabled,
		Models: matrixModel,
		Group:  "default",
	}).Insert())

	RefreshPricing()
	pricing := GetPricingByModel(matrixModel)

	require.NotNil(t, pricing)
	require.NotNil(t, pricing.PricingConfig)
	require.Equal(t, types.BillingUnitPer1MTokens, pricing.PricingConfig.BillingUnit)
	require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
	require.Equal(t, 28.0, pricing.ModelPrice)
	require.Equal(t, []constant.EndpointType{constant.EndpointTypeDoubaoVideo}, pricing.SupportedEndpointTypes)
}

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

	const displayModel = "display-pricing-model"
	require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+displayModel+`":10}`))
	require.NoError(t, ratio_setting.SetModelDisplayPricing(displayModel, []types.ModelDisplayPricingItem{
		{Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1},
	}))

	require.NoError(t, (&Model{
		ModelName: displayModel,
		Status:    1,
	}).Insert())
	require.NoError(t, (&Channel{
		Id:     1,
		Type:   constant.ChannelTypeOpenAI,
		Key:    "k",
		Name:   "channel-a",
		Status: common.ChannelStatusEnabled,
		Models: displayModel,
		Group:  "default",
	}).Insert())

	RefreshPricing()
	pricing := GetPricingByModel(displayModel)

	require.NotNil(t, pricing)
	require.Len(t, pricing.DisplayPricing, 1)
	require.Equal(t, "768p-6s", pricing.DisplayPricing[0].Specification)
	require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
	require.Equal(t, 10.0, pricing.ModelRatio)
	require.Equal(t, 0.0, pricing.ModelPrice)
	require.Nil(t, pricing.PricingConfig)
}

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