|
- package controller
-
- import (
- "bytes"
- "encoding/json"
- "net/http"
- "net/http/httptest"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "github.com/glebarez/sqlite"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupUserRateLimitControllerDB(t *testing.T) {
- t.Helper()
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := model.DB
- model.DB = db
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{}))
-
- t.Cleanup(func() {
- model.DB = origDB
- sqlDB.Close()
- })
- }
-
- func setupUserRateLimitRouter() *gin.Engine {
- gin.SetMode(gin.TestMode)
- r := gin.New()
- g := r.Group("/api/user")
- {
- g.GET("/:id/rate-limits", GetUserRateLimits)
- g.PUT("/:id/rate-limits", SetUserRateLimits)
- }
- return r
- }
-
- func TestGetUserRateLimits_Empty(t *testing.T) {
- setupUserRateLimitControllerDB(t)
- router := setupUserRateLimitRouter()
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest(http.MethodGet, "/api/user/1/rate-limits", nil)
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- // data 应为空列表(非 nil)
- data, ok := resp["data"]
- require.True(t, ok, "response should have data key")
- assert.Empty(t, data)
- }
-
- func TestSetUserRateLimits_OK(t *testing.T) {
- setupUserRateLimitControllerDB(t)
- router := setupUserRateLimitRouter()
-
- items := []model.UserModelRateLimit{
- {Model: "gpt-4", Rpm: 60},
- }
- body, _ := json.Marshal(items)
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body))
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
-
- // 验证数据确实写入
- list, err := model.GetUserModelRateLimits(1)
- require.NoError(t, err)
- require.Len(t, list, 1)
- assert.Equal(t, "gpt-4", list[0].Model)
- assert.Equal(t, 60, list[0].Rpm)
- }
-
- func TestSetUserRateLimits_NegativeRpm(t *testing.T) {
- setupUserRateLimitControllerDB(t)
- router := setupUserRateLimitRouter()
-
- items := []model.UserModelRateLimit{
- {Model: "gpt-4", Rpm: -1},
- }
- body, _ := json.Marshal(items)
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body))
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.False(t, resp["success"].(bool))
- }
-
- func TestGetUserRateLimits_InvalidId(t *testing.T) {
- setupUserRateLimitControllerDB(t)
- router := setupUserRateLimitRouter()
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest(http.MethodGet, "/api/user/notanid/rate-limits", nil)
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.False(t, resp["success"].(bool))
- }
-
- func TestSetUserRateLimits_EmptySlice(t *testing.T) {
- setupUserRateLimitControllerDB(t)
- router := setupUserRateLimitRouter()
-
- // 先写一条
- _ = model.SetUserModelRateLimits(2, []model.UserModelRateLimit{
- {UserId: 2, Model: "claude-3-opus", Rpm: 10},
- })
-
- // 用空 slice 清空
- body, _ := json.Marshal([]model.UserModelRateLimit{})
- w := httptest.NewRecorder()
- req, _ := http.NewRequest(http.MethodPut, "/api/user/2/rate-limits", bytes.NewReader(body))
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
-
- list, err := model.GetUserModelRateLimits(2)
- require.NoError(t, err)
- assert.Empty(t, list)
- }
|