No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.
 
 
 

800 líneas
23 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 fakeMigrationCandidateClient struct {
  41. resp *user_migration.ListRemoteUsersResponse
  42. err error
  43. gotKeyword string
  44. gotPage int
  45. gotPageSize int
  46. }
  47. func (f *fakeMigrationCandidateClient) ListUsers(page, pageSize int, keyword string) (*user_migration.ListRemoteUsersResponse, error) {
  48. f.gotKeyword = keyword
  49. f.gotPage = page
  50. f.gotPageSize = pageSize
  51. if f.err != nil {
  52. return nil, f.err
  53. }
  54. return f.resp, nil
  55. }
  56. type fakeMigrationVerifier struct {
  57. results map[int]*user_migration.VerifyResult
  58. errs map[int]error
  59. }
  60. func (f *fakeMigrationVerifier) VerifyItem(item *model.UserMigrationItem) (*user_migration.VerifyResult, error) {
  61. if err, ok := f.errs[item.SourceUserId]; ok {
  62. return nil, err
  63. }
  64. if result, ok := f.results[item.SourceUserId]; ok {
  65. return result, nil
  66. }
  67. return nil, errors.New("unexpected source user")
  68. }
  69. func TestCreateUserMigrationBatch_Success(t *testing.T) {
  70. gin.SetMode(gin.TestMode)
  71. router := gin.New()
  72. router.Use(testRootAuthMiddleware())
  73. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  74. db := setupUserMigrationControllerDB(t)
  75. orig := model.DB
  76. model.DB = db
  77. defer func() {
  78. model.DB = orig
  79. }()
  80. body := map[string]any{
  81. "name": "wave-root-1",
  82. "source_region": "overseas",
  83. "target_region": "cn",
  84. }
  85. bodyBytes, err := json.Marshal(body)
  86. require.NoError(t, err)
  87. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBuffer(bodyBytes))
  88. req.Header.Set("Content-Type", "application/json")
  89. w := httptest.NewRecorder()
  90. router.ServeHTTP(w, req)
  91. require.Equal(t, http.StatusOK, w.Code)
  92. var resp map[string]any
  93. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  94. require.True(t, resp["success"].(bool))
  95. require.NotNil(t, resp["data"])
  96. }
  97. func TestCreateUserMigrationBatch_ExplicitIDsSuccess(t *testing.T) {
  98. gin.SetMode(gin.TestMode)
  99. router := gin.New()
  100. router.Use(testRootAuthMiddleware())
  101. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  102. db := setupUserMigrationControllerDB(t)
  103. orig := model.DB
  104. model.DB = db
  105. defer func() {
  106. model.DB = orig
  107. }()
  108. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBufferString(`{
  109. "name":"explicit-batch",
  110. "source_region":"ov",
  111. "target_region":"cn",
  112. "selection_mode":"explicit_ids",
  113. "source_user_ids":[1001,1002,1002]
  114. }`))
  115. req.Header.Set("Content-Type", "application/json")
  116. w := httptest.NewRecorder()
  117. router.ServeHTTP(w, req)
  118. require.Equal(t, http.StatusOK, w.Code)
  119. var batch model.UserMigrationBatch
  120. require.NoError(t, db.First(&batch, 1).Error)
  121. require.Equal(t, "explicit_ids", batch.SelectionMode)
  122. require.Equal(t, 2, batch.RequestedUserCount)
  123. }
  124. func TestCreateUserMigrationBatch_ExplicitIDsValidation(t *testing.T) {
  125. gin.SetMode(gin.TestMode)
  126. router := gin.New()
  127. router.Use(testRootAuthMiddleware())
  128. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  129. cases := []string{
  130. `{"name":"b1","source_region":"ov","target_region":"cn","selection_mode":"explicit_ids","source_user_ids":[]}`,
  131. `{"name":"b2","source_region":"ov","target_region":"cn","selection_mode":"explicit_ids","source_user_ids":[0]}`,
  132. `{"name":"b3","source_region":"ov","target_region":"cn","selection_mode":"invalid","source_user_ids":[1]}`,
  133. }
  134. db := setupUserMigrationControllerDB(t)
  135. orig := model.DB
  136. model.DB = db
  137. defer func() {
  138. model.DB = orig
  139. }()
  140. for _, body := range cases {
  141. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBufferString(body))
  142. req.Header.Set("Content-Type", "application/json")
  143. w := httptest.NewRecorder()
  144. router.ServeHTTP(w, req)
  145. require.Equal(t, http.StatusBadRequest, w.Code)
  146. }
  147. }
  148. func TestCreateUserMigrationBatch_ExplicitIDs_RejectsActiveBatchOverlap(t *testing.T) {
  149. gin.SetMode(gin.TestMode)
  150. router := gin.New()
  151. router.Use(testRootAuthMiddleware())
  152. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  153. db := setupUserMigrationControllerDB(t)
  154. orig := model.DB
  155. model.DB = db
  156. defer func() {
  157. model.DB = orig
  158. }()
  159. existingBatch := &model.UserMigrationBatch{
  160. Name: "existing-draft-batch",
  161. SourceRegion: "ov",
  162. TargetRegion: "cn",
  163. SelectionMode: "explicit_ids",
  164. SelectionPayload: `{"source_user_ids":[1002]}`,
  165. Status: model.UserMigrationBatchStatusDraft,
  166. OperatorId: 1,
  167. }
  168. require.NoError(t, db.Create(existingBatch).Error)
  169. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBufferString(`{
  170. "name":"explicit-batch",
  171. "source_region":"ov",
  172. "target_region":"cn",
  173. "selection_mode":"explicit_ids",
  174. "source_user_ids":[1001,1002]
  175. }`))
  176. req.Header.Set("Content-Type", "application/json")
  177. w := httptest.NewRecorder()
  178. router.ServeHTTP(w, req)
  179. require.Equal(t, http.StatusConflict, w.Code)
  180. var resp struct {
  181. Success bool `json:"success"`
  182. Error string `json:"error"`
  183. Conflicts []struct {
  184. SourceUserID int `json:"source_user_id"`
  185. BatchID int `json:"batch_id"`
  186. BatchStatus string `json:"batch_status"`
  187. } `json:"conflicts"`
  188. }
  189. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  190. require.False(t, resp.Success)
  191. require.Equal(t, "some selected users already exist in active migration batches", resp.Error)
  192. require.Len(t, resp.Conflicts, 1)
  193. require.Equal(t, 1002, resp.Conflicts[0].SourceUserID)
  194. require.Equal(t, existingBatch.Id, resp.Conflicts[0].BatchID)
  195. require.Equal(t, model.UserMigrationBatchStatusDraft, resp.Conflicts[0].BatchStatus)
  196. }
  197. func TestListMigrationCandidateUsers_Success(t *testing.T) {
  198. gin.SetMode(gin.TestMode)
  199. router := gin.New()
  200. router.Use(testRootAuthMiddleware())
  201. origFactory := newMigrationCandidateClient
  202. defer func() {
  203. newMigrationCandidateClient = origFactory
  204. }()
  205. newMigrationCandidateClient = func() migrationCandidateClient {
  206. return &fakeMigrationCandidateClient{
  207. resp: &user_migration.ListRemoteUsersResponse{
  208. Success: true,
  209. Data: []user_migration.RemoteUserSnapshot{
  210. {Id: 10000001, Username: "ov-a", Email: "a@example.com", Quota: 100},
  211. },
  212. Total: 1,
  213. },
  214. }
  215. }
  216. router.GET("/api/user-migrations/candidate-users", ListMigrationCandidateUsers)
  217. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/candidate-users?page=0&page_size=50", nil)
  218. w := httptest.NewRecorder()
  219. router.ServeHTTP(w, req)
  220. require.Equal(t, http.StatusOK, w.Code)
  221. var resp map[string]any
  222. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  223. require.True(t, resp["success"].(bool))
  224. require.EqualValues(t, 1, resp["total"])
  225. }
  226. func TestListMigrationCandidateUsers_UpstreamFailure(t *testing.T) {
  227. gin.SetMode(gin.TestMode)
  228. router := gin.New()
  229. router.Use(testRootAuthMiddleware())
  230. origFactory := newMigrationCandidateClient
  231. defer func() {
  232. newMigrationCandidateClient = origFactory
  233. }()
  234. newMigrationCandidateClient = func() migrationCandidateClient {
  235. return &fakeMigrationCandidateClient{err: errors.New("upstream unavailable")}
  236. }
  237. router.GET("/api/user-migrations/candidate-users", ListMigrationCandidateUsers)
  238. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/candidate-users?page=0&page_size=50", nil)
  239. w := httptest.NewRecorder()
  240. router.ServeHTTP(w, req)
  241. require.Equal(t, http.StatusInternalServerError, w.Code)
  242. }
  243. func TestListMigrationCandidateUsers_ForwardsKeyword(t *testing.T) {
  244. gin.SetMode(gin.TestMode)
  245. router := gin.New()
  246. router.Use(testRootAuthMiddleware())
  247. client := &fakeMigrationCandidateClient{
  248. resp: &user_migration.ListRemoteUsersResponse{
  249. Success: true,
  250. Data: []user_migration.RemoteUserSnapshot{},
  251. Total: 0,
  252. },
  253. }
  254. origFactory := newMigrationCandidateClient
  255. defer func() {
  256. newMigrationCandidateClient = origFactory
  257. }()
  258. newMigrationCandidateClient = func() migrationCandidateClient {
  259. return client
  260. }
  261. router.GET("/api/user-migrations/candidate-users", ListMigrationCandidateUsers)
  262. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/candidate-users?page=2&page_size=30&keyword=alice", nil)
  263. w := httptest.NewRecorder()
  264. router.ServeHTTP(w, req)
  265. require.Equal(t, http.StatusOK, w.Code)
  266. require.Equal(t, 2, client.gotPage)
  267. require.Equal(t, 30, client.gotPageSize)
  268. require.Equal(t, "alice", client.gotKeyword)
  269. }
  270. func TestListUserMigrationBatches_Success(t *testing.T) {
  271. gin.SetMode(gin.TestMode)
  272. router := gin.New()
  273. router.Use(testRootAuthMiddleware())
  274. router.GET("/api/user-migrations/batches", ListUserMigrationBatches)
  275. db := setupUserMigrationControllerDB(t)
  276. orig := model.DB
  277. model.DB = db
  278. defer func() {
  279. model.DB = orig
  280. }()
  281. require.NoError(t, db.Create(&model.UserMigrationBatch{
  282. Id: 1,
  283. Name: "batch-1",
  284. SourceRegion: "ov",
  285. TargetRegion: "cn",
  286. OperatorId: 1,
  287. }).Error)
  288. require.NoError(t, db.Create(&model.UserMigrationBatch{
  289. Id: 2,
  290. Name: "batch-2",
  291. SourceRegion: "ov",
  292. TargetRegion: "cn",
  293. OperatorId: 1,
  294. }).Error)
  295. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  296. w := httptest.NewRecorder()
  297. router.ServeHTTP(w, req)
  298. require.Equal(t, http.StatusOK, w.Code)
  299. var resp map[string]any
  300. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  301. require.True(t, resp["success"].(bool))
  302. require.Len(t, resp["data"], 2)
  303. }
  304. func TestGetUserMigrationBatch_Success(t *testing.T) {
  305. gin.SetMode(gin.TestMode)
  306. router := gin.New()
  307. router.Use(testRootAuthMiddleware())
  308. router.GET("/api/user-migrations/batches/:id", GetUserMigrationBatch)
  309. db := setupUserMigrationControllerDB(t)
  310. orig := model.DB
  311. model.DB = db
  312. defer func() {
  313. model.DB = orig
  314. }()
  315. require.NoError(t, db.Create(&model.UserMigrationBatch{
  316. Id: 1,
  317. Name: "detail-batch",
  318. SourceRegion: "ov",
  319. TargetRegion: "cn",
  320. OperatorId: 1,
  321. }).Error)
  322. require.NoError(t, db.Create(&model.UserMigrationItem{
  323. BatchId: 1,
  324. SourceUserId: 1001,
  325. SourceUsername: "ov-u1",
  326. Status: model.UserMigrationItemStatusReady,
  327. }).Error)
  328. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches/1?page=0&page_size=50", nil)
  329. w := httptest.NewRecorder()
  330. router.ServeHTTP(w, req)
  331. require.Equal(t, http.StatusOK, w.Code)
  332. var resp map[string]any
  333. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  334. require.True(t, resp["success"].(bool))
  335. require.NotNil(t, resp["data"])
  336. require.Len(t, resp["items"], 1)
  337. require.EqualValues(t, 1, resp["total"])
  338. }
  339. func TestGetUserMigrationBatch_NotFound(t *testing.T) {
  340. gin.SetMode(gin.TestMode)
  341. router := gin.New()
  342. router.Use(testRootAuthMiddleware())
  343. router.GET("/api/user-migrations/batches/:id", GetUserMigrationBatch)
  344. db := setupUserMigrationControllerDB(t)
  345. orig := model.DB
  346. model.DB = db
  347. defer func() {
  348. model.DB = orig
  349. }()
  350. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches/999", nil)
  351. w := httptest.NewRecorder()
  352. router.ServeHTTP(w, req)
  353. require.Equal(t, http.StatusNotFound, w.Code)
  354. }
  355. func TestExecuteUserMigrationBatch_AlreadyRunning(t *testing.T) {
  356. gin.SetMode(gin.TestMode)
  357. router := gin.New()
  358. router.Use(testRootAuthMiddleware())
  359. router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch)
  360. db := setupUserMigrationControllerDB(t)
  361. orig := model.DB
  362. model.DB = db
  363. defer func() {
  364. model.DB = orig
  365. }()
  366. batch := &model.UserMigrationBatch{
  367. Name: "b",
  368. SourceRegion: "ov",
  369. TargetRegion: "cn",
  370. OperatorId: 1,
  371. Status: model.UserMigrationBatchStatusRunning,
  372. }
  373. require.NoError(t, db.Create(batch).Error)
  374. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  375. w := httptest.NewRecorder()
  376. router.ServeHTTP(w, req)
  377. require.Equal(t, http.StatusConflict, w.Code)
  378. }
  379. func TestExecuteUserMigrationBatch_InvalidStateRejected(t *testing.T) {
  380. gin.SetMode(gin.TestMode)
  381. router := gin.New()
  382. router.Use(testRootAuthMiddleware())
  383. router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch)
  384. db := setupUserMigrationControllerDB(t)
  385. orig := model.DB
  386. model.DB = db
  387. defer func() {
  388. model.DB = orig
  389. }()
  390. for _, status := range []string{model.UserMigrationBatchStatusDraft, model.UserMigrationBatchStatusScanned} {
  391. require.NoError(t, db.Exec("DELETE FROM user_migration_batches").Error)
  392. require.NoError(t, db.Create(&model.UserMigrationBatch{
  393. Id: 1,
  394. Name: "invalid-exec",
  395. SourceRegion: "ov",
  396. TargetRegion: "cn",
  397. OperatorId: 1,
  398. Status: status,
  399. }).Error)
  400. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  401. w := httptest.NewRecorder()
  402. router.ServeHTTP(w, req)
  403. require.Equal(t, http.StatusConflict, w.Code)
  404. }
  405. }
  406. func TestResolveUserMigrationItem_InvalidStrategy(t *testing.T) {
  407. gin.SetMode(gin.TestMode)
  408. router := gin.New()
  409. router.Use(testRootAuthMiddleware())
  410. router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem)
  411. db := setupUserMigrationControllerDB(t)
  412. orig := model.DB
  413. model.DB = db
  414. defer func() {
  415. model.DB = orig
  416. }()
  417. require.NoError(t, db.Create(&model.UserMigrationItem{
  418. BatchId: 1,
  419. SourceUserId: 10000001,
  420. SourceUsername: "u",
  421. Status: model.UserMigrationItemStatusConflict,
  422. }).Error)
  423. req := httptest.NewRequest(
  424. http.MethodPost,
  425. "/api/user-migrations/items/1/resolve",
  426. bytes.NewBufferString(`{"resolution_strategy":"invalid_strategy"}`),
  427. )
  428. req.Header.Set("Content-Type", "application/json")
  429. w := httptest.NewRecorder()
  430. router.ServeHTTP(w, req)
  431. require.Equal(t, http.StatusBadRequest, w.Code)
  432. }
  433. func TestResolveUserMigrationItem_RefreshesBatchStatusToReady(t *testing.T) {
  434. gin.SetMode(gin.TestMode)
  435. router := gin.New()
  436. router.Use(testRootAuthMiddleware())
  437. router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem)
  438. db := setupUserMigrationControllerDB(t)
  439. orig := model.DB
  440. model.DB = db
  441. defer func() {
  442. model.DB = orig
  443. }()
  444. batch := &model.UserMigrationBatch{
  445. Name: "resolve-ready",
  446. SourceRegion: "ov",
  447. TargetRegion: "cn",
  448. OperatorId: 1,
  449. Status: model.UserMigrationBatchStatusScanned,
  450. }
  451. require.NoError(t, model.CreateUserMigrationBatch(batch))
  452. require.NoError(t, db.Create(&model.UserMigrationItem{
  453. BatchId: batch.Id,
  454. SourceUserId: 1001,
  455. SourceUsername: "ov-u1",
  456. Status: model.UserMigrationItemStatusConflict,
  457. }).Error)
  458. req := httptest.NewRequest(
  459. http.MethodPost,
  460. "/api/user-migrations/items/1/resolve",
  461. bytes.NewBufferString(`{"resolution_strategy":"create_new"}`),
  462. )
  463. req.Header.Set("Content-Type", "application/json")
  464. w := httptest.NewRecorder()
  465. router.ServeHTTP(w, req)
  466. require.Equal(t, http.StatusOK, w.Code)
  467. var refreshed model.UserMigrationBatch
  468. require.NoError(t, db.First(&refreshed, batch.Id).Error)
  469. require.Equal(t, model.UserMigrationBatchStatusReady, refreshed.Status)
  470. }
  471. func TestRetryUserMigrationBatch_OnlyFailedItemsReset(t *testing.T) {
  472. gin.SetMode(gin.TestMode)
  473. router := gin.New()
  474. router.Use(testRootAuthMiddleware())
  475. router.POST("/api/user-migrations/batches/:id/retry", RetryUserMigrationBatch)
  476. db := setupUserMigrationControllerDB(t)
  477. orig := model.DB
  478. model.DB = db
  479. defer func() {
  480. model.DB = orig
  481. }()
  482. require.NoError(t, db.Create(&model.UserMigrationBatch{
  483. Id: 1,
  484. Name: "retry-batch",
  485. SourceRegion: "ov",
  486. TargetRegion: "cn",
  487. OperatorId: 1,
  488. Status: model.UserMigrationBatchStatusRunning,
  489. }).Error)
  490. require.NoError(t, db.Create(&model.UserMigrationItem{
  491. Id: 1,
  492. BatchId: 1,
  493. SourceUserId: 1001,
  494. SourceUsername: "failed-user",
  495. Status: model.UserMigrationItemStatusFailed,
  496. ErrorMessage: "boom",
  497. }).Error)
  498. require.NoError(t, db.Create(&model.UserMigrationItem{
  499. Id: 2,
  500. BatchId: 1,
  501. SourceUserId: 1002,
  502. SourceUsername: "migrated-user",
  503. Status: model.UserMigrationItemStatusMigrated,
  504. }).Error)
  505. require.NoError(t, db.Create(&model.UserMigrationItem{
  506. Id: 3,
  507. BatchId: 1,
  508. SourceUserId: 1003,
  509. SourceUsername: "skipped-user",
  510. Status: model.UserMigrationItemStatusSkipped,
  511. }).Error)
  512. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/retry", nil)
  513. w := httptest.NewRecorder()
  514. router.ServeHTTP(w, req)
  515. require.Equal(t, http.StatusConflict, w.Code)
  516. var failedItem, migratedItem, skippedItem model.UserMigrationItem
  517. require.NoError(t, db.First(&failedItem, 1).Error)
  518. require.NoError(t, db.First(&migratedItem, 2).Error)
  519. require.NoError(t, db.First(&skippedItem, 3).Error)
  520. require.Equal(t, model.UserMigrationItemStatusReady, failedItem.Status)
  521. require.Empty(t, failedItem.ErrorMessage)
  522. require.Equal(t, model.UserMigrationItemStatusMigrated, migratedItem.Status)
  523. require.Equal(t, model.UserMigrationItemStatusSkipped, skippedItem.Status)
  524. }
  525. func TestCancelUserMigrationBatch_SoftDeleteSuccess(t *testing.T) {
  526. gin.SetMode(gin.TestMode)
  527. router := gin.New()
  528. router.Use(testRootAuthMiddleware())
  529. router.POST("/api/user-migrations/batches/:id/cancel", CancelUserMigrationBatch)
  530. db := setupUserMigrationControllerDB(t)
  531. orig := model.DB
  532. model.DB = db
  533. defer func() {
  534. model.DB = orig
  535. }()
  536. require.NoError(t, db.Create(&model.UserMigrationBatch{
  537. Id: 1,
  538. Name: "cancel-me",
  539. SourceRegion: "ov",
  540. TargetRegion: "cn",
  541. SelectionMode: "explicit_ids",
  542. SelectionPayload: `{"source_user_ids":[11,12]}`,
  543. OperatorId: 1,
  544. Status: model.UserMigrationBatchStatusScanned,
  545. }).Error)
  546. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  547. w := httptest.NewRecorder()
  548. router.ServeHTTP(w, req)
  549. require.Equal(t, http.StatusOK, w.Code)
  550. var batch model.UserMigrationBatch
  551. require.NoError(t, db.First(&batch, 1).Error)
  552. require.Equal(t, model.UserMigrationBatchStatusCancelled, batch.Status)
  553. }
  554. func TestCancelUserMigrationBatch_InvalidStateRejected(t *testing.T) {
  555. gin.SetMode(gin.TestMode)
  556. router := gin.New()
  557. router.Use(testRootAuthMiddleware())
  558. router.POST("/api/user-migrations/batches/:id/cancel", CancelUserMigrationBatch)
  559. db := setupUserMigrationControllerDB(t)
  560. orig := model.DB
  561. model.DB = db
  562. defer func() {
  563. model.DB = orig
  564. }()
  565. for _, status := range []string{
  566. model.UserMigrationBatchStatusScanning,
  567. model.UserMigrationBatchStatusRunning,
  568. model.UserMigrationBatchStatusCompleted,
  569. } {
  570. require.NoError(t, db.Exec("DELETE FROM user_migration_batches").Error)
  571. require.NoError(t, db.Create(&model.UserMigrationBatch{
  572. Id: 1,
  573. Name: "cannot-cancel",
  574. SourceRegion: "ov",
  575. TargetRegion: "cn",
  576. OperatorId: 1,
  577. Status: status,
  578. }).Error)
  579. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  580. w := httptest.NewRecorder()
  581. router.ServeHTTP(w, req)
  582. require.Equal(t, http.StatusConflict, w.Code)
  583. }
  584. }
  585. func TestCreateUserMigrationBatch_ExplicitIDs_AllowsReuseAfterCancel(t *testing.T) {
  586. gin.SetMode(gin.TestMode)
  587. router := gin.New()
  588. router.Use(testRootAuthMiddleware())
  589. router.POST("/api/user-migrations/batches", CreateUserMigrationBatch)
  590. router.POST("/api/user-migrations/batches/:id/cancel", CancelUserMigrationBatch)
  591. db := setupUserMigrationControllerDB(t)
  592. orig := model.DB
  593. model.DB = db
  594. defer func() {
  595. model.DB = orig
  596. }()
  597. existingBatch := &model.UserMigrationBatch{
  598. Name: "existing-scanned-batch",
  599. SourceRegion: "ov",
  600. TargetRegion: "cn",
  601. SelectionMode: "explicit_ids",
  602. SelectionPayload: `{"source_user_ids":[11,12]}`,
  603. Status: model.UserMigrationBatchStatusScanned,
  604. OperatorId: 1,
  605. }
  606. require.NoError(t, db.Create(existingBatch).Error)
  607. cancelReq := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  608. cancelW := httptest.NewRecorder()
  609. router.ServeHTTP(cancelW, cancelReq)
  610. require.Equal(t, http.StatusOK, cancelW.Code)
  611. createReq := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBufferString(`{
  612. "name":"recreated-batch",
  613. "source_region":"ov",
  614. "target_region":"cn",
  615. "selection_mode":"explicit_ids",
  616. "source_user_ids":[11,12]
  617. }`))
  618. createReq.Header.Set("Content-Type", "application/json")
  619. createW := httptest.NewRecorder()
  620. router.ServeHTTP(createW, createReq)
  621. require.Equal(t, http.StatusOK, createW.Code)
  622. }
  623. func TestVerifyUserMigrationBatch_AggregatesResult(t *testing.T) {
  624. gin.SetMode(gin.TestMode)
  625. router := gin.New()
  626. router.Use(testRootAuthMiddleware())
  627. router.POST("/api/user-migrations/batches/:id/verify", VerifyUserMigrationBatch)
  628. db := setupUserMigrationControllerDB(t)
  629. orig := model.DB
  630. origFactory := newMigrationVerifier
  631. model.DB = db
  632. defer func() {
  633. model.DB = orig
  634. newMigrationVerifier = origFactory
  635. }()
  636. require.NoError(t, db.Create(&model.UserMigrationItem{
  637. Id: 1,
  638. BatchId: 1,
  639. SourceUserId: 1001,
  640. TargetUserId: 1,
  641. Status: model.UserMigrationItemStatusMigrated,
  642. }).Error)
  643. require.NoError(t, db.Create(&model.UserMigrationItem{
  644. Id: 2,
  645. BatchId: 1,
  646. SourceUserId: 1002,
  647. TargetUserId: 2,
  648. Status: model.UserMigrationItemStatusMigrated,
  649. }).Error)
  650. require.NoError(t, db.Create(&model.UserMigrationItem{
  651. Id: 3,
  652. BatchId: 1,
  653. SourceUserId: 1003,
  654. TargetUserId: 3,
  655. Status: model.UserMigrationItemStatusMigrated,
  656. }).Error)
  657. newMigrationVerifier = func() migrationVerifier {
  658. return &fakeMigrationVerifier{
  659. results: map[int]*user_migration.VerifyResult{
  660. 1001: {
  661. TargetUserExists: true,
  662. RemoteConverted: true,
  663. QuotaMatched: true,
  664. },
  665. 1002: {
  666. TargetUserExists: false,
  667. RemoteConverted: true,
  668. QuotaMatched: false,
  669. },
  670. },
  671. errs: map[int]error{
  672. 1003: errors.New("remote user not found"),
  673. },
  674. }
  675. }
  676. req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/verify", nil)
  677. w := httptest.NewRecorder()
  678. router.ServeHTTP(w, req)
  679. require.Equal(t, http.StatusOK, w.Code)
  680. var resp struct {
  681. Success bool `json:"success"`
  682. Data []map[string]any `json:"data"`
  683. }
  684. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  685. require.True(t, resp.Success)
  686. require.Len(t, resp.Data, 3)
  687. bySource := make(map[int]map[string]any, len(resp.Data))
  688. for _, entry := range resp.Data {
  689. bySource[int(entry["source_user_id"].(float64))] = entry
  690. }
  691. require.NotNil(t, bySource[1001]["result"])
  692. require.NotNil(t, bySource[1002]["result"])
  693. require.Equal(t, "remote user not found", bySource[1003]["error"])
  694. }