Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

554 linhas
15 KiB

  1. package user_migration
  2. import (
  3. "fmt"
  4. "testing"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/model"
  7. "github.com/glebarez/sqlite"
  8. "github.com/stretchr/testify/require"
  9. "gorm.io/gorm"
  10. )
  11. type fakeClient struct {
  12. listPages map[int][]RemoteUserSnapshot
  13. getUserResp *RemoteUserSnapshot
  14. getUserByID map[int]*RemoteUserSnapshot
  15. hasCopyResult bool
  16. hasCopyByID map[int]bool
  17. convertCalls []int
  18. convertReqs map[int]*ConvertRemoteUserRequest
  19. }
  20. func (f *fakeClient) ListUsers(page, pageSize int) (*ListRemoteUsersResponse, error) {
  21. users := f.listPages[page]
  22. return &ListRemoteUsersResponse{
  23. Success: true,
  24. Data: users,
  25. Total: len(users),
  26. }, nil
  27. }
  28. func (f *fakeClient) GetUser(userId int) (*RemoteUserSnapshot, error) {
  29. if snapshot, ok := f.getUserByID[userId]; ok {
  30. return snapshot, nil
  31. }
  32. if f.getUserResp != nil {
  33. return f.getUserResp, nil
  34. }
  35. return nil, fmt.Errorf("user %d not found", userId)
  36. }
  37. func (f *fakeClient) ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error {
  38. f.convertCalls = append(f.convertCalls, userId)
  39. if f.convertReqs == nil {
  40. f.convertReqs = make(map[int]*ConvertRemoteUserRequest)
  41. }
  42. if req != nil {
  43. copyReq := *req
  44. f.convertReqs[userId] = &copyReq
  45. if snapshot, ok := f.getUserByID[userId]; ok && snapshot != nil {
  46. snapshot.Source = common.UserSourceSynced
  47. snapshot.RemoteUserId = req.RemoteUserId
  48. }
  49. if f.getUserResp != nil && f.getUserResp.Id == userId {
  50. f.getUserResp.Source = common.UserSourceSynced
  51. f.getUserResp.RemoteUserId = req.RemoteUserId
  52. }
  53. }
  54. return nil
  55. }
  56. func (f *fakeClient) HasSyncedCopy(cnUserId int) (bool, error) {
  57. if value, ok := f.hasCopyByID[cnUserId]; ok {
  58. return value, nil
  59. }
  60. return f.hasCopyResult, nil
  61. }
  62. func setupMigrationServiceDB(t *testing.T) *gorm.DB {
  63. t.Helper()
  64. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  65. require.NoError(t, err)
  66. sqlDB, err := db.DB()
  67. require.NoError(t, err)
  68. sqlDB.SetMaxOpenConns(1)
  69. require.NoError(t, db.AutoMigrate(
  70. &model.User{},
  71. &model.CustomOAuthProvider{},
  72. &model.UserOAuthBinding{},
  73. &model.UserMigrationBatch{},
  74. &model.UserMigrationItem{},
  75. &model.MigrationQuotaGrant{},
  76. ))
  77. return db
  78. }
  79. func TestScanBatch_NoConflict_MarksReady(t *testing.T) {
  80. db := setupMigrationServiceDB(t)
  81. orig := model.DB
  82. model.DB = db
  83. defer func() {
  84. model.DB = orig
  85. }()
  86. batch := &model.UserMigrationBatch{Name: "wave-1", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  87. require.NoError(t, model.CreateUserMigrationBatch(batch))
  88. client := &fakeClient{
  89. listPages: map[int][]RemoteUserSnapshot{
  90. 0: {{
  91. Id: 10000001,
  92. Username: "ov-bob",
  93. Email: "bob@example.com",
  94. Quota: 456,
  95. AffCode: "A002",
  96. }},
  97. },
  98. }
  99. svc := &Service{client: client}
  100. require.NoError(t, svc.ScanBatch(batch.Id))
  101. items, err := model.ListUserMigrationItemsByBatch(batch.Id)
  102. require.NoError(t, err)
  103. require.Len(t, items, 1)
  104. require.Equal(t, model.UserMigrationStrategyCreateNew, items[0].ResolutionStrategy)
  105. require.Equal(t, model.UserMigrationItemStatusReady, items[0].Status)
  106. var savedBatch model.UserMigrationBatch
  107. require.NoError(t, db.First(&savedBatch, batch.Id).Error)
  108. require.Equal(t, model.UserMigrationBatchStatusReady, savedBatch.Status)
  109. }
  110. func TestScanBatch_EmailConflict_NoCNSyncedCopy_RecommendsMerge(t *testing.T) {
  111. db := setupMigrationServiceDB(t)
  112. orig := model.DB
  113. model.DB = db
  114. defer func() {
  115. model.DB = orig
  116. }()
  117. require.NoError(t, db.Create(&model.User{
  118. Id: 1,
  119. Username: "cn-alice",
  120. Password: "hash",
  121. Email: "alice@example.com",
  122. AffCode: "C001",
  123. }).Error)
  124. batch := &model.UserMigrationBatch{Name: "wave-2", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  125. require.NoError(t, model.CreateUserMigrationBatch(batch))
  126. client := &fakeClient{
  127. listPages: map[int][]RemoteUserSnapshot{
  128. 0: {{
  129. Id: 10000001,
  130. Username: "ov-alice",
  131. Email: "alice@example.com",
  132. Quota: 123,
  133. AffCode: "A001",
  134. }},
  135. },
  136. hasCopyResult: false,
  137. }
  138. svc := &Service{client: client}
  139. require.NoError(t, svc.ScanBatch(batch.Id))
  140. items, err := model.ListUserMigrationItemsByBatch(batch.Id)
  141. require.NoError(t, err)
  142. require.Len(t, items, 1)
  143. require.Equal(t, model.UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy)
  144. require.Equal(t, model.UserMigrationItemStatusReady, items[0].Status)
  145. require.Equal(t, 1, items[0].TargetUserId)
  146. }
  147. func TestScanBatch_EmailConflict_CNHasSyncedCopy_MarksHardConflict(t *testing.T) {
  148. db := setupMigrationServiceDB(t)
  149. orig := model.DB
  150. model.DB = db
  151. defer func() {
  152. model.DB = orig
  153. }()
  154. require.NoError(t, db.Create(&model.User{
  155. Id: 456,
  156. Username: "cn-user456",
  157. Password: "hash",
  158. Email: "shared@example.com",
  159. AffCode: "C456",
  160. }).Error)
  161. batch := &model.UserMigrationBatch{Name: "wave-3", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  162. require.NoError(t, model.CreateUserMigrationBatch(batch))
  163. client := &fakeClient{
  164. listPages: map[int][]RemoteUserSnapshot{
  165. 0: {{
  166. Id: 10000003,
  167. Username: "ov-x",
  168. Email: "shared@example.com",
  169. Quota: 50,
  170. AffCode: "OVX",
  171. }},
  172. },
  173. hasCopyResult: true,
  174. }
  175. svc := &Service{client: client}
  176. require.NoError(t, svc.ScanBatch(batch.Id))
  177. items, err := model.ListUserMigrationItemsByBatch(batch.Id)
  178. require.NoError(t, err)
  179. require.Len(t, items, 1)
  180. require.Equal(t, model.UserMigrationItemStatusConflict, items[0].Status)
  181. require.Contains(t, items[0].ConflictFlags, "cn_already_synced_to_ov")
  182. }
  183. func TestScanBatch_Paginated(t *testing.T) {
  184. db := setupMigrationServiceDB(t)
  185. orig := model.DB
  186. model.DB = db
  187. defer func() {
  188. model.DB = orig
  189. }()
  190. batch := &model.UserMigrationBatch{Name: "wave-big", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  191. require.NoError(t, model.CreateUserMigrationBatch(batch))
  192. page0 := make([]RemoteUserSnapshot, scanPageSize)
  193. for i := range page0 {
  194. page0[i] = RemoteUserSnapshot{
  195. Id: 10000001 + i,
  196. Username: fmt.Sprintf("u%d", i),
  197. AffCode: fmt.Sprintf("A%d", i),
  198. }
  199. }
  200. page1 := []RemoteUserSnapshot{{
  201. Id: 10000001 + scanPageSize,
  202. Username: "last",
  203. AffCode: "LAST",
  204. }}
  205. client := &fakeClient{
  206. listPages: map[int][]RemoteUserSnapshot{
  207. 0: page0,
  208. 1: page1,
  209. },
  210. }
  211. svc := &Service{client: client}
  212. require.NoError(t, svc.ScanBatch(batch.Id))
  213. items, err := model.ListUserMigrationItemsByBatch(batch.Id)
  214. require.NoError(t, err)
  215. require.Len(t, items, scanPageSize+1)
  216. }
  217. func TestScanBatch_Rescan_PreservesHumanStrategy(t *testing.T) {
  218. db := setupMigrationServiceDB(t)
  219. orig := model.DB
  220. model.DB = db
  221. defer func() {
  222. model.DB = orig
  223. }()
  224. batch := &model.UserMigrationBatch{Name: "wave-rescan", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  225. require.NoError(t, model.CreateUserMigrationBatch(batch))
  226. client := &fakeClient{
  227. listPages: map[int][]RemoteUserSnapshot{
  228. 0: {{
  229. Id: 10000001,
  230. Username: "ov-dave",
  231. Email: "dave@example.com",
  232. Quota: 100,
  233. AffCode: "OVD1",
  234. }},
  235. },
  236. hasCopyResult: false,
  237. }
  238. svc := &Service{client: client}
  239. require.NoError(t, svc.ScanBatch(batch.Id))
  240. require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
  241. Update("resolution_strategy", model.UserMigrationStrategySkip).Error)
  242. require.NoError(t, svc.ScanBatch(batch.Id))
  243. items, err := model.ListUserMigrationItemsByBatch(batch.Id)
  244. require.NoError(t, err)
  245. require.Len(t, items, 1)
  246. require.Equal(t, model.UserMigrationStrategySkip, items[0].ResolutionStrategy)
  247. }
  248. func TestExecuteBatch_CreateNew_CreatesUserAndConvertsRemote(t *testing.T) {
  249. db := setupMigrationServiceDB(t)
  250. orig := model.DB
  251. model.DB = db
  252. defer func() {
  253. model.DB = orig
  254. }()
  255. batch := &model.UserMigrationBatch{
  256. Name: "wave-exec",
  257. Status: model.UserMigrationBatchStatusReady,
  258. }
  259. batch.SourceRegion = "ov"
  260. batch.TargetRegion = "cn"
  261. batch.OperatorId = 1
  262. require.NoError(t, model.CreateUserMigrationBatch(batch))
  263. item := &model.UserMigrationItem{
  264. BatchId: batch.Id,
  265. SourceUserId: 10000001,
  266. SourceUsername: "ov-alice",
  267. SourceEmail: "alice@example.com",
  268. SourceQuota: 300,
  269. ResolutionStrategy: model.UserMigrationStrategyCreateNew,
  270. Status: model.UserMigrationItemStatusReady,
  271. MigrationToken: "batch-1-user-10000001",
  272. }
  273. require.NoError(t, model.UpsertUserMigrationItem(item))
  274. client := &fakeClient{
  275. getUserResp: &RemoteUserSnapshot{
  276. Id: 10000001,
  277. Username: "ov-alice",
  278. Password: "hashed-password",
  279. Email: "alice@example.com",
  280. DisplayName: "Alice",
  281. Status: 1,
  282. Role: 1,
  283. Group: "default",
  284. Quota: 300,
  285. AffCode: "OVA1",
  286. },
  287. }
  288. svc := &Service{client: client}
  289. locked, err := model.TryLockBatchForExecution(batch.Id)
  290. require.NoError(t, err)
  291. require.True(t, locked)
  292. require.NoError(t, svc.ExecuteBatch(batch.Id))
  293. var cnUser model.User
  294. require.NoError(t, db.Where("username = ?", "ov-alice").First(&cnUser).Error)
  295. require.Equal(t, 300, cnUser.Quota)
  296. require.Less(t, cnUser.Id, 10000000)
  297. require.Len(t, client.convertCalls, 1)
  298. require.Equal(t, 10000001, client.convertCalls[0])
  299. }
  300. func TestExecuteBatch_CreateNew_IdempotentOnRetry(t *testing.T) {
  301. db := setupMigrationServiceDB(t)
  302. orig := model.DB
  303. model.DB = db
  304. defer func() {
  305. model.DB = orig
  306. }()
  307. batch := &model.UserMigrationBatch{
  308. Name: "wave-retry",
  309. Status: model.UserMigrationBatchStatusReady,
  310. }
  311. batch.SourceRegion = "ov"
  312. batch.TargetRegion = "cn"
  313. batch.OperatorId = 1
  314. require.NoError(t, model.CreateUserMigrationBatch(batch))
  315. item := &model.UserMigrationItem{
  316. BatchId: batch.Id,
  317. SourceUserId: 10000002,
  318. SourceUsername: "ov-carol",
  319. SourceEmail: "carol@example.com",
  320. SourceQuota: 150,
  321. ResolutionStrategy: model.UserMigrationStrategyCreateNew,
  322. Status: model.UserMigrationItemStatusReady,
  323. MigrationToken: "batch-1-user-10000002",
  324. }
  325. require.NoError(t, model.UpsertUserMigrationItem(item))
  326. snapshot := &RemoteUserSnapshot{
  327. Id: 10000002,
  328. Username: "ov-carol",
  329. Password: "hash",
  330. Email: "carol@example.com",
  331. Quota: 150,
  332. AffCode: "OVC1",
  333. }
  334. client := &fakeClient{getUserResp: snapshot}
  335. svc := &Service{client: client}
  336. locked, err := model.TryLockBatchForExecution(batch.Id)
  337. require.NoError(t, err)
  338. require.True(t, locked)
  339. require.NoError(t, svc.ExecuteBatch(batch.Id))
  340. require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
  341. Update("status", model.UserMigrationBatchStatusFailed).Error)
  342. require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
  343. Update("status", model.UserMigrationItemStatusReady).Error)
  344. locked, err = model.TryLockBatchForExecution(batch.Id)
  345. require.NoError(t, err)
  346. require.True(t, locked)
  347. require.NoError(t, svc.ExecuteBatch(batch.Id))
  348. var count int64
  349. require.NoError(t, db.Model(&model.User{}).Where("username = ?", "ov-carol").Count(&count).Error)
  350. require.EqualValues(t, 1, count)
  351. require.Len(t, client.convertCalls, 1)
  352. }
  353. func TestExecuteBatch_MergeExisting_IdempotentQuota(t *testing.T) {
  354. db := setupMigrationServiceDB(t)
  355. orig := model.DB
  356. model.DB = db
  357. defer func() {
  358. model.DB = orig
  359. }()
  360. require.NoError(t, db.Create(&model.User{
  361. Id: 1,
  362. Username: "cn-bob",
  363. Password: "hash",
  364. Email: "bob@example.com",
  365. AffCode: "CNB1",
  366. Quota: 500,
  367. }).Error)
  368. batch := &model.UserMigrationBatch{
  369. Name: "wave-merge",
  370. Status: model.UserMigrationBatchStatusReady,
  371. }
  372. batch.SourceRegion = "ov"
  373. batch.TargetRegion = "cn"
  374. batch.OperatorId = 1
  375. require.NoError(t, model.CreateUserMigrationBatch(batch))
  376. item := &model.UserMigrationItem{
  377. BatchId: batch.Id,
  378. SourceUserId: 10000002,
  379. TargetUserId: 1,
  380. SourceUsername: "ov-bob",
  381. SourceEmail: "bob@example.com",
  382. SourceQuota: 200,
  383. ResolutionStrategy: model.UserMigrationStrategyMergeExisting,
  384. Status: model.UserMigrationItemStatusReady,
  385. MigrationToken: "batch-1-user-10000002",
  386. }
  387. require.NoError(t, model.UpsertUserMigrationItem(item))
  388. client := &fakeClient{
  389. getUserResp: &RemoteUserSnapshot{
  390. Id: 10000002,
  391. Username: "ov-bob",
  392. Password: "hashed",
  393. Email: "bob@example.com",
  394. Quota: 200,
  395. },
  396. }
  397. svc := &Service{client: client}
  398. for i := 0; i < 2; i++ {
  399. require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
  400. Update("status", model.UserMigrationBatchStatusReady).Error)
  401. require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
  402. Update("status", model.UserMigrationItemStatusReady).Error)
  403. locked, err := model.TryLockBatchForExecution(batch.Id)
  404. require.NoError(t, err)
  405. require.True(t, locked)
  406. require.NoError(t, svc.ExecuteBatch(batch.Id))
  407. }
  408. var saved model.User
  409. require.NoError(t, db.First(&saved, 1).Error)
  410. require.Equal(t, 700, saved.Quota)
  411. require.Len(t, client.convertCalls, 1)
  412. }
  413. func TestExecuteBatch_SkipItem_DoesNotConvert(t *testing.T) {
  414. db := setupMigrationServiceDB(t)
  415. orig := model.DB
  416. model.DB = db
  417. defer func() {
  418. model.DB = orig
  419. }()
  420. batch := &model.UserMigrationBatch{
  421. Name: "skip-test",
  422. SourceRegion: "ov",
  423. TargetRegion: "cn",
  424. OperatorId: 1,
  425. Status: model.UserMigrationBatchStatusReady,
  426. }
  427. require.NoError(t, model.CreateUserMigrationBatch(batch))
  428. require.NoError(t, model.UpsertUserMigrationItem(&model.UserMigrationItem{
  429. BatchId: batch.Id,
  430. SourceUserId: 10000099,
  431. SourceUsername: "skip-user",
  432. SourceQuota: 50,
  433. ResolutionStrategy: model.UserMigrationStrategySkip,
  434. Status: model.UserMigrationItemStatusReady,
  435. MigrationToken: "batch-1-user-10000099",
  436. }))
  437. client := &fakeClient{}
  438. svc := &Service{client: client}
  439. locked, err := model.TryLockBatchForExecution(batch.Id)
  440. require.NoError(t, err)
  441. require.True(t, locked)
  442. require.NoError(t, svc.ExecuteBatch(batch.Id))
  443. require.Len(t, client.convertCalls, 0)
  444. items, err := model.ListUserMigrationItemsByBatch(batch.Id)
  445. require.NoError(t, err)
  446. require.Len(t, items, 1)
  447. require.Equal(t, model.UserMigrationItemStatusSkipped, items[0].Status)
  448. }
  449. func TestVerifyItem_CorrectlyChecksRemoteConverted(t *testing.T) {
  450. db := setupMigrationServiceDB(t)
  451. orig := model.DB
  452. model.DB = db
  453. defer func() {
  454. model.DB = orig
  455. }()
  456. require.NoError(t, db.Create(&model.User{
  457. Id: 789,
  458. Username: "cn-migrated",
  459. Password: "hash",
  460. AffCode: "CM1",
  461. Quota: 300,
  462. }).Error)
  463. item := &model.UserMigrationItem{
  464. BatchId: 1,
  465. SourceUserId: 10000001,
  466. TargetUserId: 789,
  467. SourceQuota: 300,
  468. Status: model.UserMigrationItemStatusMigrated,
  469. }
  470. client := &fakeClient{
  471. getUserResp: &RemoteUserSnapshot{
  472. Id: 10000001,
  473. Source: common.UserSourceSynced,
  474. RemoteUserId: 789,
  475. },
  476. }
  477. svc := &Service{client: client}
  478. result, err := svc.VerifyItem(item)
  479. require.NoError(t, err)
  480. require.True(t, result.TargetUserExists)
  481. require.True(t, result.RemoteConverted)
  482. require.True(t, result.QuotaMatched)
  483. }