|
- package model
-
- import (
- "testing"
-
- "github.com/QuantumNous/new-api/setting/ratio_setting"
- "github.com/stretchr/testify/require"
- )
-
- const testPricingModel = "test-pricing-model-apply"
- const testPricingCallModel = "test-pricing-call-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 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}`))
-
- p := &Pricing{}
- applyGlobalDefault(p, testPricingModel)
-
- 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)
- 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)
- }
|