|
|
@@ -4,6 +4,7 @@ import ( |
|
|
"fmt" |
|
|
"fmt" |
|
|
"testing" |
|
|
"testing" |
|
|
|
|
|
|
|
|
|
|
|
"github.com/QuantumNous/new-api/common" |
|
|
"github.com/QuantumNous/new-api/model" |
|
|
"github.com/QuantumNous/new-api/model" |
|
|
"github.com/glebarez/sqlite" |
|
|
"github.com/glebarez/sqlite" |
|
|
"github.com/stretchr/testify/require" |
|
|
"github.com/stretchr/testify/require" |
|
|
@@ -47,6 +48,14 @@ func (f *fakeClient) ConvertToSynced(userId int, req *ConvertRemoteUserRequest) |
|
|
if req != nil { |
|
|
if req != nil { |
|
|
copyReq := *req |
|
|
copyReq := *req |
|
|
f.convertReqs[userId] = ©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 |
|
|
return nil |
|
|
} |
|
|
} |
|
|
@@ -237,3 +246,230 @@ func TestScanBatch_Paginated(t *testing.T) { |
|
|
require.NoError(t, err) |
|
|
require.NoError(t, err) |
|
|
require.Len(t, items, scanPageSize+1) |
|
|
require.Len(t, items, scanPageSize+1) |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
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 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) |
|
|
|
|
|
} |