Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

156 строки
4.1 KiB

  1. package router
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "testing"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/controller"
  8. "github.com/QuantumNous/new-api/middleware"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/gin-contrib/sessions"
  11. "github.com/gin-contrib/sessions/cookie"
  12. "github.com/gin-gonic/gin"
  13. "github.com/glebarez/sqlite"
  14. "github.com/stretchr/testify/require"
  15. "gorm.io/gorm"
  16. )
  17. func setupUserMigrationRouterDB(t *testing.T) *gorm.DB {
  18. t.Helper()
  19. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  20. require.NoError(t, err)
  21. sqlDB, err := db.DB()
  22. require.NoError(t, err)
  23. sqlDB.SetMaxOpenConns(1)
  24. require.NoError(t, db.AutoMigrate(
  25. &model.User{},
  26. &model.CustomOAuthProvider{},
  27. &model.UserOAuthBinding{},
  28. &model.UserMigrationBatch{},
  29. &model.UserMigrationItem{},
  30. &model.MigrationQuotaGrant{},
  31. ))
  32. return db
  33. }
  34. func setupUserMigrationRootRouter() *gin.Engine {
  35. gin.SetMode(gin.TestMode)
  36. r := gin.New()
  37. store := cookie.NewStore([]byte("test-session-secret"))
  38. store.Options(sessions.Options{
  39. Path: "/",
  40. MaxAge: 3600,
  41. HttpOnly: true,
  42. Secure: false,
  43. })
  44. r.Use(sessions.Sessions("session", store))
  45. userMigrationRoute := r.Group("/api/user-migrations")
  46. userMigrationRoute.Use(func(c *gin.Context) {
  47. session := sessions.Default(c)
  48. session.Set("username", "root")
  49. session.Set("role", common.RoleRootUser)
  50. session.Set("id", 1)
  51. session.Set("status", common.UserStatusEnabled)
  52. session.Save()
  53. c.Request.Header.Set("New-Api-User", "1")
  54. c.Next()
  55. }, middleware.RootAuth())
  56. {
  57. userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
  58. userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
  59. }
  60. return r
  61. }
  62. func setupUserMigrationAdminRouter() *gin.Engine {
  63. gin.SetMode(gin.TestMode)
  64. r := gin.New()
  65. store := cookie.NewStore([]byte("test-session-secret"))
  66. store.Options(sessions.Options{
  67. Path: "/",
  68. MaxAge: 3600,
  69. HttpOnly: true,
  70. Secure: false,
  71. })
  72. r.Use(sessions.Sessions("session", store))
  73. userMigrationRoute := r.Group("/api/user-migrations")
  74. userMigrationRoute.Use(func(c *gin.Context) {
  75. session := sessions.Default(c)
  76. session.Set("username", "admin")
  77. session.Set("role", common.RoleAdminUser)
  78. session.Set("id", 2)
  79. session.Set("status", common.UserStatusEnabled)
  80. session.Save()
  81. c.Request.Header.Set("New-Api-User", "2")
  82. c.Next()
  83. }, middleware.RootAuth())
  84. {
  85. userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
  86. userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
  87. }
  88. return r
  89. }
  90. func TestUserMigrationRoutes_RootCanAccess(t *testing.T) {
  91. db := setupUserMigrationRouterDB(t)
  92. orig := model.DB
  93. model.DB = db
  94. defer func() {
  95. model.DB = orig
  96. }()
  97. require.NoError(t, db.Create(&model.UserMigrationBatch{
  98. Id: 1,
  99. Name: "router-batch",
  100. SourceRegion: "ov",
  101. TargetRegion: "cn",
  102. OperatorId: 1,
  103. Status: model.UserMigrationBatchStatusDraft,
  104. }).Error)
  105. router := setupUserMigrationRootRouter()
  106. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  107. w := httptest.NewRecorder()
  108. router.ServeHTTP(w, req)
  109. require.Equal(t, http.StatusOK, w.Code)
  110. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  111. w = httptest.NewRecorder()
  112. router.ServeHTTP(w, req)
  113. require.NotEqual(t, http.StatusUnauthorized, w.Code)
  114. }
  115. func TestUserMigrationRoutes_AdminForbidden(t *testing.T) {
  116. db := setupUserMigrationRouterDB(t)
  117. orig := model.DB
  118. model.DB = db
  119. defer func() {
  120. model.DB = orig
  121. }()
  122. router := setupUserMigrationAdminRouter()
  123. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  124. w := httptest.NewRecorder()
  125. router.ServeHTTP(w, req)
  126. require.Equal(t, http.StatusOK, w.Code)
  127. require.Contains(t, w.Body.String(), "success\":false")
  128. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  129. w = httptest.NewRecorder()
  130. router.ServeHTTP(w, req)
  131. require.Equal(t, http.StatusOK, w.Code)
  132. require.Contains(t, w.Body.String(), "success\":false")
  133. }