Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

286 wiersze
7.5 KiB

  1. package model
  2. import (
  3. "encoding/json"
  4. "testing"
  5. "github.com/glebarez/sqlite"
  6. "github.com/stretchr/testify/require"
  7. "gorm.io/gorm"
  8. )
  9. func setupUserMigrationTestDB(t *testing.T) *gorm.DB {
  10. t.Helper()
  11. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  12. require.NoError(t, err)
  13. sqlDB, err := db.DB()
  14. require.NoError(t, err)
  15. sqlDB.SetMaxOpenConns(1)
  16. require.NoError(t, db.AutoMigrate(
  17. &User{},
  18. &CustomOAuthProvider{},
  19. &UserOAuthBinding{},
  20. &UserMigrationBatch{},
  21. &UserMigrationItem{},
  22. &MigrationQuotaGrant{},
  23. ))
  24. return db
  25. }
  26. func TestUserMigrationModels_AutoMigrate(t *testing.T) {
  27. db := setupUserMigrationTestDB(t)
  28. require.True(t, db.Migrator().HasTable(&UserMigrationBatch{}))
  29. require.True(t, db.Migrator().HasTable(&UserMigrationItem{}))
  30. require.True(t, db.Migrator().HasTable(&MigrationQuotaGrant{}))
  31. }
  32. func TestCreateUserMigrationBatch_DefaultStatus(t *testing.T) {
  33. db := setupUserMigrationTestDB(t)
  34. orig := DB
  35. DB = db
  36. defer func() {
  37. DB = orig
  38. }()
  39. batch := &UserMigrationBatch{
  40. Name: "ov-to-cn-wave-1",
  41. SourceRegion: "overseas",
  42. TargetRegion: "cn",
  43. OperatorId: 1,
  44. }
  45. require.NoError(t, CreateUserMigrationBatch(batch))
  46. require.Equal(t, UserMigrationBatchStatusDraft, batch.Status)
  47. require.Equal(t, "all", batch.SelectionMode)
  48. require.Greater(t, batch.Id, 0)
  49. }
  50. func TestCreateUserMigrationBatch_StoresExplicitIDsSelection(t *testing.T) {
  51. db := setupUserMigrationTestDB(t)
  52. orig := DB
  53. DB = db
  54. defer func() {
  55. DB = orig
  56. }()
  57. payload := map[string]any{"source_user_ids": []int{1001, 1002, 1003}}
  58. payloadJSON, err := json.Marshal(payload)
  59. require.NoError(t, err)
  60. batch := &UserMigrationBatch{
  61. Name: "partial-wave",
  62. SourceRegion: "ov",
  63. TargetRegion: "cn",
  64. SelectionMode: "explicit_ids",
  65. SelectionPayload: string(payloadJSON),
  66. RequestedUserCount: 3,
  67. OperatorId: 1,
  68. }
  69. require.NoError(t, CreateUserMigrationBatch(batch))
  70. var saved UserMigrationBatch
  71. require.NoError(t, DB.First(&saved, batch.Id).Error)
  72. require.Equal(t, "explicit_ids", saved.SelectionMode)
  73. require.Equal(t, string(payloadJSON), saved.SelectionPayload)
  74. require.Equal(t, 3, saved.RequestedUserCount)
  75. require.Equal(t, 0, saved.MatchedUserCount)
  76. require.Equal(t, 0, saved.ExcludedUserCount)
  77. }
  78. func TestUpsertMigrationItem_PreservesHumanSetStrategy(t *testing.T) {
  79. db := setupUserMigrationTestDB(t)
  80. orig := DB
  81. DB = db
  82. defer func() {
  83. DB = orig
  84. }()
  85. batch := &UserMigrationBatch{
  86. Name: "b1",
  87. SourceRegion: "ov",
  88. TargetRegion: "cn",
  89. OperatorId: 1,
  90. }
  91. require.NoError(t, CreateUserMigrationBatch(batch))
  92. first := &UserMigrationItem{
  93. BatchId: batch.Id,
  94. SourceUserId: 10000001,
  95. SourceUsername: "alice",
  96. SourceEmail: "alice@example.com",
  97. SourceQuota: 12345,
  98. ResolutionStrategy: "",
  99. Status: UserMigrationItemStatusConflict,
  100. }
  101. require.NoError(t, UpsertUserMigrationItem(first))
  102. require.NoError(t, DB.Model(&UserMigrationItem{}).
  103. Where("batch_id = ? AND source_user_id = ?", batch.Id, 10000001).
  104. Update("resolution_strategy", UserMigrationStrategyMergeExisting).Error)
  105. second := &UserMigrationItem{
  106. BatchId: batch.Id,
  107. SourceUserId: 10000001,
  108. SourceUsername: "alice",
  109. SourceEmail: "alice@example.com",
  110. SourceQuota: 99999,
  111. ResolutionStrategy: UserMigrationStrategyCreateNew,
  112. Status: UserMigrationItemStatusReady,
  113. }
  114. require.NoError(t, UpsertUserMigrationItem(second))
  115. items, err := ListUserMigrationItemsByBatch(batch.Id)
  116. require.NoError(t, err)
  117. require.Len(t, items, 1)
  118. require.Equal(t, UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy)
  119. require.Equal(t, UserMigrationItemStatusReady, items[0].Status)
  120. require.Equal(t, 99999, items[0].SourceQuota)
  121. }
  122. func TestAtomicGrantQuotaIfNotExists_Idempotent(t *testing.T) {
  123. db := setupUserMigrationTestDB(t)
  124. orig := DB
  125. DB = db
  126. defer func() {
  127. DB = orig
  128. }()
  129. require.NoError(t, db.Create(&User{
  130. Id: 1,
  131. Username: "cn-user",
  132. Password: "hash",
  133. AffCode: "AFF1",
  134. Quota: 500,
  135. }).Error)
  136. var alreadyDone bool
  137. err := db.Transaction(func(tx *gorm.DB) error {
  138. var innerErr error
  139. alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
  140. return innerErr
  141. })
  142. require.NoError(t, err)
  143. require.False(t, alreadyDone)
  144. var user User
  145. require.NoError(t, db.First(&user, 1).Error)
  146. require.Equal(t, 700, user.Quota)
  147. err = db.Transaction(func(tx *gorm.DB) error {
  148. var innerErr error
  149. alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
  150. return innerErr
  151. })
  152. require.NoError(t, err)
  153. require.True(t, alreadyDone)
  154. require.NoError(t, db.First(&user, 1).Error)
  155. require.Equal(t, 700, user.Quota)
  156. }
  157. func TestRefreshBatchStats_ReadyOnlyWhenNoPendingScan(t *testing.T) {
  158. db := setupUserMigrationTestDB(t)
  159. orig := DB
  160. DB = db
  161. defer func() {
  162. DB = orig
  163. }()
  164. batch := &UserMigrationBatch{Name: "stats-test", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  165. require.NoError(t, CreateUserMigrationBatch(batch))
  166. require.NoError(t, DB.Create(&UserMigrationItem{
  167. BatchId: batch.Id,
  168. SourceUserId: 1,
  169. SourceUsername: "u1",
  170. Status: UserMigrationItemStatusReady,
  171. }).Error)
  172. require.NoError(t, DB.Create(&UserMigrationItem{
  173. BatchId: batch.Id,
  174. SourceUserId: 2,
  175. SourceUsername: "u2",
  176. Status: UserMigrationItemStatusPendingScan,
  177. }).Error)
  178. require.NoError(t, RefreshUserMigrationBatchStats(batch.Id))
  179. var refreshed UserMigrationBatch
  180. require.NoError(t, DB.First(&refreshed, batch.Id).Error)
  181. require.Equal(t, UserMigrationBatchStatusScanned, refreshed.Status)
  182. }
  183. func TestRefreshUserMigrationBatchExecutionStats_MigratingLeakHandled(t *testing.T) {
  184. db := setupUserMigrationTestDB(t)
  185. orig := DB
  186. DB = db
  187. defer func() {
  188. DB = orig
  189. }()
  190. batch := &UserMigrationBatch{
  191. Name: "migrating-leak",
  192. SourceRegion: "ov",
  193. TargetRegion: "cn",
  194. OperatorId: 1,
  195. Status: UserMigrationBatchStatusRunning,
  196. }
  197. require.NoError(t, CreateUserMigrationBatch(batch))
  198. require.NoError(t, DB.Create(&UserMigrationItem{
  199. BatchId: batch.Id,
  200. SourceUserId: 101,
  201. SourceUsername: "stuck-user",
  202. Status: UserMigrationItemStatusMigrating,
  203. }).Error)
  204. require.NoError(t, RefreshUserMigrationBatchExecutionStats(batch.Id))
  205. var refreshed UserMigrationBatch
  206. require.NoError(t, DB.First(&refreshed, batch.Id).Error)
  207. require.Equal(t, UserMigrationBatchStatusFailed, refreshed.Status)
  208. require.Equal(t, 0, refreshed.MigratedUsers)
  209. require.Equal(t, 0, refreshed.FailedUsers)
  210. }
  211. func TestCancelUserMigrationBatch_ReleasesBatch(t *testing.T) {
  212. db := setupUserMigrationTestDB(t)
  213. orig := DB
  214. DB = db
  215. defer func() {
  216. DB = orig
  217. }()
  218. batch := &UserMigrationBatch{
  219. Name: "cancel-batch",
  220. SourceRegion: "ov",
  221. TargetRegion: "cn",
  222. SelectionMode: "explicit_ids",
  223. SelectionPayload: `{"source_user_ids":[1001,1002]}`,
  224. OperatorId: 1,
  225. Status: UserMigrationBatchStatusScanned,
  226. }
  227. require.NoError(t, CreateUserMigrationBatch(batch))
  228. conflicts, err := FindActiveMigrationUserConflicts([]int{1001, 1002})
  229. require.NoError(t, err)
  230. require.Len(t, conflicts, 2)
  231. cancelled, err := CancelUserMigrationBatch(batch.Id)
  232. require.NoError(t, err)
  233. require.True(t, cancelled)
  234. var refreshed UserMigrationBatch
  235. require.NoError(t, DB.First(&refreshed, batch.Id).Error)
  236. require.Equal(t, UserMigrationBatchStatusCancelled, refreshed.Status)
  237. conflicts, err = FindActiveMigrationUserConflicts([]int{1001, 1002})
  238. require.NoError(t, err)
  239. require.Len(t, conflicts, 0)
  240. }