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

138 строки
3.3 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/QuantumNous/new-api/model"
  9. "github.com/gin-gonic/gin"
  10. "github.com/glebarez/sqlite"
  11. "github.com/stretchr/testify/require"
  12. "gorm.io/gorm"
  13. )
  14. func setupUserMigrationControllerDB(t *testing.T) *gorm.DB {
  15. t.Helper()
  16. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  17. require.NoError(t, err)
  18. sqlDB, err := db.DB()
  19. require.NoError(t, err)
  20. sqlDB.SetMaxOpenConns(1)
  21. require.NoError(t, db.AutoMigrate(
  22. &model.User{},
  23. &model.UserMigrationBatch{},
  24. &model.UserMigrationItem{},
  25. &model.MigrationQuotaGrant{},
  26. ))
  27. return db
  28. }
  29. func testRootAuthMiddleware() gin.HandlerFunc {
  30. return func(c *gin.Context) {
  31. c.Set("id", 1)
  32. c.Set("role", 100)
  33. c.Next()
  34. }
  35. }
  36. func TestCreateUserMigrationBatch_Success(t *testing.T) {
  37. gin.SetMode(gin.TestMode)
  38. router := gin.New()
  39. router.Use(testRootAuthMiddleware())
  40. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  41. db := setupUserMigrationControllerDB(t)
  42. orig := model.DB
  43. model.DB = db
  44. defer func() {
  45. model.DB = orig
  46. }()
  47. body := map[string]any{
  48. "name": "wave-root-1",
  49. "source_region": "overseas",
  50. "target_region": "cn",
  51. }
  52. bodyBytes, err := json.Marshal(body)
  53. require.NoError(t, err)
  54. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBuffer(bodyBytes))
  55. req.Header.Set("Content-Type", "application/json")
  56. w := httptest.NewRecorder()
  57. router.ServeHTTP(w, req)
  58. require.Equal(t, http.StatusOK, w.Code)
  59. var resp map[string]any
  60. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  61. require.True(t, resp["success"].(bool))
  62. require.NotNil(t, resp["data"])
  63. }
  64. func TestExecuteUserMigrationBatch_AlreadyRunning(t *testing.T) {
  65. gin.SetMode(gin.TestMode)
  66. router := gin.New()
  67. router.Use(testRootAuthMiddleware())
  68. router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch)
  69. db := setupUserMigrationControllerDB(t)
  70. orig := model.DB
  71. model.DB = db
  72. defer func() {
  73. model.DB = orig
  74. }()
  75. batch := &model.UserMigrationBatch{
  76. Name: "b",
  77. SourceRegion: "ov",
  78. TargetRegion: "cn",
  79. OperatorId: 1,
  80. Status: model.UserMigrationBatchStatusRunning,
  81. }
  82. require.NoError(t, db.Create(batch).Error)
  83. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  84. w := httptest.NewRecorder()
  85. router.ServeHTTP(w, req)
  86. require.Equal(t, http.StatusConflict, w.Code)
  87. }
  88. func TestResolveUserMigrationItem_InvalidStrategy(t *testing.T) {
  89. gin.SetMode(gin.TestMode)
  90. router := gin.New()
  91. router.Use(testRootAuthMiddleware())
  92. router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem)
  93. db := setupUserMigrationControllerDB(t)
  94. orig := model.DB
  95. model.DB = db
  96. defer func() {
  97. model.DB = orig
  98. }()
  99. require.NoError(t, db.Create(&model.UserMigrationItem{
  100. BatchId: 1,
  101. SourceUserId: 10000001,
  102. SourceUsername: "u",
  103. Status: model.UserMigrationItemStatusConflict,
  104. }).Error)
  105. req := httptest.NewRequest(
  106. http.MethodPost,
  107. "/api/user-migrations/items/1/resolve",
  108. bytes.NewBufferString(`{"resolution_strategy":"invalid_strategy"}`),
  109. )
  110. req.Header.Set("Content-Type", "application/json")
  111. w := httptest.NewRecorder()
  112. router.ServeHTTP(w, req)
  113. require.Equal(t, http.StatusBadRequest, w.Code)
  114. }