|
- 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
- }
|