|
- package controller
-
- import (
- "encoding/json"
- "net/http"
- "net/http/httptest"
- "strings"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- user_migration "github.com/QuantumNous/new-api/service/user_migration"
- "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)
- require.NoError(t, db.Create(&model.User{
- Id: 10000003,
- Username: "root",
- Password: "hash",
- AffCode: "A3",
- Source: common.UserSourceLocal,
- Role: common.RoleRootUser,
- }).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 TestListMigrationUsers_FiltersByKeyword(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: 10000011,
- Username: "alice-local",
- Password: "hash",
- Email: "alice@example.com",
- AffCode: "A11",
- Source: common.UserSourceLocal,
- }).Error)
- require.NoError(t, db.Create(&model.User{
- Id: 10000012,
- Username: "bob-local",
- Password: "hash",
- Email: "bob@example.com",
- AffCode: "A12",
- Source: common.UserSourceLocal,
- }).Error)
-
- req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users?page=0&page_size=20&keyword=alice", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusOK, w.Code)
- var body struct {
- Success bool `json:"success"`
- Data []user_migration.RemoteUserSnapshot `json:"data"`
- Total int `json:"total"`
- }
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
- require.True(t, body.Success)
- require.Equal(t, 1, body.Total)
- require.Len(t, body.Data, 1)
- require.Equal(t, 10000011, body.Data[0].Id)
- }
-
- func TestQueryMigrationUsers_OnlyReturnsSelectedLocalUsers(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.POST("/api/internal/migration/users/query", QueryMigrationUsers)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{Id: 10000001, Username: "u1", Password: "hash", AffCode: "A1", Source: common.UserSourceLocal}).Error)
- require.NoError(t, db.Create(&model.User{Id: 10000002, Username: "u2", Password: "hash", AffCode: "A2", Source: common.UserSourceLocal, Role: common.RoleRootUser}).Error)
- require.NoError(t, db.Create(&model.User{Id: 10000003, Username: "u3", Password: "hash", AffCode: "A3", Source: common.UserSourceSynced}).Error)
-
- req := httptest.NewRequest(http.MethodPost, "/api/internal/migration/users/query", strings.NewReader(`{"page":0,"page_size":50,"selection_mode":"explicit_ids","source_user_ids":[10000001,10000002,10000003,10000004]}`))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusOK, w.Code)
- var body struct {
- Success bool `json:"success"`
- Data []user_migration.RemoteUserSnapshot `json:"data"`
- Total int `json:"total"`
- SelectionSummary user_migration.SelectionSummary `json:"selection_summary"`
- }
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
- require.True(t, body.Success)
- require.Len(t, body.Data, 1)
- require.Equal(t, 10000001, body.Data[0].Id)
- require.Equal(t, 1, body.Total)
- require.Equal(t, 4, body.SelectionSummary.Requested)
- require.Equal(t, 1, body.SelectionSummary.Matched)
- require.Len(t, body.SelectionSummary.Excluded, 3)
- }
-
- func TestQueryMigrationUsers_InvalidSelectionModeRejected(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.POST("/api/internal/migration/users/query", QueryMigrationUsers)
-
- req := httptest.NewRequest(http.MethodPost, "/api/internal/migration/users/query", strings.NewReader(`{"page":0,"page_size":50,"selection_mode":"all","source_user_ids":[1]}`))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusBadRequest, w.Code)
- }
-
- func TestGetMigrationUser_RootExcluded(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.GET("/api/internal/migration/users/:id", GetMigrationUser)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 10000011,
- Username: "root",
- Password: "hash",
- AffCode: "ROOT1",
- Source: common.UserSourceLocal,
- Role: common.RoleRootUser,
- }).Error)
-
- req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/10000011", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusNotFound, w.Code)
- }
-
- func TestGetMigrationUser_SyncedUserStillReadableForVerify(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.GET("/api/internal/migration/users/:id", GetMigrationUser)
-
- db := setupInternalMigrationDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 10000021,
- Username: "synced-user",
- Password: "hash",
- AffCode: "SYNC1",
- Source: common.UserSourceSynced,
- RemoteUserId: 88,
- SyncedQuota: 123,
- Role: common.RoleCommonUser,
- }).Error)
-
- req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/10000021", 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))
- }
-
- func TestGetMigrationUser_InvalidID(t *testing.T) {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(testSyncAuthMiddleware())
- router.GET("/api/internal/migration/users/:id", GetMigrationUser)
-
- req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/not-a-number", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusBadRequest, w.Code)
- }
-
- 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 TestConvertMigrationUserToSynced_ConflictUsesSentinelErrors(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: 10000013,
- Username: "ov-13",
- Password: "hash",
- AffCode: "A13",
- Source: common.UserSourceSynced,
- RemoteUserId: 5001,
- }).Error)
-
- body := `{"remote_user_id":9999,"synced_quota":100}`
- req := httptest.NewRequest(http.MethodPost,
- "/api/internal/migration/users/10000013/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 TestConvertMigrationUserToSynced_RejectsRootUser(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: 10000012,
- Username: "root",
- Password: "hash",
- AffCode: "ROOT2",
- Source: common.UserSourceLocal,
- Role: common.RoleRootUser,
- }).Error)
-
- body := `{"remote_user_id":9999,"synced_quota":100}`
- req := httptest.NewRequest(http.MethodPost,
- "/api/internal/migration/users/10000012/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 TestConvertMigrationUserToSynced_InvalidRemoteUserID(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: 10000014,
- Username: "ov-14",
- Password: "hash",
- AffCode: "A14",
- Source: common.UserSourceLocal,
- }).Error)
-
- body := `{"remote_user_id":0,"synced_quota":100}`
- req := httptest.NewRequest(http.MethodPost,
- "/api/internal/migration/users/10000014/convert-to-synced",
- strings.NewReader(body))
- req.Header.Set("Content-Type", "application/json")
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusBadRequest, w.Code)
- }
-
- func TestConvertMigrationUserToSynced_PersistsSyncedFields(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: 10000015,
- Username: "ov-15",
- Password: "hash",
- AffCode: "A15",
- Source: common.UserSourceLocal,
- Quota: 88,
- }).Error)
-
- body := `{"remote_user_id":2015,"synced_quota":166}`
- req := httptest.NewRequest(http.MethodPost,
- "/api/internal/migration/users/10000015/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)
-
- var saved model.User
- require.NoError(t, db.First(&saved, 10000015).Error)
- require.Equal(t, common.UserSourceSynced, saved.Source)
- require.Equal(t, 2015, saved.RemoteUserId)
- require.Equal(t, 166, saved.SyncedQuota)
- require.Greater(t, saved.LastSyncAt, int64(0))
- }
-
- 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))
- }
|