|
- 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/gin-gonic/gin"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupUserControllerDB(t *testing.T) *gorm.DB {
- t.Helper()
-
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
-
- sqlDB, err := db.DB()
- require.NoError(t, err)
- sqlDB.SetMaxOpenConns(1)
-
- require.NoError(t, db.AutoMigrate(&model.User{}))
- return db
- }
-
- func TestUpdateUser_SyncedUserAllowsNonQuotaEditsWithoutOverwritingQuota(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testRootAuthMiddleware())
- router.PUT("/api/user/", UpdateUser)
-
- db := setupUserControllerDB(t)
- orig := model.DB
- origRedisEnabled := common.RedisEnabled
- model.DB = db
- defer func() {
- model.DB = orig
- common.RedisEnabled = origRedisEnabled
- }()
- common.RedisEnabled = false
-
- require.NoError(t, db.Create(&model.User{
- Id: 1001,
- Username: "synced-user",
- Password: "hashed-password",
- DisplayName: "Old Name",
- Role: common.RoleCommonUser,
- Status: common.UserStatusEnabled,
- Group: "default",
- Quota: 0,
- SyncedQuota: 100000,
- Source: common.UserSourceSynced,
- AffCode: "SYNC1",
- }).Error)
-
- body := map[string]any{
- "id": 1001,
- "username": "synced-user",
- "display_name": "New Name",
- "password": "",
- "group": "default",
- "quota": 100000,
- }
- bodyBytes, err := json.Marshal(body)
- require.NoError(t, err)
-
- req := httptest.NewRequest(http.MethodPut, "/api/user/", bytes.NewBuffer(bodyBytes))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusOK, w.Code)
- var resp map[string]any
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- require.True(t, resp["success"].(bool))
-
- var saved model.User
- require.NoError(t, db.First(&saved, 1001).Error)
- require.Equal(t, "New Name", saved.DisplayName)
- require.Equal(t, 0, saved.Quota)
- require.Equal(t, 100000, saved.SyncedQuota)
- }
-
- func TestUpdateUser_SyncedUserQuotaChangeRejected(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testRootAuthMiddleware())
- router.PUT("/api/user/", UpdateUser)
-
- db := setupUserControllerDB(t)
- orig := model.DB
- origRedisEnabled := common.RedisEnabled
- model.DB = db
- defer func() {
- model.DB = orig
- common.RedisEnabled = origRedisEnabled
- }()
- common.RedisEnabled = false
-
- require.NoError(t, db.Create(&model.User{
- Id: 1002,
- Username: "synced-user-2",
- Password: "hashed-password",
- DisplayName: "Synced User 2",
- Role: common.RoleCommonUser,
- Status: common.UserStatusEnabled,
- Group: "default",
- Quota: 0,
- SyncedQuota: 100000,
- Source: common.UserSourceSynced,
- AffCode: "SYNC2",
- }).Error)
-
- body := map[string]any{
- "id": 1002,
- "username": "synced-user-2",
- "display_name": "Synced User 2",
- "password": "",
- "group": "default",
- "quota": 100001,
- }
- bodyBytes, err := json.Marshal(body)
- require.NoError(t, err)
-
- req := httptest.NewRequest(http.MethodPut, "/api/user/", bytes.NewBuffer(bodyBytes))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusOK, w.Code)
- var resp map[string]any
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- require.False(t, resp["success"].(bool))
- require.Equal(t, "user.synced_user_quota_cannot_modify", resp["message"])
-
- var saved model.User
- require.NoError(t, db.First(&saved, 1002).Error)
- require.Equal(t, 0, saved.Quota)
- require.Equal(t, 100000, saved.SyncedQuota)
- }
|