package helpers import ( "fmt" "net/http" "testing" "time" "github.com/QuantumNous/new-api/model" userMigration "github.com/QuantumNous/new-api/service/user_migration" "github.com/stretchr/testify/require" ) type BatchDetail struct { Batch model.UserMigrationBatch Items []model.UserMigrationItem Total int64 } type batchDetailResponse struct { Success bool `json:"success"` Data model.UserMigrationBatch `json:"data"` Items []model.UserMigrationItem `json:"items"` Total int64 `json:"total"` } type VerifyEntry struct { SourceUserID int `json:"source_user_id"` TargetUserID int `json:"target_user_id"` Result *userMigration.VerifyResult `json:"result,omitempty"` Error string `json:"error,omitempty"` } type verifyResponse struct { Success bool `json:"success"` Data []VerifyEntry `json:"data"` } type remoteUserResponse struct { Success bool `json:"success"` Data userMigration.RemoteUserSnapshot `json:"data"` } func GetBatchDetail(client *APIClient, batchID int) (*BatchDetail, error) { var resp batchDetailResponse if err := client.GetJSON(fmt.Sprintf("/api/user-migrations/batches/%d?page=0&page_size=100", batchID), http.StatusOK, &resp); err != nil { return nil, err } if !resp.Success { return nil, fmt.Errorf("get batch detail %d failed", batchID) } return &BatchDetail{ Batch: resp.Data, Items: resp.Items, Total: resp.Total, }, nil } func WaitBatchStatus(client *APIClient, batchID int, expectedStatus string, timeout time.Duration) (*BatchDetail, error) { deadline := time.Now().Add(timeout) lastStatus := "" for time.Now().Before(deadline) { detail, err := GetBatchDetail(client, batchID) if err == nil { lastStatus = detail.Batch.Status if detail.Batch.Status == expectedStatus { return detail, nil } } time.Sleep(250 * time.Millisecond) } return nil, fmt.Errorf("batch %d did not reach status %q within %s, last status=%q", batchID, expectedStatus, timeout, lastStatus) } func VerifyBatch(client *APIClient, batchID int) ([]VerifyEntry, error) { var resp verifyResponse if err := client.PostJSON(fmt.Sprintf("/api/user-migrations/batches/%d/verify", batchID), map[string]any{}, http.StatusOK, &resp); err != nil { return nil, err } if !resp.Success { return nil, fmt.Errorf("verify batch %d failed", batchID) } return resp.Data, nil } func FindItemBySourceUserID(items []model.UserMigrationItem, sourceUserID int) *model.UserMigrationItem { for i := range items { if items[i].SourceUserId == sourceUserID { return &items[i] } } return nil } func GetRemoteUserSnapshot(baseURL string, userID int) (*userMigration.RemoteUserSnapshot, error) { client, err := NewAPIClient(baseURL) if err != nil { return nil, err } var resp remoteUserResponse if err := client.DoJSONWithHeaders( http.MethodGet, fmt.Sprintf("/api/internal/migration/users/%d", userID), nil, http.StatusOK, map[string]string{ "X-Sync-API-Key": DefaultSyncAPIKey, "X-Sync-Node": "e2e-assert", }, &resp, ); err != nil { return nil, err } if !resp.Success { return nil, fmt.Errorf("get remote user snapshot %d failed", userID) } return &resp.Data, nil } func AssertCNUserExists(t *testing.T, dbPath, username string) *model.User { t.Helper() db, err := OpenSQLite(dbPath) require.NoError(t, err) defer func() { require.NoError(t, CloseSQLite(db)) }() var user model.User require.NoError(t, db.Where("username = ?", username).First(&user).Error) return &user } func LoadCNUserByEmail(t *testing.T, dbPath, email string) *model.User { t.Helper() db, err := OpenSQLite(dbPath) require.NoError(t, err) defer func() { require.NoError(t, CloseSQLite(db)) }() var user model.User require.NoError(t, db.Where("email = ?", email).First(&user).Error) return &user } func AssertOVUserSynced(t *testing.T, ovBaseURL string, userID, targetUserID, syncedQuota int) { t.Helper() snapshot, err := GetRemoteUserSnapshot(ovBaseURL, userID) require.NoError(t, err) require.Equal(t, "synced", snapshot.Source) require.Equal(t, targetUserID, snapshot.RemoteUserId) require.Equal(t, syncedQuota, snapshot.SyncedQuota) } func AssertQuota(t *testing.T, dbPath string, userID, quota int) { t.Helper() db, err := OpenSQLite(dbPath) require.NoError(t, err) defer func() { require.NoError(t, CloseSQLite(db)) }() var user model.User require.NoError(t, db.Select("id", "quota").First(&user, userID).Error) require.Equal(t, quota, user.Quota) } func AssertNoDuplicateQuotaGrant(t *testing.T, dbPath, migrationToken string) { t.Helper() db, err := OpenSQLite(dbPath) require.NoError(t, err) defer func() { require.NoError(t, CloseSQLite(db)) }() var count int64 require.NoError(t, db.Model(&model.MigrationQuotaGrant{}).Where("migration_token = ?", migrationToken).Count(&count).Error) require.EqualValues(t, 1, count) } func CountUsersByUsername(t *testing.T, dbPath, username string) int64 { t.Helper() db, err := OpenSQLite(dbPath) require.NoError(t, err) defer func() { require.NoError(t, CloseSQLite(db)) }() var count int64 require.NoError(t, db.Model(&model.User{}).Where("username = ?", username).Count(&count).Error) return count } func CountUsersByEmail(t *testing.T, dbPath, email string) int64 { t.Helper() db, err := OpenSQLite(dbPath) require.NoError(t, err) defer func() { require.NoError(t, CloseSQLite(db)) }() var count int64 require.NoError(t, db.Model(&model.User{}).Where("email = ?", email).Count(&count).Error) return count }