package router import ( "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/controller" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/gin-contrib/sessions" "github.com/gin-contrib/sessions/cookie" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupUserMigrationRouterDB(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{}, &model.UserMigrationBatch{}, &model.UserMigrationItem{}, &model.MigrationQuotaGrant{}, )) return db } func setupUserMigrationRootRouter() *gin.Engine { gin.SetMode(gin.TestMode) r := gin.New() store := cookie.NewStore([]byte("test-session-secret")) store.Options(sessions.Options{ Path: "/", MaxAge: 3600, HttpOnly: true, Secure: false, }) r.Use(sessions.Sessions("session", store)) userMigrationRoute := r.Group("/api/user-migrations") userMigrationRoute.Use(func(c *gin.Context) { session := sessions.Default(c) session.Set("username", "root") session.Set("role", common.RoleRootUser) session.Set("id", 1) session.Set("status", common.UserStatusEnabled) session.Save() c.Request.Header.Set("New-Api-User", "1") c.Next() }, middleware.RootAuth()) { userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches) userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch) } return r } func setupUserMigrationAdminRouter() *gin.Engine { gin.SetMode(gin.TestMode) r := gin.New() store := cookie.NewStore([]byte("test-session-secret")) store.Options(sessions.Options{ Path: "/", MaxAge: 3600, HttpOnly: true, Secure: false, }) r.Use(sessions.Sessions("session", store)) userMigrationRoute := r.Group("/api/user-migrations") userMigrationRoute.Use(func(c *gin.Context) { session := sessions.Default(c) session.Set("username", "admin") session.Set("role", common.RoleAdminUser) session.Set("id", 2) session.Set("status", common.UserStatusEnabled) session.Save() c.Request.Header.Set("New-Api-User", "2") c.Next() }, middleware.RootAuth()) { userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches) userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch) } return r } func TestUserMigrationRoutes_RootCanAccess(t *testing.T) { db := setupUserMigrationRouterDB(t) orig := model.DB model.DB = db defer func() { model.DB = orig }() require.NoError(t, db.Create(&model.UserMigrationBatch{ Id: 1, Name: "router-batch", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1, Status: model.UserMigrationBatchStatusDraft, }).Error) router := setupUserMigrationRootRouter() req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil) w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusOK, w.Code) req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil) w = httptest.NewRecorder() router.ServeHTTP(w, req) require.NotEqual(t, http.StatusUnauthorized, w.Code) } func TestUserMigrationRoutes_AdminForbidden(t *testing.T) { db := setupUserMigrationRouterDB(t) orig := model.DB model.DB = db defer func() { model.DB = orig }() router := setupUserMigrationAdminRouter() req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil) w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusOK, w.Code) require.Contains(t, w.Body.String(), "success\":false") req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil) w = httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusOK, w.Code) require.Contains(t, w.Body.String(), "success\":false") }