You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

471 lines
13 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "errors"
  6. "net/http"
  7. "net/http/httptest"
  8. "testing"
  9. "github.com/QuantumNous/new-api/model"
  10. user_migration "github.com/QuantumNous/new-api/service/user_migration"
  11. "github.com/gin-gonic/gin"
  12. "github.com/glebarez/sqlite"
  13. "github.com/stretchr/testify/require"
  14. "gorm.io/gorm"
  15. )
  16. func setupUserMigrationControllerDB(t *testing.T) *gorm.DB {
  17. t.Helper()
  18. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  19. require.NoError(t, err)
  20. sqlDB, err := db.DB()
  21. require.NoError(t, err)
  22. sqlDB.SetMaxOpenConns(1)
  23. require.NoError(t, db.AutoMigrate(
  24. &model.User{},
  25. &model.CustomOAuthProvider{},
  26. &model.UserOAuthBinding{},
  27. &model.UserMigrationBatch{},
  28. &model.UserMigrationItem{},
  29. &model.MigrationQuotaGrant{},
  30. ))
  31. return db
  32. }
  33. func testRootAuthMiddleware() gin.HandlerFunc {
  34. return func(c *gin.Context) {
  35. c.Set("id", 1)
  36. c.Set("role", 100)
  37. c.Next()
  38. }
  39. }
  40. type fakeMigrationVerifier struct {
  41. results map[int]*user_migration.VerifyResult
  42. errs map[int]error
  43. }
  44. func (f *fakeMigrationVerifier) VerifyItem(item *model.UserMigrationItem) (*user_migration.VerifyResult, error) {
  45. if err, ok := f.errs[item.SourceUserId]; ok {
  46. return nil, err
  47. }
  48. if result, ok := f.results[item.SourceUserId]; ok {
  49. return result, nil
  50. }
  51. return nil, errors.New("unexpected source user")
  52. }
  53. func TestCreateUserMigrationBatch_Success(t *testing.T) {
  54. gin.SetMode(gin.TestMode)
  55. router := gin.New()
  56. router.Use(testRootAuthMiddleware())
  57. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  58. db := setupUserMigrationControllerDB(t)
  59. orig := model.DB
  60. model.DB = db
  61. defer func() {
  62. model.DB = orig
  63. }()
  64. body := map[string]any{
  65. "name": "wave-root-1",
  66. "source_region": "overseas",
  67. "target_region": "cn",
  68. }
  69. bodyBytes, err := json.Marshal(body)
  70. require.NoError(t, err)
  71. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBuffer(bodyBytes))
  72. req.Header.Set("Content-Type", "application/json")
  73. w := httptest.NewRecorder()
  74. router.ServeHTTP(w, req)
  75. require.Equal(t, http.StatusOK, w.Code)
  76. var resp map[string]any
  77. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  78. require.True(t, resp["success"].(bool))
  79. require.NotNil(t, resp["data"])
  80. }
  81. func TestListUserMigrationBatches_Success(t *testing.T) {
  82. gin.SetMode(gin.TestMode)
  83. router := gin.New()
  84. router.Use(testRootAuthMiddleware())
  85. router.GET("/api/user-migrations/batches", ListUserMigrationBatches)
  86. db := setupUserMigrationControllerDB(t)
  87. orig := model.DB
  88. model.DB = db
  89. defer func() {
  90. model.DB = orig
  91. }()
  92. require.NoError(t, db.Create(&model.UserMigrationBatch{
  93. Id: 1,
  94. Name: "batch-1",
  95. SourceRegion: "ov",
  96. TargetRegion: "cn",
  97. OperatorId: 1,
  98. }).Error)
  99. require.NoError(t, db.Create(&model.UserMigrationBatch{
  100. Id: 2,
  101. Name: "batch-2",
  102. SourceRegion: "ov",
  103. TargetRegion: "cn",
  104. OperatorId: 1,
  105. }).Error)
  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. var resp map[string]any
  111. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  112. require.True(t, resp["success"].(bool))
  113. require.Len(t, resp["data"], 2)
  114. }
  115. func TestGetUserMigrationBatch_Success(t *testing.T) {
  116. gin.SetMode(gin.TestMode)
  117. router := gin.New()
  118. router.Use(testRootAuthMiddleware())
  119. router.GET("/api/user-migrations/batches/:id", GetUserMigrationBatch)
  120. db := setupUserMigrationControllerDB(t)
  121. orig := model.DB
  122. model.DB = db
  123. defer func() {
  124. model.DB = orig
  125. }()
  126. require.NoError(t, db.Create(&model.UserMigrationBatch{
  127. Id: 1,
  128. Name: "detail-batch",
  129. SourceRegion: "ov",
  130. TargetRegion: "cn",
  131. OperatorId: 1,
  132. }).Error)
  133. require.NoError(t, db.Create(&model.UserMigrationItem{
  134. BatchId: 1,
  135. SourceUserId: 1001,
  136. SourceUsername: "ov-u1",
  137. Status: model.UserMigrationItemStatusReady,
  138. }).Error)
  139. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches/1?page=0&page_size=50", nil)
  140. w := httptest.NewRecorder()
  141. router.ServeHTTP(w, req)
  142. require.Equal(t, http.StatusOK, w.Code)
  143. var resp map[string]any
  144. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  145. require.True(t, resp["success"].(bool))
  146. require.NotNil(t, resp["data"])
  147. require.Len(t, resp["items"], 1)
  148. require.EqualValues(t, 1, resp["total"])
  149. }
  150. func TestGetUserMigrationBatch_NotFound(t *testing.T) {
  151. gin.SetMode(gin.TestMode)
  152. router := gin.New()
  153. router.Use(testRootAuthMiddleware())
  154. router.GET("/api/user-migrations/batches/:id", GetUserMigrationBatch)
  155. db := setupUserMigrationControllerDB(t)
  156. orig := model.DB
  157. model.DB = db
  158. defer func() {
  159. model.DB = orig
  160. }()
  161. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches/999", nil)
  162. w := httptest.NewRecorder()
  163. router.ServeHTTP(w, req)
  164. require.Equal(t, http.StatusNotFound, w.Code)
  165. }
  166. func TestExecuteUserMigrationBatch_AlreadyRunning(t *testing.T) {
  167. gin.SetMode(gin.TestMode)
  168. router := gin.New()
  169. router.Use(testRootAuthMiddleware())
  170. router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch)
  171. db := setupUserMigrationControllerDB(t)
  172. orig := model.DB
  173. model.DB = db
  174. defer func() {
  175. model.DB = orig
  176. }()
  177. batch := &model.UserMigrationBatch{
  178. Name: "b",
  179. SourceRegion: "ov",
  180. TargetRegion: "cn",
  181. OperatorId: 1,
  182. Status: model.UserMigrationBatchStatusRunning,
  183. }
  184. require.NoError(t, db.Create(batch).Error)
  185. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  186. w := httptest.NewRecorder()
  187. router.ServeHTTP(w, req)
  188. require.Equal(t, http.StatusConflict, w.Code)
  189. }
  190. func TestExecuteUserMigrationBatch_InvalidStateRejected(t *testing.T) {
  191. gin.SetMode(gin.TestMode)
  192. router := gin.New()
  193. router.Use(testRootAuthMiddleware())
  194. router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch)
  195. db := setupUserMigrationControllerDB(t)
  196. orig := model.DB
  197. model.DB = db
  198. defer func() {
  199. model.DB = orig
  200. }()
  201. for _, status := range []string{model.UserMigrationBatchStatusDraft, model.UserMigrationBatchStatusScanned} {
  202. require.NoError(t, db.Exec("DELETE FROM user_migration_batches").Error)
  203. require.NoError(t, db.Create(&model.UserMigrationBatch{
  204. Id: 1,
  205. Name: "invalid-exec",
  206. SourceRegion: "ov",
  207. TargetRegion: "cn",
  208. OperatorId: 1,
  209. Status: status,
  210. }).Error)
  211. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  212. w := httptest.NewRecorder()
  213. router.ServeHTTP(w, req)
  214. require.Equal(t, http.StatusConflict, w.Code)
  215. }
  216. }
  217. func TestResolveUserMigrationItem_InvalidStrategy(t *testing.T) {
  218. gin.SetMode(gin.TestMode)
  219. router := gin.New()
  220. router.Use(testRootAuthMiddleware())
  221. router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem)
  222. db := setupUserMigrationControllerDB(t)
  223. orig := model.DB
  224. model.DB = db
  225. defer func() {
  226. model.DB = orig
  227. }()
  228. require.NoError(t, db.Create(&model.UserMigrationItem{
  229. BatchId: 1,
  230. SourceUserId: 10000001,
  231. SourceUsername: "u",
  232. Status: model.UserMigrationItemStatusConflict,
  233. }).Error)
  234. req := httptest.NewRequest(
  235. http.MethodPost,
  236. "/api/user-migrations/items/1/resolve",
  237. bytes.NewBufferString(`{"resolution_strategy":"invalid_strategy"}`),
  238. )
  239. req.Header.Set("Content-Type", "application/json")
  240. w := httptest.NewRecorder()
  241. router.ServeHTTP(w, req)
  242. require.Equal(t, http.StatusBadRequest, w.Code)
  243. }
  244. func TestResolveUserMigrationItem_RefreshesBatchStatusToReady(t *testing.T) {
  245. gin.SetMode(gin.TestMode)
  246. router := gin.New()
  247. router.Use(testRootAuthMiddleware())
  248. router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem)
  249. db := setupUserMigrationControllerDB(t)
  250. orig := model.DB
  251. model.DB = db
  252. defer func() {
  253. model.DB = orig
  254. }()
  255. batch := &model.UserMigrationBatch{
  256. Name: "resolve-ready",
  257. SourceRegion: "ov",
  258. TargetRegion: "cn",
  259. OperatorId: 1,
  260. Status: model.UserMigrationBatchStatusScanned,
  261. }
  262. require.NoError(t, model.CreateUserMigrationBatch(batch))
  263. require.NoError(t, db.Create(&model.UserMigrationItem{
  264. BatchId: batch.Id,
  265. SourceUserId: 1001,
  266. SourceUsername: "ov-u1",
  267. Status: model.UserMigrationItemStatusConflict,
  268. }).Error)
  269. req := httptest.NewRequest(
  270. http.MethodPost,
  271. "/api/user-migrations/items/1/resolve",
  272. bytes.NewBufferString(`{"resolution_strategy":"create_new"}`),
  273. )
  274. req.Header.Set("Content-Type", "application/json")
  275. w := httptest.NewRecorder()
  276. router.ServeHTTP(w, req)
  277. require.Equal(t, http.StatusOK, w.Code)
  278. var refreshed model.UserMigrationBatch
  279. require.NoError(t, db.First(&refreshed, batch.Id).Error)
  280. require.Equal(t, model.UserMigrationBatchStatusReady, refreshed.Status)
  281. }
  282. func TestRetryUserMigrationBatch_OnlyFailedItemsReset(t *testing.T) {
  283. gin.SetMode(gin.TestMode)
  284. router := gin.New()
  285. router.Use(testRootAuthMiddleware())
  286. router.POST("/api/user-migrations/batches/:id/retry", RetryUserMigrationBatch)
  287. db := setupUserMigrationControllerDB(t)
  288. orig := model.DB
  289. model.DB = db
  290. defer func() {
  291. model.DB = orig
  292. }()
  293. require.NoError(t, db.Create(&model.UserMigrationBatch{
  294. Id: 1,
  295. Name: "retry-batch",
  296. SourceRegion: "ov",
  297. TargetRegion: "cn",
  298. OperatorId: 1,
  299. Status: model.UserMigrationBatchStatusRunning,
  300. }).Error)
  301. require.NoError(t, db.Create(&model.UserMigrationItem{
  302. Id: 1,
  303. BatchId: 1,
  304. SourceUserId: 1001,
  305. SourceUsername: "failed-user",
  306. Status: model.UserMigrationItemStatusFailed,
  307. ErrorMessage: "boom",
  308. }).Error)
  309. require.NoError(t, db.Create(&model.UserMigrationItem{
  310. Id: 2,
  311. BatchId: 1,
  312. SourceUserId: 1002,
  313. SourceUsername: "migrated-user",
  314. Status: model.UserMigrationItemStatusMigrated,
  315. }).Error)
  316. require.NoError(t, db.Create(&model.UserMigrationItem{
  317. Id: 3,
  318. BatchId: 1,
  319. SourceUserId: 1003,
  320. SourceUsername: "skipped-user",
  321. Status: model.UserMigrationItemStatusSkipped,
  322. }).Error)
  323. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/retry", nil)
  324. w := httptest.NewRecorder()
  325. router.ServeHTTP(w, req)
  326. require.Equal(t, http.StatusConflict, w.Code)
  327. var failedItem, migratedItem, skippedItem model.UserMigrationItem
  328. require.NoError(t, db.First(&failedItem, 1).Error)
  329. require.NoError(t, db.First(&migratedItem, 2).Error)
  330. require.NoError(t, db.First(&skippedItem, 3).Error)
  331. require.Equal(t, model.UserMigrationItemStatusReady, failedItem.Status)
  332. require.Empty(t, failedItem.ErrorMessage)
  333. require.Equal(t, model.UserMigrationItemStatusMigrated, migratedItem.Status)
  334. require.Equal(t, model.UserMigrationItemStatusSkipped, skippedItem.Status)
  335. }
  336. func TestVerifyUserMigrationBatch_AggregatesResult(t *testing.T) {
  337. gin.SetMode(gin.TestMode)
  338. router := gin.New()
  339. router.Use(testRootAuthMiddleware())
  340. router.POST("/api/user-migrations/batches/:id/verify", VerifyUserMigrationBatch)
  341. db := setupUserMigrationControllerDB(t)
  342. orig := model.DB
  343. origFactory := newMigrationVerifier
  344. model.DB = db
  345. defer func() {
  346. model.DB = orig
  347. newMigrationVerifier = origFactory
  348. }()
  349. require.NoError(t, db.Create(&model.UserMigrationItem{
  350. Id: 1,
  351. BatchId: 1,
  352. SourceUserId: 1001,
  353. TargetUserId: 1,
  354. Status: model.UserMigrationItemStatusMigrated,
  355. }).Error)
  356. require.NoError(t, db.Create(&model.UserMigrationItem{
  357. Id: 2,
  358. BatchId: 1,
  359. SourceUserId: 1002,
  360. TargetUserId: 2,
  361. Status: model.UserMigrationItemStatusMigrated,
  362. }).Error)
  363. require.NoError(t, db.Create(&model.UserMigrationItem{
  364. Id: 3,
  365. BatchId: 1,
  366. SourceUserId: 1003,
  367. TargetUserId: 3,
  368. Status: model.UserMigrationItemStatusMigrated,
  369. }).Error)
  370. newMigrationVerifier = func() migrationVerifier {
  371. return &fakeMigrationVerifier{
  372. results: map[int]*user_migration.VerifyResult{
  373. 1001: {
  374. TargetUserExists: true,
  375. RemoteConverted: true,
  376. QuotaMatched: true,
  377. },
  378. 1002: {
  379. TargetUserExists: false,
  380. RemoteConverted: true,
  381. QuotaMatched: false,
  382. },
  383. },
  384. errs: map[int]error{
  385. 1003: errors.New("remote user not found"),
  386. },
  387. }
  388. }
  389. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/verify", nil)
  390. w := httptest.NewRecorder()
  391. router.ServeHTTP(w, req)
  392. require.Equal(t, http.StatusOK, w.Code)
  393. var resp struct {
  394. Success bool `json:"success"`
  395. Data []map[string]any `json:"data"`
  396. }
  397. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  398. require.True(t, resp.Success)
  399. require.Len(t, resp.Data, 3)
  400. bySource := make(map[int]map[string]any, len(resp.Data))
  401. for _, entry := range resp.Data {
  402. bySource[int(entry["source_user_id"].(float64))] = entry
  403. }
  404. require.NotNil(t, bySource[1001]["result"])
  405. require.NotNil(t, bySource[1002]["result"])
  406. require.Equal(t, "remote user not found", bySource[1003]["error"])
  407. }