package model import ( "encoding/json" "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.Equal(t, "all", batch.SelectionMode) require.Greater(t, batch.Id, 0) } func TestCreateUserMigrationBatch_StoresExplicitIDsSelection(t *testing.T) { db := setupUserMigrationTestDB(t) orig := DB DB = db defer func() { DB = orig }() payload := map[string]any{"source_user_ids": []int{1001, 1002, 1003}} payloadJSON, err := json.Marshal(payload) require.NoError(t, err) batch := &UserMigrationBatch{ Name: "partial-wave", SourceRegion: "ov", TargetRegion: "cn", SelectionMode: "explicit_ids", SelectionPayload: string(payloadJSON), RequestedUserCount: 3, OperatorId: 1, } require.NoError(t, CreateUserMigrationBatch(batch)) var saved UserMigrationBatch require.NoError(t, DB.First(&saved, batch.Id).Error) require.Equal(t, "explicit_ids", saved.SelectionMode) require.Equal(t, string(payloadJSON), saved.SelectionPayload) require.Equal(t, 3, saved.RequestedUserCount) require.Equal(t, 0, saved.MatchedUserCount) require.Equal(t, 0, saved.ExcludedUserCount) } 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) } func TestRefreshBatchStats_ReadyOnlyWhenNoPendingScan(t *testing.T) { db := setupUserMigrationTestDB(t) orig := DB DB = db defer func() { DB = orig }() batch := &UserMigrationBatch{Name: "stats-test", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1} require.NoError(t, CreateUserMigrationBatch(batch)) require.NoError(t, DB.Create(&UserMigrationItem{ BatchId: batch.Id, SourceUserId: 1, SourceUsername: "u1", Status: UserMigrationItemStatusReady, }).Error) require.NoError(t, DB.Create(&UserMigrationItem{ BatchId: batch.Id, SourceUserId: 2, SourceUsername: "u2", Status: UserMigrationItemStatusPendingScan, }).Error) require.NoError(t, RefreshUserMigrationBatchStats(batch.Id)) var refreshed UserMigrationBatch require.NoError(t, DB.First(&refreshed, batch.Id).Error) require.Equal(t, UserMigrationBatchStatusScanned, refreshed.Status) } func TestRefreshUserMigrationBatchExecutionStats_MigratingLeakHandled(t *testing.T) { db := setupUserMigrationTestDB(t) orig := DB DB = db defer func() { DB = orig }() batch := &UserMigrationBatch{ Name: "migrating-leak", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1, Status: UserMigrationBatchStatusRunning, } require.NoError(t, CreateUserMigrationBatch(batch)) require.NoError(t, DB.Create(&UserMigrationItem{ BatchId: batch.Id, SourceUserId: 101, SourceUsername: "stuck-user", Status: UserMigrationItemStatusMigrating, }).Error) require.NoError(t, RefreshUserMigrationBatchExecutionStats(batch.Id)) var refreshed UserMigrationBatch require.NoError(t, DB.First(&refreshed, batch.Id).Error) require.Equal(t, UserMigrationBatchStatusFailed, refreshed.Status) require.Equal(t, 0, refreshed.MigratedUsers) require.Equal(t, 0, refreshed.FailedUsers) } func TestCancelUserMigrationBatch_ReleasesBatch(t *testing.T) { db := setupUserMigrationTestDB(t) orig := DB DB = db defer func() { DB = orig }() batch := &UserMigrationBatch{ Name: "cancel-batch", SourceRegion: "ov", TargetRegion: "cn", SelectionMode: "explicit_ids", SelectionPayload: `{"source_user_ids":[1001,1002]}`, OperatorId: 1, Status: UserMigrationBatchStatusScanned, } require.NoError(t, CreateUserMigrationBatch(batch)) conflicts, err := FindActiveMigrationUserConflicts([]int{1001, 1002}) require.NoError(t, err) require.Len(t, conflicts, 2) cancelled, err := CancelUserMigrationBatch(batch.Id) require.NoError(t, err) require.True(t, cancelled) var refreshed UserMigrationBatch require.NoError(t, DB.First(&refreshed, batch.Id).Error) require.Equal(t, UserMigrationBatchStatusCancelled, refreshed.Status) conflicts, err = FindActiveMigrationUserConflicts([]int{1001, 1002}) require.NoError(t, err) require.Len(t, conflicts, 0) }