|
- package user_migration
-
- import (
- "fmt"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- type fakeClient struct {
- listPages map[int][]RemoteUserSnapshot
- getUserResp *RemoteUserSnapshot
- getUserByID map[int]*RemoteUserSnapshot
- hasCopyResult bool
- hasCopyByID map[int]bool
- convertCalls []int
- convertReqs map[int]*ConvertRemoteUserRequest
- }
-
- func (f *fakeClient) ListUsers(page, pageSize int) (*ListRemoteUsersResponse, error) {
- users := f.listPages[page]
- return &ListRemoteUsersResponse{
- Success: true,
- Data: users,
- Total: len(users),
- }, nil
- }
-
- func (f *fakeClient) GetUser(userId int) (*RemoteUserSnapshot, error) {
- if snapshot, ok := f.getUserByID[userId]; ok {
- return snapshot, nil
- }
- if f.getUserResp != nil {
- return f.getUserResp, nil
- }
- return nil, fmt.Errorf("user %d not found", userId)
- }
-
- func (f *fakeClient) ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error {
- f.convertCalls = append(f.convertCalls, userId)
- if f.convertReqs == nil {
- f.convertReqs = make(map[int]*ConvertRemoteUserRequest)
- }
- if req != nil {
- copyReq := *req
- f.convertReqs[userId] = ©Req
- if snapshot, ok := f.getUserByID[userId]; ok && snapshot != nil {
- snapshot.Source = common.UserSourceSynced
- snapshot.RemoteUserId = req.RemoteUserId
- }
- if f.getUserResp != nil && f.getUserResp.Id == userId {
- f.getUserResp.Source = common.UserSourceSynced
- f.getUserResp.RemoteUserId = req.RemoteUserId
- }
- }
- return nil
- }
-
- func (f *fakeClient) HasSyncedCopy(cnUserId int) (bool, error) {
- if value, ok := f.hasCopyByID[cnUserId]; ok {
- return value, nil
- }
- return f.hasCopyResult, nil
- }
-
- func setupMigrationServiceDB(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(
- &model.User{},
- &model.CustomOAuthProvider{},
- &model.UserOAuthBinding{},
- &model.UserMigrationBatch{},
- &model.UserMigrationItem{},
- &model.MigrationQuotaGrant{},
- ))
-
- return db
- }
-
- func TestScanBatch_NoConflict_MarksReady(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- batch := &model.UserMigrationBatch{Name: "wave-1", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- client := &fakeClient{
- listPages: map[int][]RemoteUserSnapshot{
- 0: {{
- Id: 10000001,
- Username: "ov-bob",
- Email: "bob@example.com",
- Quota: 456,
- AffCode: "A002",
- }},
- },
- }
- svc := &Service{client: client}
- require.NoError(t, svc.ScanBatch(batch.Id))
-
- items, err := model.ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, 1)
- require.Equal(t, model.UserMigrationStrategyCreateNew, items[0].ResolutionStrategy)
- require.Equal(t, model.UserMigrationItemStatusReady, items[0].Status)
-
- var savedBatch model.UserMigrationBatch
- require.NoError(t, db.First(&savedBatch, batch.Id).Error)
- require.Equal(t, model.UserMigrationBatchStatusReady, savedBatch.Status)
- }
-
- func TestScanBatch_EmailConflict_NoCNSyncedCopy_RecommendsMerge(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 1,
- Username: "cn-alice",
- Password: "hash",
- Email: "alice@example.com",
- AffCode: "C001",
- }).Error)
-
- batch := &model.UserMigrationBatch{Name: "wave-2", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- client := &fakeClient{
- listPages: map[int][]RemoteUserSnapshot{
- 0: {{
- Id: 10000001,
- Username: "ov-alice",
- Email: "alice@example.com",
- Quota: 123,
- AffCode: "A001",
- }},
- },
- hasCopyResult: false,
- }
- svc := &Service{client: client}
- require.NoError(t, svc.ScanBatch(batch.Id))
-
- items, err := model.ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, 1)
- require.Equal(t, model.UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy)
- require.Equal(t, model.UserMigrationItemStatusReady, items[0].Status)
- require.Equal(t, 1, items[0].TargetUserId)
- }
-
- func TestScanBatch_EmailConflict_CNHasSyncedCopy_MarksHardConflict(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 456,
- Username: "cn-user456",
- Password: "hash",
- Email: "shared@example.com",
- AffCode: "C456",
- }).Error)
-
- batch := &model.UserMigrationBatch{Name: "wave-3", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- client := &fakeClient{
- listPages: map[int][]RemoteUserSnapshot{
- 0: {{
- Id: 10000003,
- Username: "ov-x",
- Email: "shared@example.com",
- Quota: 50,
- AffCode: "OVX",
- }},
- },
- hasCopyResult: true,
- }
- svc := &Service{client: client}
- require.NoError(t, svc.ScanBatch(batch.Id))
-
- items, err := model.ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, 1)
- require.Equal(t, model.UserMigrationItemStatusConflict, items[0].Status)
- require.Contains(t, items[0].ConflictFlags, "cn_already_synced_to_ov")
- }
-
- func TestScanBatch_Paginated(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- batch := &model.UserMigrationBatch{Name: "wave-big", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- page0 := make([]RemoteUserSnapshot, scanPageSize)
- for i := range page0 {
- page0[i] = RemoteUserSnapshot{
- Id: 10000001 + i,
- Username: fmt.Sprintf("u%d", i),
- AffCode: fmt.Sprintf("A%d", i),
- }
- }
- page1 := []RemoteUserSnapshot{{
- Id: 10000001 + scanPageSize,
- Username: "last",
- AffCode: "LAST",
- }}
-
- client := &fakeClient{
- listPages: map[int][]RemoteUserSnapshot{
- 0: page0,
- 1: page1,
- },
- }
- svc := &Service{client: client}
- require.NoError(t, svc.ScanBatch(batch.Id))
-
- items, err := model.ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, scanPageSize+1)
- }
-
- func TestScanBatch_Rescan_PreservesHumanStrategy(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- batch := &model.UserMigrationBatch{Name: "wave-rescan", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1}
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- client := &fakeClient{
- listPages: map[int][]RemoteUserSnapshot{
- 0: {{
- Id: 10000001,
- Username: "ov-dave",
- Email: "dave@example.com",
- Quota: 100,
- AffCode: "OVD1",
- }},
- },
- hasCopyResult: false,
- }
- svc := &Service{client: client}
- require.NoError(t, svc.ScanBatch(batch.Id))
-
- require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
- Update("resolution_strategy", model.UserMigrationStrategySkip).Error)
-
- require.NoError(t, svc.ScanBatch(batch.Id))
-
- items, err := model.ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, 1)
- require.Equal(t, model.UserMigrationStrategySkip, items[0].ResolutionStrategy)
- }
-
- func TestExecuteBatch_CreateNew_CreatesUserAndConvertsRemote(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- batch := &model.UserMigrationBatch{
- Name: "wave-exec",
- Status: model.UserMigrationBatchStatusReady,
- }
- batch.SourceRegion = "ov"
- batch.TargetRegion = "cn"
- batch.OperatorId = 1
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- item := &model.UserMigrationItem{
- BatchId: batch.Id,
- SourceUserId: 10000001,
- SourceUsername: "ov-alice",
- SourceEmail: "alice@example.com",
- SourceQuota: 300,
- ResolutionStrategy: model.UserMigrationStrategyCreateNew,
- Status: model.UserMigrationItemStatusReady,
- MigrationToken: "batch-1-user-10000001",
- }
- require.NoError(t, model.UpsertUserMigrationItem(item))
-
- client := &fakeClient{
- getUserResp: &RemoteUserSnapshot{
- Id: 10000001,
- Username: "ov-alice",
- Password: "hashed-password",
- Email: "alice@example.com",
- DisplayName: "Alice",
- Status: 1,
- Role: 1,
- Group: "default",
- Quota: 300,
- AffCode: "OVA1",
- },
- }
- svc := &Service{client: client}
-
- locked, err := model.TryLockBatchForExecution(batch.Id)
- require.NoError(t, err)
- require.True(t, locked)
- require.NoError(t, svc.ExecuteBatch(batch.Id))
-
- var cnUser model.User
- require.NoError(t, db.Where("username = ?", "ov-alice").First(&cnUser).Error)
- require.Equal(t, 300, cnUser.Quota)
- require.Less(t, cnUser.Id, 10000000)
-
- require.Len(t, client.convertCalls, 1)
- require.Equal(t, 10000001, client.convertCalls[0])
- }
-
- func TestExecuteBatch_CreateNew_IdempotentOnRetry(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- batch := &model.UserMigrationBatch{
- Name: "wave-retry",
- Status: model.UserMigrationBatchStatusReady,
- }
- batch.SourceRegion = "ov"
- batch.TargetRegion = "cn"
- batch.OperatorId = 1
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- item := &model.UserMigrationItem{
- BatchId: batch.Id,
- SourceUserId: 10000002,
- SourceUsername: "ov-carol",
- SourceEmail: "carol@example.com",
- SourceQuota: 150,
- ResolutionStrategy: model.UserMigrationStrategyCreateNew,
- Status: model.UserMigrationItemStatusReady,
- MigrationToken: "batch-1-user-10000002",
- }
- require.NoError(t, model.UpsertUserMigrationItem(item))
-
- snapshot := &RemoteUserSnapshot{
- Id: 10000002,
- Username: "ov-carol",
- Password: "hash",
- Email: "carol@example.com",
- Quota: 150,
- AffCode: "OVC1",
- }
- client := &fakeClient{getUserResp: snapshot}
- svc := &Service{client: client}
-
- locked, err := model.TryLockBatchForExecution(batch.Id)
- require.NoError(t, err)
- require.True(t, locked)
- require.NoError(t, svc.ExecuteBatch(batch.Id))
-
- require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
- Update("status", model.UserMigrationBatchStatusFailed).Error)
- require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
- Update("status", model.UserMigrationItemStatusReady).Error)
-
- locked, err = model.TryLockBatchForExecution(batch.Id)
- require.NoError(t, err)
- require.True(t, locked)
- require.NoError(t, svc.ExecuteBatch(batch.Id))
-
- var count int64
- require.NoError(t, db.Model(&model.User{}).Where("username = ?", "ov-carol").Count(&count).Error)
- require.EqualValues(t, 1, count)
- require.Len(t, client.convertCalls, 1)
- }
-
- func TestExecuteBatch_MergeExisting_IdempotentQuota(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 1,
- Username: "cn-bob",
- Password: "hash",
- Email: "bob@example.com",
- AffCode: "CNB1",
- Quota: 500,
- }).Error)
-
- batch := &model.UserMigrationBatch{
- Name: "wave-merge",
- Status: model.UserMigrationBatchStatusReady,
- }
- batch.SourceRegion = "ov"
- batch.TargetRegion = "cn"
- batch.OperatorId = 1
- require.NoError(t, model.CreateUserMigrationBatch(batch))
-
- item := &model.UserMigrationItem{
- BatchId: batch.Id,
- SourceUserId: 10000002,
- TargetUserId: 1,
- SourceUsername: "ov-bob",
- SourceEmail: "bob@example.com",
- SourceQuota: 200,
- ResolutionStrategy: model.UserMigrationStrategyMergeExisting,
- Status: model.UserMigrationItemStatusReady,
- MigrationToken: "batch-1-user-10000002",
- }
- require.NoError(t, model.UpsertUserMigrationItem(item))
-
- client := &fakeClient{
- getUserResp: &RemoteUserSnapshot{
- Id: 10000002,
- Username: "ov-bob",
- Password: "hashed",
- Email: "bob@example.com",
- Quota: 200,
- },
- }
- svc := &Service{client: client}
-
- for i := 0; i < 2; i++ {
- require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
- Update("status", model.UserMigrationBatchStatusReady).Error)
- require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
- Update("status", model.UserMigrationItemStatusReady).Error)
-
- locked, err := model.TryLockBatchForExecution(batch.Id)
- require.NoError(t, err)
- require.True(t, locked)
- require.NoError(t, svc.ExecuteBatch(batch.Id))
- }
-
- var saved model.User
- require.NoError(t, db.First(&saved, 1).Error)
- require.Equal(t, 700, saved.Quota)
- require.Len(t, client.convertCalls, 1)
- }
-
- func TestExecuteBatch_SkipItem_DoesNotConvert(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- batch := &model.UserMigrationBatch{
- Name: "skip-test",
- SourceRegion: "ov",
- TargetRegion: "cn",
- OperatorId: 1,
- Status: model.UserMigrationBatchStatusReady,
- }
- require.NoError(t, model.CreateUserMigrationBatch(batch))
- require.NoError(t, model.UpsertUserMigrationItem(&model.UserMigrationItem{
- BatchId: batch.Id,
- SourceUserId: 10000099,
- SourceUsername: "skip-user",
- SourceQuota: 50,
- ResolutionStrategy: model.UserMigrationStrategySkip,
- Status: model.UserMigrationItemStatusReady,
- MigrationToken: "batch-1-user-10000099",
- }))
-
- client := &fakeClient{}
- svc := &Service{client: client}
-
- locked, err := model.TryLockBatchForExecution(batch.Id)
- require.NoError(t, err)
- require.True(t, locked)
- require.NoError(t, svc.ExecuteBatch(batch.Id))
-
- require.Len(t, client.convertCalls, 0)
- items, err := model.ListUserMigrationItemsByBatch(batch.Id)
- require.NoError(t, err)
- require.Len(t, items, 1)
- require.Equal(t, model.UserMigrationItemStatusSkipped, items[0].Status)
- }
-
- func TestVerifyItem_CorrectlyChecksRemoteConverted(t *testing.T) {
- db := setupMigrationServiceDB(t)
- orig := model.DB
- model.DB = db
- defer func() {
- model.DB = orig
- }()
-
- require.NoError(t, db.Create(&model.User{
- Id: 789,
- Username: "cn-migrated",
- Password: "hash",
- AffCode: "CM1",
- Quota: 300,
- }).Error)
-
- item := &model.UserMigrationItem{
- BatchId: 1,
- SourceUserId: 10000001,
- TargetUserId: 789,
- SourceQuota: 300,
- Status: model.UserMigrationItemStatusMigrated,
- }
-
- client := &fakeClient{
- getUserResp: &RemoteUserSnapshot{
- Id: 10000001,
- Source: common.UserSourceSynced,
- RemoteUserId: 789,
- },
- }
- svc := &Service{client: client}
- result, err := svc.VerifyItem(item)
- require.NoError(t, err)
- require.True(t, result.TargetUserExists)
- require.True(t, result.RemoteConverted)
- require.True(t, result.QuotaMatched)
- }
|