package controller import ( "bytes" "encoding/json" "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/model" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupUserMigrationControllerDB(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.UserMigrationBatch{}, &model.UserMigrationItem{}, &model.MigrationQuotaGrant{}, )) return db } func testRootAuthMiddleware() gin.HandlerFunc { return func(c *gin.Context) { c.Set("id", 1) c.Set("role", 100) c.Next() } } func TestCreateUserMigrationBatch_Success(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() router.Use(testRootAuthMiddleware()) router.POST("/api/user-migrations/batches", CreateUserMigrationBatch) db := setupUserMigrationControllerDB(t) orig := model.DB model.DB = db defer func() { model.DB = orig }() body := map[string]any{ "name": "wave-root-1", "source_region": "overseas", "target_region": "cn", } bodyBytes, err := json.Marshal(body) require.NoError(t, err) req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", 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)) require.NotNil(t, resp["data"]) } func TestExecuteUserMigrationBatch_AlreadyRunning(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() router.Use(testRootAuthMiddleware()) router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch) db := setupUserMigrationControllerDB(t) orig := model.DB model.DB = db defer func() { model.DB = orig }() batch := &model.UserMigrationBatch{ Name: "b", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1, Status: model.UserMigrationBatchStatusRunning, } require.NoError(t, db.Create(batch).Error) req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil) w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusConflict, w.Code) } func TestResolveUserMigrationItem_InvalidStrategy(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() router.Use(testRootAuthMiddleware()) router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem) db := setupUserMigrationControllerDB(t) orig := model.DB model.DB = db defer func() { model.DB = orig }() require.NoError(t, db.Create(&model.UserMigrationItem{ BatchId: 1, SourceUserId: 10000001, SourceUsername: "u", Status: model.UserMigrationItemStatusConflict, }).Error) req := httptest.NewRequest( http.MethodPost, "/api/user-migrations/items/1/resolve", bytes.NewBufferString(`{"resolution_strategy":"invalid_strategy"}`), ) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusBadRequest, w.Code) } func TestResolveUserMigrationItem_RefreshesBatchStatusToReady(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() router.Use(testRootAuthMiddleware()) router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem) db := setupUserMigrationControllerDB(t) orig := model.DB model.DB = db defer func() { model.DB = orig }() batch := &model.UserMigrationBatch{ Name: "resolve-ready", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1, Status: model.UserMigrationBatchStatusScanned, } require.NoError(t, model.CreateUserMigrationBatch(batch)) require.NoError(t, db.Create(&model.UserMigrationItem{ BatchId: batch.Id, SourceUserId: 1001, SourceUsername: "ov-u1", Status: model.UserMigrationItemStatusConflict, }).Error) req := httptest.NewRequest( http.MethodPost, "/api/user-migrations/items/1/resolve", bytes.NewBufferString(`{"resolution_strategy":"create_new"}`), ) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusOK, w.Code) var refreshed model.UserMigrationBatch require.NoError(t, db.First(&refreshed, batch.Id).Error) require.Equal(t, model.UserMigrationBatchStatusReady, refreshed.Status) }