From 3ac94e2e6f75ec7de10daa6efc027c9881b96ad1 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 17 Jun 2026 19:38:32 +0800 Subject: [PATCH] test(task-pricing): cover model pricing api and remix snapshots --- controller/model_pricing_controller_test.go | 135 ++++++++++++++++++ relay/relay_task_test.go | 148 ++++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 controller/model_pricing_controller_test.go create mode 100644 relay/relay_task_test.go diff --git a/controller/model_pricing_controller_test.go b/controller/model_pricing_controller_test.go new file mode 100644 index 0000000..15c182a --- /dev/null +++ b/controller/model_pricing_controller_test.go @@ -0,0 +1,135 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupModelPricingControllerTest(t *testing.T) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Option{})) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + origUsingSQLite := common.UsingSQLite + origRedisEnabled := common.RedisEnabled + model.DB = db + common.UsingSQLite = true + common.RedisEnabled = false + model.InitOptionMap() + require.NoError(t, ratio_setting.UpdateModelPricingRulesByJSONString(`{}`)) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"video-test":0.25}`)) + + t.Cleanup(func() { + model.DB = origDB + common.UsingSQLite = origUsingSQLite + common.RedisEnabled = origRedisEnabled + require.NoError(t, ratio_setting.UpdateModelPricingRulesByJSONString(`{}`)) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) + require.NoError(t, sqlDB.Close()) + }) + + r := gin.New() + r.GET("/api/option/model_pricing", GetModelPricingRules) + r.GET("/api/option/model_pricing/:model", GetModelPricingRule) + r.PUT("/api/option/model_pricing/:model", UpdateModelPricingRule) + r.DELETE("/api/option/model_pricing/:model", DeleteModelPricingRule) + return r +} + +func TestModelPricingController_CRUD(t *testing.T) { + r := setupModelPricingControllerTest(t) + cfg := validControllerPricingConfig() + + w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/video-test", cfg) + require.Equal(t, http.StatusOK, w.Code) + assert.True(t, responseSuccess(t, w.Body.Bytes())) + + w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/video-test", nil) + require.Equal(t, http.StatusOK, w.Code) + assert.True(t, responseSuccess(t, w.Body.Bytes())) + assert.Contains(t, w.Body.String(), `"billing_unit":"per_call"`) + + w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing", nil) + require.Equal(t, http.StatusOK, w.Code) + assert.True(t, responseSuccess(t, w.Body.Bytes())) + assert.Contains(t, w.Body.String(), "video-test") + + w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_pricing/video-test", nil) + require.Equal(t, http.StatusOK, w.Code) + assert.True(t, responseSuccess(t, w.Body.Bytes())) + + w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/video-test", nil) + assert.Equal(t, http.StatusNotFound, w.Code) + legacyPrice, ok := ratio_setting.GetModelPrice("video-test", false) + require.True(t, ok) + assert.Equal(t, 0.25, legacyPrice) +} + +func TestModelPricingController_PutRejectsBadConfig(t *testing.T) { + r := setupModelPricingControllerTest(t) + cfg := validControllerPricingConfig() + cfg.Table = nil + + w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/video-test", cfg) + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "table must not be empty") +} + +func performJSONRequest(t *testing.T, r *gin.Engine, method, path string, body any) *httptest.ResponseRecorder { + t.Helper() + var payload []byte + if body != nil { + var err error + payload, err = common.Marshal(body) + require.NoError(t, err) + } + req := httptest.NewRequest(method, path, bytes.NewReader(payload)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + return w +} + +func responseSuccess(t *testing.T, payload []byte) bool { + t.Helper() + var body struct { + Success bool `json:"success"` + } + require.NoError(t, common.Unmarshal(payload, &body)) + return body.Success +} + +func validControllerPricingConfig() types.PricingConfig { + return types.PricingConfig{ + SchemaVersion: 1, + Scope: types.PricingScopeModel, + BillingUnit: types.BillingUnitPerCall, + PreconsumeStrategy: types.PreconsumeStrategyExact, + Dimensions: []types.PricingDimension{ + {Key: "resolution", Source: "request.resolution", Type: "string"}, + }, + Table: []types.PricingRow{ + {"resolution": "720P", "price": 0.1, "source": types.PricingRowSourceManual}, + }, + Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject}, + } +} diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go new file mode 100644 index 0000000..b00f40c --- /dev/null +++ b/relay/relay_task_test.go @@ -0,0 +1,148 @@ +package relay + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupRelayTaskTestDB(t *testing.T) { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Task{}, &model.Channel{})) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + origUsingSQLite := common.UsingSQLite + origRedisEnabled := common.RedisEnabled + model.DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + t.Cleanup(func() { + model.DB = origDB + common.UsingSQLite = origUsingSQLite + common.RedisEnabled = origRedisEnabled + require.NoError(t, sqlDB.Close()) + }) +} + +func TestResolveOriginTask_RestoresMatrixOriginPricing(t *testing.T) { + setupRelayTaskTestDB(t) + require.NoError(t, model.DB.Create(&model.Channel{ + Id: 11, + Type: constant.ChannelTypeMiniMax, + Key: "test-key", + Status: common.ChannelStatusEnabled, + Name: "hailuo", + Group: "default", + Models: "matrix-video", + }).Error) + require.NoError(t, model.DB.Create(&model.Task{ + TaskID: "task_origin", + UserId: 7, + ChannelId: 11, + Properties: model.Properties{ + OriginModelName: "matrix-video", + }, + PrivateData: model.TaskPrivateData{ + BillingContext: &model.TaskBillingContext{ + BillingMode: types.BillingModeMatrix, + BillingUnit: types.BillingUnitPer1MTokens, + TokenUnitPriceUSD: 0.5, + PricingSnapshot: map[string]any{ + "resolution": "720P", + "price": 0.5, + "billing_unit": types.BillingUnitPer1MTokens, + "pricing_mode": string(types.PricingModeMatched), + }, + PerCallBilling: false, + }, + }, + }).Error) + + c := newRelayTaskRemixContext(t, "task_origin") + info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} + + taskErr := ResolveOriginTask(c, info) + require.Nil(t, taskErr) + require.NotNil(t, info.OriginPricing) + require.NotNil(t, info.LockedChannel) + assert.Equal(t, constant.TaskActionRemix, info.Action) + assert.Equal(t, "matrix-video", info.OriginModelName) + assert.Equal(t, types.BillingModeMatrix, info.OriginPricing.BillingMode) + assert.Equal(t, types.BillingUnitPer1MTokens, info.OriginPricing.BillingUnit) + assert.Equal(t, 0.5, info.OriginPricing.TokenUnitPriceUSD) + assert.Equal(t, "720P", info.OriginPricing.Snapshot["resolution"]) + lockedChannel, ok := info.LockedChannel.(*model.Channel) + require.True(t, ok) + assert.Equal(t, 11, lockedChannel.Id) +} + +func TestResolveOriginTask_RestoresLegacyOriginPricing(t *testing.T) { + setupRelayTaskTestDB(t) + require.NoError(t, model.DB.Create(&model.Channel{ + Id: 12, + Type: constant.ChannelTypeOpenAI, + Key: "test-key", + Status: common.ChannelStatusEnabled, + Name: "sora", + Group: "default", + Models: "legacy-video", + }).Error) + require.NoError(t, model.DB.Create(&model.Task{ + TaskID: "task_legacy", + UserId: 7, + ChannelId: 12, + Properties: model.Properties{ + OriginModelName: "legacy-video", + }, + PrivateData: model.TaskPrivateData{ + BillingContext: &model.TaskBillingContext{ + ModelPrice: 0.2, + OtherRatios: map[string]float64{"seconds": 10, "size": 1.666667}, + PerCallBilling: true, + }, + }, + }).Error) + + c := newRelayTaskRemixContext(t, "task_legacy") + info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} + + taskErr := ResolveOriginTask(c, info) + require.Nil(t, taskErr) + require.NotNil(t, info.OriginPricing) + require.NotNil(t, info.LockedChannel) + assert.Equal(t, "legacy-video", info.OriginModelName) + assert.Equal(t, 0.2, info.OriginPricing.PriceUSD) + assert.Equal(t, map[string]float64{"seconds": 10, "size": 1.666667}, info.OriginPricing.OtherRatios) + assert.True(t, info.OriginPricing.PerCallBilling) + lockedChannel, ok := info.LockedChannel.(*model.Channel) + require.True(t, ok) + assert.Equal(t, 12, lockedChannel.Id) +} + +func newRelayTaskRemixContext(t *testing.T, originTaskID string) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/"+originTaskID+"/remix", nil) + c.Params = gin.Params{{Key: "video_id", Value: originTaskID}} + return c +}