new-api/controller/model_pricing_controller_test.go

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