Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

215 строки
5.5 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/glebarez/sqlite"
  5. "github.com/stretchr/testify/require"
  6. "gorm.io/gorm"
  7. )
  8. func setupUserMigrationTestDB(t *testing.T) *gorm.DB {
  9. t.Helper()
  10. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  11. require.NoError(t, err)
  12. sqlDB, err := db.DB()
  13. require.NoError(t, err)
  14. sqlDB.SetMaxOpenConns(1)
  15. require.NoError(t, db.AutoMigrate(
  16. &User{},
  17. &CustomOAuthProvider{},
  18. &UserOAuthBinding{},
  19. &UserMigrationBatch{},
  20. &UserMigrationItem{},
  21. &MigrationQuotaGrant{},
  22. ))
  23. return db
  24. }
  25. func TestUserMigrationModels_AutoMigrate(t *testing.T) {
  26. db := setupUserMigrationTestDB(t)
  27. require.True(t, db.Migrator().HasTable(&UserMigrationBatch{}))
  28. require.True(t, db.Migrator().HasTable(&UserMigrationItem{}))
  29. require.True(t, db.Migrator().HasTable(&MigrationQuotaGrant{}))
  30. }
  31. func TestCreateUserMigrationBatch_DefaultStatus(t *testing.T) {
  32. db := setupUserMigrationTestDB(t)
  33. orig := DB
  34. DB = db
  35. defer func() {
  36. DB = orig
  37. }()
  38. batch := &UserMigrationBatch{
  39. Name: "ov-to-cn-wave-1",
  40. SourceRegion: "overseas",
  41. TargetRegion: "cn",
  42. OperatorId: 1,
  43. }
  44. require.NoError(t, CreateUserMigrationBatch(batch))
  45. require.Equal(t, UserMigrationBatchStatusDraft, batch.Status)
  46. require.Greater(t, batch.Id, 0)
  47. }
  48. func TestUpsertMigrationItem_PreservesHumanSetStrategy(t *testing.T) {
  49. db := setupUserMigrationTestDB(t)
  50. orig := DB
  51. DB = db
  52. defer func() {
  53. DB = orig
  54. }()
  55. batch := &UserMigrationBatch{
  56. Name: "b1",
  57. SourceRegion: "ov",
  58. TargetRegion: "cn",
  59. OperatorId: 1,
  60. }
  61. require.NoError(t, CreateUserMigrationBatch(batch))
  62. first := &UserMigrationItem{
  63. BatchId: batch.Id,
  64. SourceUserId: 10000001,
  65. SourceUsername: "alice",
  66. SourceEmail: "alice@example.com",
  67. SourceQuota: 12345,
  68. ResolutionStrategy: "",
  69. Status: UserMigrationItemStatusConflict,
  70. }
  71. require.NoError(t, UpsertUserMigrationItem(first))
  72. require.NoError(t, DB.Model(&UserMigrationItem{}).
  73. Where("batch_id = ? AND source_user_id = ?", batch.Id, 10000001).
  74. Update("resolution_strategy", UserMigrationStrategyMergeExisting).Error)
  75. second := &UserMigrationItem{
  76. BatchId: batch.Id,
  77. SourceUserId: 10000001,
  78. SourceUsername: "alice",
  79. SourceEmail: "alice@example.com",
  80. SourceQuota: 99999,
  81. ResolutionStrategy: UserMigrationStrategyCreateNew,
  82. Status: UserMigrationItemStatusReady,
  83. }
  84. require.NoError(t, UpsertUserMigrationItem(second))
  85. items, err := ListUserMigrationItemsByBatch(batch.Id)
  86. require.NoError(t, err)
  87. require.Len(t, items, 1)
  88. require.Equal(t, UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy)
  89. require.Equal(t, UserMigrationItemStatusReady, items[0].Status)
  90. require.Equal(t, 99999, items[0].SourceQuota)
  91. }
  92. func TestAtomicGrantQuotaIfNotExists_Idempotent(t *testing.T) {
  93. db := setupUserMigrationTestDB(t)
  94. orig := DB
  95. DB = db
  96. defer func() {
  97. DB = orig
  98. }()
  99. require.NoError(t, db.Create(&User{
  100. Id: 1,
  101. Username: "cn-user",
  102. Password: "hash",
  103. AffCode: "AFF1",
  104. Quota: 500,
  105. }).Error)
  106. var alreadyDone bool
  107. err := db.Transaction(func(tx *gorm.DB) error {
  108. var innerErr error
  109. alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
  110. return innerErr
  111. })
  112. require.NoError(t, err)
  113. require.False(t, alreadyDone)
  114. var user User
  115. require.NoError(t, db.First(&user, 1).Error)
  116. require.Equal(t, 700, user.Quota)
  117. err = db.Transaction(func(tx *gorm.DB) error {
  118. var innerErr error
  119. alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
  120. return innerErr
  121. })
  122. require.NoError(t, err)
  123. require.True(t, alreadyDone)
  124. require.NoError(t, db.First(&user, 1).Error)
  125. require.Equal(t, 700, user.Quota)
  126. }
  127. func TestRefreshBatchStats_ReadyOnlyWhenNoPendingScan(t *testing.T) {
  128. db := setupUserMigrationTestDB(t)
  129. orig := DB
  130. DB = db
  131. defer func() {
  132. DB = orig
  133. }()
  134. batch := &UserMigrationBatch{Name: "stats-test", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
  135. require.NoError(t, CreateUserMigrationBatch(batch))
  136. require.NoError(t, DB.Create(&UserMigrationItem{
  137. BatchId: batch.Id,
  138. SourceUserId: 1,
  139. SourceUsername: "u1",
  140. Status: UserMigrationItemStatusReady,
  141. }).Error)
  142. require.NoError(t, DB.Create(&UserMigrationItem{
  143. BatchId: batch.Id,
  144. SourceUserId: 2,
  145. SourceUsername: "u2",
  146. Status: UserMigrationItemStatusPendingScan,
  147. }).Error)
  148. require.NoError(t, RefreshUserMigrationBatchStats(batch.Id))
  149. var refreshed UserMigrationBatch
  150. require.NoError(t, DB.First(&refreshed, batch.Id).Error)
  151. require.Equal(t, UserMigrationBatchStatusScanned, refreshed.Status)
  152. }
  153. func TestRefreshUserMigrationBatchExecutionStats_MigratingLeakHandled(t *testing.T) {
  154. db := setupUserMigrationTestDB(t)
  155. orig := DB
  156. DB = db
  157. defer func() {
  158. DB = orig
  159. }()
  160. batch := &UserMigrationBatch{
  161. Name: "migrating-leak",
  162. SourceRegion: "ov",
  163. TargetRegion: "cn",
  164. OperatorId: 1,
  165. Status: UserMigrationBatchStatusRunning,
  166. }
  167. require.NoError(t, CreateUserMigrationBatch(batch))
  168. require.NoError(t, DB.Create(&UserMigrationItem{
  169. BatchId: batch.Id,
  170. SourceUserId: 101,
  171. SourceUsername: "stuck-user",
  172. Status: UserMigrationItemStatusMigrating,
  173. }).Error)
  174. require.NoError(t, RefreshUserMigrationBatchExecutionStats(batch.Id))
  175. var refreshed UserMigrationBatch
  176. require.NoError(t, DB.First(&refreshed, batch.Id).Error)
  177. require.Equal(t, UserMigrationBatchStatusFailed, refreshed.Status)
  178. require.Equal(t, 0, refreshed.MigratedUsers)
  179. require.Equal(t, 0, refreshed.FailedUsers)
  180. }