Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

173 строки
4.4 KiB

  1. package controller
  2. import (
  3. "encoding/json"
  4. "net/http"
  5. "net/http/httptest"
  6. "strings"
  7. "testing"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/gin-gonic/gin"
  11. "github.com/glebarez/sqlite"
  12. "github.com/stretchr/testify/require"
  13. "gorm.io/gorm"
  14. )
  15. func setupInternalMigrationDB(t *testing.T) *gorm.DB {
  16. t.Helper()
  17. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  18. require.NoError(t, err)
  19. sqlDB, err := db.DB()
  20. require.NoError(t, err)
  21. sqlDB.SetMaxOpenConns(1)
  22. require.NoError(t, db.AutoMigrate(&model.User{}, &model.CustomOAuthProvider{}, &model.UserOAuthBinding{}))
  23. return db
  24. }
  25. func TestListMigrationUsers_ReturnsLocalOnly(t *testing.T) {
  26. gin.SetMode(gin.TestMode)
  27. router := gin.New()
  28. router.Use(testSyncAuthMiddleware())
  29. router.GET("/api/internal/migration/users", ListMigrationUsers)
  30. db := setupInternalMigrationDB(t)
  31. orig := model.DB
  32. model.DB = db
  33. defer func() {
  34. model.DB = orig
  35. }()
  36. require.NoError(t, db.Create(&model.User{
  37. Id: 10000001,
  38. Username: "ov-1",
  39. Password: "hash",
  40. Email: "u1@example.com",
  41. AffCode: "A1",
  42. Source: common.UserSourceLocal,
  43. Quota: 99,
  44. }).Error)
  45. require.NoError(t, db.Create(&model.User{
  46. Id: 10000002,
  47. Username: "ov-2",
  48. Password: "hash",
  49. AffCode: "A2",
  50. Source: common.UserSourceSynced,
  51. RemoteUserId: 1,
  52. }).Error)
  53. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users?page=0&page_size=20", nil)
  54. w := httptest.NewRecorder()
  55. router.ServeHTTP(w, req)
  56. require.Equal(t, http.StatusOK, w.Code)
  57. var body map[string]any
  58. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  59. require.True(t, body["success"].(bool))
  60. require.EqualValues(t, 1, body["total"])
  61. }
  62. func TestConvertMigrationUserToSynced_Idempotent(t *testing.T) {
  63. gin.SetMode(gin.TestMode)
  64. router := gin.New()
  65. router.Use(testSyncAuthMiddleware())
  66. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  67. db := setupInternalMigrationDB(t)
  68. orig := model.DB
  69. model.DB = db
  70. defer func() {
  71. model.DB = orig
  72. }()
  73. require.NoError(t, db.Create(&model.User{
  74. Id: 10000002,
  75. Username: "ov-2",
  76. Password: "hash",
  77. AffCode: "A2",
  78. Source: common.UserSourceLocal,
  79. Quota: 66,
  80. }).Error)
  81. body := `{"remote_user_id":2002,"synced_quota":66}`
  82. for i := 0; i < 2; i++ {
  83. req := httptest.NewRequest(http.MethodPost,
  84. "/api/internal/migration/users/10000002/convert-to-synced",
  85. strings.NewReader(body))
  86. req.Header.Set("Content-Type", "application/json")
  87. w := httptest.NewRecorder()
  88. router.ServeHTTP(w, req)
  89. require.Equal(t, http.StatusOK, w.Code)
  90. }
  91. }
  92. func TestConvertMigrationUserToSynced_ConflictRemoteId(t *testing.T) {
  93. gin.SetMode(gin.TestMode)
  94. router := gin.New()
  95. router.Use(testSyncAuthMiddleware())
  96. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  97. db := setupInternalMigrationDB(t)
  98. orig := model.DB
  99. model.DB = db
  100. defer func() {
  101. model.DB = orig
  102. }()
  103. require.NoError(t, db.Create(&model.User{
  104. Id: 10000003,
  105. Username: "ov-3",
  106. Password: "hash",
  107. AffCode: "A3",
  108. Source: common.UserSourceSynced,
  109. RemoteUserId: 5000,
  110. }).Error)
  111. body := `{"remote_user_id":9999,"synced_quota":100}`
  112. req := httptest.NewRequest(http.MethodPost,
  113. "/api/internal/migration/users/10000003/convert-to-synced",
  114. strings.NewReader(body))
  115. req.Header.Set("Content-Type", "application/json")
  116. w := httptest.NewRecorder()
  117. router.ServeHTTP(w, req)
  118. require.Equal(t, http.StatusConflict, w.Code)
  119. }
  120. func TestCheckSyncedCopy(t *testing.T) {
  121. gin.SetMode(gin.TestMode)
  122. router := gin.New()
  123. router.Use(testSyncAuthMiddleware())
  124. router.GET("/api/internal/migration/users/check-synced-copy", CheckSyncedCopy)
  125. db := setupInternalMigrationDB(t)
  126. orig := model.DB
  127. model.DB = db
  128. defer func() {
  129. model.DB = orig
  130. }()
  131. require.NoError(t, db.Create(&model.User{
  132. Id: 10000010,
  133. Username: "synced",
  134. Password: "hash",
  135. AffCode: "S1",
  136. Source: common.UserSourceSynced,
  137. RemoteUserId: 456,
  138. }).Error)
  139. req := httptest.NewRequest(http.MethodGet,
  140. "/api/internal/migration/users/check-synced-copy?cn_user_id=456", nil)
  141. w := httptest.NewRecorder()
  142. router.ServeHTTP(w, req)
  143. require.Equal(t, http.StatusOK, w.Code)
  144. var body map[string]any
  145. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  146. require.True(t, body["has_copy"].(bool))
  147. }