|
- 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(), middleware.MasterNodeOnly())
- {
- userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
- userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
- userMigrationRoute.POST("/batches/:id/cancel", controller.CancelUserMigrationBatch)
- }
-
- 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(), middleware.MasterNodeOnly())
- {
- userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
- userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
- userMigrationRoute.POST("/batches/:id/cancel", controller.CancelUserMigrationBatch)
- }
-
- return r
- }
-
- func TestUserMigrationRoutes_RootCanAccess(t *testing.T) {
- db := setupUserMigrationRouterDB(t)
- orig := model.DB
- origMaster := common.IsMasterNode
- model.DB = db
- defer func() {
- model.DB = orig
- common.IsMasterNode = origMaster
- }()
- common.IsMasterNode = true
-
- 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)
-
- req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", 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
- origMaster := common.IsMasterNode
- model.DB = db
- defer func() {
- model.DB = orig
- common.IsMasterNode = origMaster
- }()
- common.IsMasterNode = true
-
- 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")
-
- req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
- w = httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusOK, w.Code)
- require.Contains(t, w.Body.String(), "success\":false")
- }
-
- func TestUserMigrationRoutes_RootForbiddenOnSlaveNode(t *testing.T) {
- db := setupUserMigrationRouterDB(t)
- orig := model.DB
- origMaster := common.IsMasterNode
- model.DB = db
- defer func() {
- model.DB = orig
- common.IsMasterNode = origMaster
- }()
- common.IsMasterNode = false
-
- router := setupUserMigrationRootRouter()
-
- req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusForbidden, w.Code)
- require.Contains(t, w.Body.String(), "only master node")
- }
|