Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.
 
 
 

151 rader
3.7 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. }