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ů.
 
 
 

476 řádky
13 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 TestExecuteBatch_CreateNew_CreatesUserAndConvertsRemote(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{
  225. Name: "wave-exec",
  226. Status: model.UserMigrationBatchStatusReady,
  227. }
  228. batch.SourceRegion = "ov"
  229. batch.TargetRegion = "cn"
  230. batch.OperatorId = 1
  231. require.NoError(t, model.CreateUserMigrationBatch(batch))
  232. item := &model.UserMigrationItem{
  233. BatchId: batch.Id,
  234. SourceUserId: 10000001,
  235. SourceUsername: "ov-alice",
  236. SourceEmail: "alice@example.com",
  237. SourceQuota: 300,
  238. ResolutionStrategy: model.UserMigrationStrategyCreateNew,
  239. Status: model.UserMigrationItemStatusReady,
  240. MigrationToken: "batch-1-user-10000001",
  241. }
  242. require.NoError(t, model.UpsertUserMigrationItem(item))
  243. client := &fakeClient{
  244. getUserResp: &RemoteUserSnapshot{
  245. Id: 10000001,
  246. Username: "ov-alice",
  247. Password: "hashed-password",
  248. Email: "alice@example.com",
  249. DisplayName: "Alice",
  250. Status: 1,
  251. Role: 1,
  252. Group: "default",
  253. Quota: 300,
  254. AffCode: "OVA1",
  255. },
  256. }
  257. svc := &Service{client: client}
  258. locked, err := model.TryLockBatchForExecution(batch.Id)
  259. require.NoError(t, err)
  260. require.True(t, locked)
  261. require.NoError(t, svc.ExecuteBatch(batch.Id))
  262. var cnUser model.User
  263. require.NoError(t, db.Where("username = ?", "ov-alice").First(&cnUser).Error)
  264. require.Equal(t, 300, cnUser.Quota)
  265. require.Less(t, cnUser.Id, 10000000)
  266. require.Len(t, client.convertCalls, 1)
  267. require.Equal(t, 10000001, client.convertCalls[0])
  268. }
  269. func TestExecuteBatch_CreateNew_IdempotentOnRetry(t *testing.T) {
  270. db := setupMigrationServiceDB(t)
  271. orig := model.DB
  272. model.DB = db
  273. defer func() {
  274. model.DB = orig
  275. }()
  276. batch := &model.UserMigrationBatch{
  277. Name: "wave-retry",
  278. Status: model.UserMigrationBatchStatusReady,
  279. }
  280. batch.SourceRegion = "ov"
  281. batch.TargetRegion = "cn"
  282. batch.OperatorId = 1
  283. require.NoError(t, model.CreateUserMigrationBatch(batch))
  284. item := &model.UserMigrationItem{
  285. BatchId: batch.Id,
  286. SourceUserId: 10000002,
  287. SourceUsername: "ov-carol",
  288. SourceEmail: "carol@example.com",
  289. SourceQuota: 150,
  290. ResolutionStrategy: model.UserMigrationStrategyCreateNew,
  291. Status: model.UserMigrationItemStatusReady,
  292. MigrationToken: "batch-1-user-10000002",
  293. }
  294. require.NoError(t, model.UpsertUserMigrationItem(item))
  295. snapshot := &RemoteUserSnapshot{
  296. Id: 10000002,
  297. Username: "ov-carol",
  298. Password: "hash",
  299. Email: "carol@example.com",
  300. Quota: 150,
  301. AffCode: "OVC1",
  302. }
  303. client := &fakeClient{getUserResp: snapshot}
  304. svc := &Service{client: client}
  305. locked, err := model.TryLockBatchForExecution(batch.Id)
  306. require.NoError(t, err)
  307. require.True(t, locked)
  308. require.NoError(t, svc.ExecuteBatch(batch.Id))
  309. require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
  310. Update("status", model.UserMigrationBatchStatusFailed).Error)
  311. require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
  312. Update("status", model.UserMigrationItemStatusReady).Error)
  313. locked, err = model.TryLockBatchForExecution(batch.Id)
  314. require.NoError(t, err)
  315. require.True(t, locked)
  316. require.NoError(t, svc.ExecuteBatch(batch.Id))
  317. var count int64
  318. require.NoError(t, db.Model(&model.User{}).Where("username = ?", "ov-carol").Count(&count).Error)
  319. require.EqualValues(t, 1, count)
  320. require.Len(t, client.convertCalls, 1)
  321. }
  322. func TestExecuteBatch_MergeExisting_IdempotentQuota(t *testing.T) {
  323. db := setupMigrationServiceDB(t)
  324. orig := model.DB
  325. model.DB = db
  326. defer func() {
  327. model.DB = orig
  328. }()
  329. require.NoError(t, db.Create(&model.User{
  330. Id: 1,
  331. Username: "cn-bob",
  332. Password: "hash",
  333. Email: "bob@example.com",
  334. AffCode: "CNB1",
  335. Quota: 500,
  336. }).Error)
  337. batch := &model.UserMigrationBatch{
  338. Name: "wave-merge",
  339. Status: model.UserMigrationBatchStatusReady,
  340. }
  341. batch.SourceRegion = "ov"
  342. batch.TargetRegion = "cn"
  343. batch.OperatorId = 1
  344. require.NoError(t, model.CreateUserMigrationBatch(batch))
  345. item := &model.UserMigrationItem{
  346. BatchId: batch.Id,
  347. SourceUserId: 10000002,
  348. TargetUserId: 1,
  349. SourceUsername: "ov-bob",
  350. SourceEmail: "bob@example.com",
  351. SourceQuota: 200,
  352. ResolutionStrategy: model.UserMigrationStrategyMergeExisting,
  353. Status: model.UserMigrationItemStatusReady,
  354. MigrationToken: "batch-1-user-10000002",
  355. }
  356. require.NoError(t, model.UpsertUserMigrationItem(item))
  357. client := &fakeClient{
  358. getUserResp: &RemoteUserSnapshot{
  359. Id: 10000002,
  360. Username: "ov-bob",
  361. Password: "hashed",
  362. Email: "bob@example.com",
  363. Quota: 200,
  364. },
  365. }
  366. svc := &Service{client: client}
  367. for i := 0; i < 2; i++ {
  368. require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
  369. Update("status", model.UserMigrationBatchStatusReady).Error)
  370. require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
  371. Update("status", model.UserMigrationItemStatusReady).Error)
  372. locked, err := model.TryLockBatchForExecution(batch.Id)
  373. require.NoError(t, err)
  374. require.True(t, locked)
  375. require.NoError(t, svc.ExecuteBatch(batch.Id))
  376. }
  377. var saved model.User
  378. require.NoError(t, db.First(&saved, 1).Error)
  379. require.Equal(t, 700, saved.Quota)
  380. require.Len(t, client.convertCalls, 1)
  381. }
  382. func TestVerifyItem_CorrectlyChecksRemoteConverted(t *testing.T) {
  383. db := setupMigrationServiceDB(t)
  384. orig := model.DB
  385. model.DB = db
  386. defer func() {
  387. model.DB = orig
  388. }()
  389. require.NoError(t, db.Create(&model.User{
  390. Id: 789,
  391. Username: "cn-migrated",
  392. Password: "hash",
  393. AffCode: "CM1",
  394. Quota: 300,
  395. }).Error)
  396. item := &model.UserMigrationItem{
  397. BatchId: 1,
  398. SourceUserId: 10000001,
  399. TargetUserId: 789,
  400. SourceQuota: 300,
  401. Status: model.UserMigrationItemStatusMigrated,
  402. }
  403. client := &fakeClient{
  404. getUserResp: &RemoteUserSnapshot{
  405. Id: 10000001,
  406. Source: common.UserSourceSynced,
  407. RemoteUserId: 789,
  408. },
  409. }
  410. svc := &Service{client: client}
  411. result, err := svc.VerifyItem(item)
  412. require.NoError(t, err)
  413. require.True(t, result.TargetUserExists)
  414. require.True(t, result.RemoteConverted)
  415. require.True(t, result.QuotaMatched)
  416. }