Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

169 řádky
4.7 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. userMigrationRoute.POST("/batches/:id/cancel", controller.CancelUserMigrationBatch)
  60. }
  61. return r
  62. }
  63. func setupUserMigrationAdminRouter() *gin.Engine {
  64. gin.SetMode(gin.TestMode)
  65. r := gin.New()
  66. store := cookie.NewStore([]byte("test-session-secret"))
  67. store.Options(sessions.Options{
  68. Path: "/",
  69. MaxAge: 3600,
  70. HttpOnly: true,
  71. Secure: false,
  72. })
  73. r.Use(sessions.Sessions("session", store))
  74. userMigrationRoute := r.Group("/api/user-migrations")
  75. userMigrationRoute.Use(func(c *gin.Context) {
  76. session := sessions.Default(c)
  77. session.Set("username", "admin")
  78. session.Set("role", common.RoleAdminUser)
  79. session.Set("id", 2)
  80. session.Set("status", common.UserStatusEnabled)
  81. session.Save()
  82. c.Request.Header.Set("New-Api-User", "2")
  83. c.Next()
  84. }, middleware.RootAuth())
  85. {
  86. userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
  87. userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
  88. userMigrationRoute.POST("/batches/:id/cancel", controller.CancelUserMigrationBatch)
  89. }
  90. return r
  91. }
  92. func TestUserMigrationRoutes_RootCanAccess(t *testing.T) {
  93. db := setupUserMigrationRouterDB(t)
  94. orig := model.DB
  95. model.DB = db
  96. defer func() {
  97. model.DB = orig
  98. }()
  99. require.NoError(t, db.Create(&model.UserMigrationBatch{
  100. Id: 1,
  101. Name: "router-batch",
  102. SourceRegion: "ov",
  103. TargetRegion: "cn",
  104. OperatorId: 1,
  105. Status: model.UserMigrationBatchStatusDraft,
  106. }).Error)
  107. router := setupUserMigrationRootRouter()
  108. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  109. w := httptest.NewRecorder()
  110. router.ServeHTTP(w, req)
  111. require.Equal(t, http.StatusOK, w.Code)
  112. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  113. w = httptest.NewRecorder()
  114. router.ServeHTTP(w, req)
  115. require.NotEqual(t, http.StatusUnauthorized, w.Code)
  116. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  117. w = httptest.NewRecorder()
  118. router.ServeHTTP(w, req)
  119. require.NotEqual(t, http.StatusUnauthorized, w.Code)
  120. }
  121. func TestUserMigrationRoutes_AdminForbidden(t *testing.T) {
  122. db := setupUserMigrationRouterDB(t)
  123. orig := model.DB
  124. model.DB = db
  125. defer func() {
  126. model.DB = orig
  127. }()
  128. router := setupUserMigrationAdminRouter()
  129. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  130. w := httptest.NewRecorder()
  131. router.ServeHTTP(w, req)
  132. require.Equal(t, http.StatusOK, w.Code)
  133. require.Contains(t, w.Body.String(), "success\":false")
  134. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  135. w = httptest.NewRecorder()
  136. router.ServeHTTP(w, req)
  137. require.Equal(t, http.StatusOK, w.Code)
  138. require.Contains(t, w.Body.String(), "success\":false")
  139. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  140. w = httptest.NewRecorder()
  141. router.ServeHTTP(w, req)
  142. require.Equal(t, http.StatusOK, w.Code)
  143. require.Contains(t, w.Body.String(), "success\":false")
  144. }