|
- package model
-
- import (
- "testing"
-
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupUserMigrationTestDB(t *testing.T) *gorm.DB {
- t.Helper()
-
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
-
- sqlDB, err := db.DB()
- require.NoError(t, err)
- sqlDB.SetMaxOpenConns(1)
-
- require.NoError(t, db.AutoMigrate(
- &User{},
- &CustomOAuthProvider{},
- &UserOAuthBinding{},
- &UserMigrationBatch{},
- &UserMigrationItem{},
- &MigrationQuotaGrant{},
- ))
-
- return db
- }
-
- func TestUserMigrationModels_AutoMigrate(t *testing.T) {
- db := setupUserMigrationTestDB(t)
-
- require.True(t, db.Migrator().HasTable(&UserMigrationBatch{}))
- require.True(t, db.Migrator().HasTable(&UserMigrationItem{}))
- require.True(t, db.Migrator().HasTable(&MigrationQuotaGrant{}))
- }
-
- func TestCreateUserMigrationBatch_DefaultStatus(t *testing.T) {
- db := setupUserMigrationTestDB(t)
- orig := DB
- DB = db
- defer func() {
- DB = orig
- }()
-
- batch := &UserMigrationBatch{
- Name: "ov-to-cn-wave-1",
- SourceRegion: "overseas",
- TargetRegion: "cn",
- OperatorId: 1,
- }
-
- require.NoError(t, CreateUserMigrationBatch(batch))
- require.Equal(t, UserMigrationBatchStatusDraft, batch.Status)
- require.Greater(t, batch.Id, 0)
- }
-
- func TestUpsertMigrationItem_PreservesHumanSetStrategy(t *testing.T) {
- db := setupUserMigrationTestDB(t)
- orig := DB
- DB = db
- defer func() {
- DB = orig
- }()
-
- batch := &UserMigrationBatch{
- Name: "b1",
- SourceRegion: "ov",
- TargetRegion: "cn",
- OperatorId: 1,
- }
- require.NoError(t, CreateUserMigrationBatch(batch))
-
- first := &UserMigrationItem{
- BatchId: batch.Id,
- SourceUserId: 10000001,
- SourceUsername: "alice",
- SourceEmail: "alice@example.com",
- SourceQuota: 12345,
- ResolutionStrategy: "",
- Status: UserMigrationItemStatusConflict,
- }
- require.NoError(t, UpsertUserMigrationItem(first))
-
- require.NoError(t, DB.Model(&UserMigrationItem{}).
- Where("batch_id = ? AND source_user_id = ?", batch.Id, 10000001).
- Update("resolution_strategy", UserMigrationStrategyMergeExisting).Error)
-
- second := &UserMigrationItem{
- BatchId: batch.Id,
- SourceUserId: 10000001,
- SourceUsername: "alice",
- SourceEmail: "alice@example.com",
- SourceQuota: 99999,
- ResolutionStrategy: UserMigrationStrategyCreateNew,
- Status: UserMigrationItemStatusReady,
- }
- require.NoError(t, UpsertUserMigrationItem(second))
-
- items, err := ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, 1)
- require.Equal(t, UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy)
- require.Equal(t, UserMigrationItemStatusReady, items[0].Status)
- require.Equal(t, 99999, items[0].SourceQuota)
- }
-
- func TestAtomicGrantQuotaIfNotExists_Idempotent(t *testing.T) {
- db := setupUserMigrationTestDB(t)
- orig := DB
- DB = db
- defer func() {
- DB = orig
- }()
-
- require.NoError(t, db.Create(&User{
- Id: 1,
- Username: "cn-user",
- Password: "hash",
- AffCode: "AFF1",
- Quota: 500,
- }).Error)
-
- var alreadyDone bool
- err := db.Transaction(func(tx *gorm.DB) error {
- var innerErr error
- alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
- return innerErr
- })
- require.NoError(t, err)
- require.False(t, alreadyDone)
-
- var user User
- require.NoError(t, db.First(&user, 1).Error)
- require.Equal(t, 700, user.Quota)
-
- err = db.Transaction(func(tx *gorm.DB) error {
- var innerErr error
- alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
- return innerErr
- })
- require.NoError(t, err)
- require.True(t, alreadyDone)
-
- require.NoError(t, db.First(&user, 1).Error)
- require.Equal(t, 700, user.Quota)
- }
|