|
- 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)
- }
|