|
- package controller
-
- import (
- "encoding/json"
- "net/http"
- "net/http/httptest"
- "strings"
- "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 setupInternalMigrationDB(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{}, &model.CustomOAuthProvider{}, &model.UserOAuthBinding{}))
-
- return db
- }
-
- func TestListMigrationUsers_ReturnsLocalOnly(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.GET("/api/internal/migration/users", ListMigrationUsers)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 10000001,
- Username: "ov-1",
- Password: "hash",
- Email: "u1@example.com",
- AffCode: "A1",
- Source: common.UserSourceLocal,
- Quota: 99,
- }).Error)
- require.NoError(t, db.Create(&model.User{
- Id: 10000002,
- Username: "ov-2",
- Password: "hash",
- AffCode: "A2",
- Source: common.UserSourceSynced,
- RemoteUserId: 1,
- }).Error)
-
- req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users?page=0&page_size=20", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusOK, w.Code)
- var body map[string]any
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
- require.True(t, body["success"].(bool))
- require.EqualValues(t, 1, body["total"])
- }
-
- func TestConvertMigrationUserToSynced_Idempotent(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 10000002,
- Username: "ov-2",
- Password: "hash",
- AffCode: "A2",
- Source: common.UserSourceLocal,
- Quota: 66,
- }).Error)
-
- body := `{"remote_user_id":2002,"synced_quota":66}`
- for i := 0; i < 2; i++ {
- req := httptest.NewRequest(http.MethodPost,
- "/api/internal/migration/users/10000002/convert-to-synced",
- strings.NewReader(body))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusOK, w.Code)
- }
- }
-
- func TestConvertMigrationUserToSynced_ConflictRemoteId(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 10000003,
- Username: "ov-3",
- Password: "hash",
- AffCode: "A3",
- Source: common.UserSourceSynced,
- RemoteUserId: 5000,
- }).Error)
-
- body := `{"remote_user_id":9999,"synced_quota":100}`
- req := httptest.NewRequest(http.MethodPost,
- "/api/internal/migration/users/10000003/convert-to-synced",
- strings.NewReader(body))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusConflict, w.Code)
- }
-
- func TestCheckSyncedCopy(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.GET("/api/internal/migration/users/check-synced-copy", CheckSyncedCopy)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 10000010,
- Username: "synced",
- Password: "hash",
- AffCode: "S1",
- Source: common.UserSourceSynced,
- RemoteUserId: 456,
- }).Error)
-
- req := httptest.NewRequest(http.MethodGet,
- "/api/internal/migration/users/check-synced-copy?cn_user_id=456", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusOK, w.Code)
-
- var body map[string]any
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
- require.True(t, body["has_copy"].(bool))
- }
|