Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

216 wiersze
5.6 KiB

  1. package helpers
  2. import (
  3. "fmt"
  4. "net/http"
  5. "testing"
  6. "time"
  7. "github.com/QuantumNous/new-api/model"
  8. userMigration "github.com/QuantumNous/new-api/service/user_migration"
  9. "github.com/stretchr/testify/require"
  10. )
  11. type BatchDetail struct {
  12. Batch model.UserMigrationBatch
  13. Items []model.UserMigrationItem
  14. Total int64
  15. }
  16. type batchDetailResponse struct {
  17. Success bool `json:"success"`
  18. Data model.UserMigrationBatch `json:"data"`
  19. Items []model.UserMigrationItem `json:"items"`
  20. Total int64 `json:"total"`
  21. }
  22. type VerifyEntry struct {
  23. SourceUserID int `json:"source_user_id"`
  24. TargetUserID int `json:"target_user_id"`
  25. Result *userMigration.VerifyResult `json:"result,omitempty"`
  26. Error string `json:"error,omitempty"`
  27. }
  28. type verifyResponse struct {
  29. Success bool `json:"success"`
  30. Data []VerifyEntry `json:"data"`
  31. }
  32. type remoteUserResponse struct {
  33. Success bool `json:"success"`
  34. Data userMigration.RemoteUserSnapshot `json:"data"`
  35. }
  36. func GetBatchDetail(client *APIClient, batchID int) (*BatchDetail, error) {
  37. var resp batchDetailResponse
  38. if err := client.GetJSON(fmt.Sprintf("/api/user-migrations/batches/%d?page=0&page_size=100", batchID), http.StatusOK, &resp); err != nil {
  39. return nil, err
  40. }
  41. if !resp.Success {
  42. return nil, fmt.Errorf("get batch detail %d failed", batchID)
  43. }
  44. return &BatchDetail{
  45. Batch: resp.Data,
  46. Items: resp.Items,
  47. Total: resp.Total,
  48. }, nil
  49. }
  50. func WaitBatchStatus(client *APIClient, batchID int, expectedStatus string, timeout time.Duration) (*BatchDetail, error) {
  51. deadline := time.Now().Add(timeout)
  52. lastStatus := ""
  53. for time.Now().Before(deadline) {
  54. detail, err := GetBatchDetail(client, batchID)
  55. if err == nil {
  56. lastStatus = detail.Batch.Status
  57. if detail.Batch.Status == expectedStatus {
  58. return detail, nil
  59. }
  60. }
  61. time.Sleep(250 * time.Millisecond)
  62. }
  63. return nil, fmt.Errorf("batch %d did not reach status %q within %s, last status=%q", batchID, expectedStatus, timeout, lastStatus)
  64. }
  65. func VerifyBatch(client *APIClient, batchID int) ([]VerifyEntry, error) {
  66. var resp verifyResponse
  67. if err := client.PostJSON(fmt.Sprintf("/api/user-migrations/batches/%d/verify", batchID), map[string]any{}, http.StatusOK, &resp); err != nil {
  68. return nil, err
  69. }
  70. if !resp.Success {
  71. return nil, fmt.Errorf("verify batch %d failed", batchID)
  72. }
  73. return resp.Data, nil
  74. }
  75. func FindItemBySourceUserID(items []model.UserMigrationItem, sourceUserID int) *model.UserMigrationItem {
  76. for i := range items {
  77. if items[i].SourceUserId == sourceUserID {
  78. return &items[i]
  79. }
  80. }
  81. return nil
  82. }
  83. func GetRemoteUserSnapshot(baseURL string, userID int) (*userMigration.RemoteUserSnapshot, error) {
  84. client, err := NewAPIClient(baseURL)
  85. if err != nil {
  86. return nil, err
  87. }
  88. var resp remoteUserResponse
  89. if err := client.DoJSONWithHeaders(
  90. http.MethodGet,
  91. fmt.Sprintf("/api/internal/migration/users/%d", userID),
  92. nil,
  93. http.StatusOK,
  94. map[string]string{
  95. "X-Sync-API-Key": DefaultSyncAPIKey,
  96. "X-Sync-Node": "e2e-assert",
  97. },
  98. &resp,
  99. ); err != nil {
  100. return nil, err
  101. }
  102. if !resp.Success {
  103. return nil, fmt.Errorf("get remote user snapshot %d failed", userID)
  104. }
  105. return &resp.Data, nil
  106. }
  107. func AssertCNUserExists(t *testing.T, dbPath, username string) *model.User {
  108. t.Helper()
  109. db, err := OpenSQLite(dbPath)
  110. require.NoError(t, err)
  111. defer func() {
  112. require.NoError(t, CloseSQLite(db))
  113. }()
  114. var user model.User
  115. require.NoError(t, db.Where("username = ?", username).First(&user).Error)
  116. return &user
  117. }
  118. func LoadCNUserByEmail(t *testing.T, dbPath, email string) *model.User {
  119. t.Helper()
  120. db, err := OpenSQLite(dbPath)
  121. require.NoError(t, err)
  122. defer func() {
  123. require.NoError(t, CloseSQLite(db))
  124. }()
  125. var user model.User
  126. require.NoError(t, db.Where("email = ?", email).First(&user).Error)
  127. return &user
  128. }
  129. func AssertOVUserSynced(t *testing.T, ovBaseURL string, userID, targetUserID, syncedQuota int) {
  130. t.Helper()
  131. snapshot, err := GetRemoteUserSnapshot(ovBaseURL, userID)
  132. require.NoError(t, err)
  133. require.Equal(t, "synced", snapshot.Source)
  134. require.Equal(t, targetUserID, snapshot.RemoteUserId)
  135. require.Equal(t, syncedQuota, snapshot.SyncedQuota)
  136. }
  137. func AssertQuota(t *testing.T, dbPath string, userID, quota int) {
  138. t.Helper()
  139. db, err := OpenSQLite(dbPath)
  140. require.NoError(t, err)
  141. defer func() {
  142. require.NoError(t, CloseSQLite(db))
  143. }()
  144. var user model.User
  145. require.NoError(t, db.Select("id", "quota").First(&user, userID).Error)
  146. require.Equal(t, quota, user.Quota)
  147. }
  148. func AssertNoDuplicateQuotaGrant(t *testing.T, dbPath, migrationToken string) {
  149. t.Helper()
  150. db, err := OpenSQLite(dbPath)
  151. require.NoError(t, err)
  152. defer func() {
  153. require.NoError(t, CloseSQLite(db))
  154. }()
  155. var count int64
  156. require.NoError(t, db.Model(&model.MigrationQuotaGrant{}).Where("migration_token = ?", migrationToken).Count(&count).Error)
  157. require.EqualValues(t, 1, count)
  158. }
  159. func CountUsersByUsername(t *testing.T, dbPath, username string) int64 {
  160. t.Helper()
  161. db, err := OpenSQLite(dbPath)
  162. require.NoError(t, err)
  163. defer func() {
  164. require.NoError(t, CloseSQLite(db))
  165. }()
  166. var count int64
  167. require.NoError(t, db.Model(&model.User{}).Where("username = ?", username).Count(&count).Error)
  168. return count
  169. }
  170. func CountUsersByEmail(t *testing.T, dbPath, email string) int64 {
  171. t.Helper()
  172. db, err := OpenSQLite(dbPath)
  173. require.NoError(t, err)
  174. defer func() {
  175. require.NoError(t, CloseSQLite(db))
  176. }()
  177. var count int64
  178. require.NoError(t, db.Model(&model.User{}).Where("email = ?", email).Count(&count).Error)
  179. return count
  180. }