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

391 строка
10 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. require.NoError(t, db.Create(&model.User{
  54. Id: 10000003,
  55. Username: "root",
  56. Password: "hash",
  57. AffCode: "A3",
  58. Source: common.UserSourceLocal,
  59. Role: common.RoleRootUser,
  60. }).Error)
  61. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users?page=0&page_size=20", nil)
  62. w := httptest.NewRecorder()
  63. router.ServeHTTP(w, req)
  64. require.Equal(t, http.StatusOK, w.Code)
  65. var body map[string]any
  66. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  67. require.True(t, body["success"].(bool))
  68. require.EqualValues(t, 1, body["total"])
  69. }
  70. func TestGetMigrationUser_RootExcluded(t *testing.T) {
  71. gin.SetMode(gin.TestMode)
  72. router := gin.New()
  73. router.Use(testSyncAuthMiddleware())
  74. router.GET("/api/internal/migration/users/:id", GetMigrationUser)
  75. db := setupInternalMigrationDB(t)
  76. orig := model.DB
  77. model.DB = db
  78. defer func() {
  79. model.DB = orig
  80. }()
  81. require.NoError(t, db.Create(&model.User{
  82. Id: 10000011,
  83. Username: "root",
  84. Password: "hash",
  85. AffCode: "ROOT1",
  86. Source: common.UserSourceLocal,
  87. Role: common.RoleRootUser,
  88. }).Error)
  89. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/10000011", nil)
  90. w := httptest.NewRecorder()
  91. router.ServeHTTP(w, req)
  92. require.Equal(t, http.StatusNotFound, w.Code)
  93. }
  94. func TestGetMigrationUser_SyncedUserStillReadableForVerify(t *testing.T) {
  95. gin.SetMode(gin.TestMode)
  96. router := gin.New()
  97. router.Use(testSyncAuthMiddleware())
  98. router.GET("/api/internal/migration/users/:id", GetMigrationUser)
  99. db := setupInternalMigrationDB(t)
  100. orig := model.DB
  101. model.DB = db
  102. defer func() {
  103. model.DB = orig
  104. }()
  105. require.NoError(t, db.Create(&model.User{
  106. Id: 10000021,
  107. Username: "synced-user",
  108. Password: "hash",
  109. AffCode: "SYNC1",
  110. Source: common.UserSourceSynced,
  111. RemoteUserId: 88,
  112. SyncedQuota: 123,
  113. Role: common.RoleCommonUser,
  114. }).Error)
  115. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/10000021", nil)
  116. w := httptest.NewRecorder()
  117. router.ServeHTTP(w, req)
  118. require.Equal(t, http.StatusOK, w.Code)
  119. var body map[string]any
  120. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  121. require.True(t, body["success"].(bool))
  122. }
  123. func TestGetMigrationUser_InvalidID(t *testing.T) {
  124. gin.SetMode(gin.TestMode)
  125. router := gin.New()
  126. router.Use(testSyncAuthMiddleware())
  127. router.GET("/api/internal/migration/users/:id", GetMigrationUser)
  128. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/not-a-number", nil)
  129. w := httptest.NewRecorder()
  130. router.ServeHTTP(w, req)
  131. require.Equal(t, http.StatusBadRequest, w.Code)
  132. }
  133. func TestConvertMigrationUserToSynced_Idempotent(t *testing.T) {
  134. gin.SetMode(gin.TestMode)
  135. router := gin.New()
  136. router.Use(testSyncAuthMiddleware())
  137. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  138. db := setupInternalMigrationDB(t)
  139. orig := model.DB
  140. model.DB = db
  141. defer func() {
  142. model.DB = orig
  143. }()
  144. require.NoError(t, db.Create(&model.User{
  145. Id: 10000002,
  146. Username: "ov-2",
  147. Password: "hash",
  148. AffCode: "A2",
  149. Source: common.UserSourceLocal,
  150. Quota: 66,
  151. }).Error)
  152. body := `{"remote_user_id":2002,"synced_quota":66}`
  153. for i := 0; i < 2; i++ {
  154. req := httptest.NewRequest(http.MethodPost,
  155. "/api/internal/migration/users/10000002/convert-to-synced",
  156. strings.NewReader(body))
  157. req.Header.Set("Content-Type", "application/json")
  158. w := httptest.NewRecorder()
  159. router.ServeHTTP(w, req)
  160. require.Equal(t, http.StatusOK, w.Code)
  161. }
  162. }
  163. func TestConvertMigrationUserToSynced_ConflictRemoteId(t *testing.T) {
  164. gin.SetMode(gin.TestMode)
  165. router := gin.New()
  166. router.Use(testSyncAuthMiddleware())
  167. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  168. db := setupInternalMigrationDB(t)
  169. orig := model.DB
  170. model.DB = db
  171. defer func() {
  172. model.DB = orig
  173. }()
  174. require.NoError(t, db.Create(&model.User{
  175. Id: 10000003,
  176. Username: "ov-3",
  177. Password: "hash",
  178. AffCode: "A3",
  179. Source: common.UserSourceSynced,
  180. RemoteUserId: 5000,
  181. }).Error)
  182. body := `{"remote_user_id":9999,"synced_quota":100}`
  183. req := httptest.NewRequest(http.MethodPost,
  184. "/api/internal/migration/users/10000003/convert-to-synced",
  185. strings.NewReader(body))
  186. req.Header.Set("Content-Type", "application/json")
  187. w := httptest.NewRecorder()
  188. router.ServeHTTP(w, req)
  189. require.Equal(t, http.StatusConflict, w.Code)
  190. }
  191. func TestConvertMigrationUserToSynced_ConflictUsesSentinelErrors(t *testing.T) {
  192. gin.SetMode(gin.TestMode)
  193. router := gin.New()
  194. router.Use(testSyncAuthMiddleware())
  195. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  196. db := setupInternalMigrationDB(t)
  197. orig := model.DB
  198. model.DB = db
  199. defer func() {
  200. model.DB = orig
  201. }()
  202. require.NoError(t, db.Create(&model.User{
  203. Id: 10000013,
  204. Username: "ov-13",
  205. Password: "hash",
  206. AffCode: "A13",
  207. Source: common.UserSourceSynced,
  208. RemoteUserId: 5001,
  209. }).Error)
  210. body := `{"remote_user_id":9999,"synced_quota":100}`
  211. req := httptest.NewRequest(http.MethodPost,
  212. "/api/internal/migration/users/10000013/convert-to-synced",
  213. strings.NewReader(body))
  214. req.Header.Set("Content-Type", "application/json")
  215. w := httptest.NewRecorder()
  216. router.ServeHTTP(w, req)
  217. require.Equal(t, http.StatusConflict, w.Code)
  218. }
  219. func TestConvertMigrationUserToSynced_RejectsRootUser(t *testing.T) {
  220. gin.SetMode(gin.TestMode)
  221. router := gin.New()
  222. router.Use(testSyncAuthMiddleware())
  223. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  224. db := setupInternalMigrationDB(t)
  225. orig := model.DB
  226. model.DB = db
  227. defer func() {
  228. model.DB = orig
  229. }()
  230. require.NoError(t, db.Create(&model.User{
  231. Id: 10000012,
  232. Username: "root",
  233. Password: "hash",
  234. AffCode: "ROOT2",
  235. Source: common.UserSourceLocal,
  236. Role: common.RoleRootUser,
  237. }).Error)
  238. body := `{"remote_user_id":9999,"synced_quota":100}`
  239. req := httptest.NewRequest(http.MethodPost,
  240. "/api/internal/migration/users/10000012/convert-to-synced",
  241. strings.NewReader(body))
  242. req.Header.Set("Content-Type", "application/json")
  243. w := httptest.NewRecorder()
  244. router.ServeHTTP(w, req)
  245. require.Equal(t, http.StatusConflict, w.Code)
  246. }
  247. func TestConvertMigrationUserToSynced_InvalidRemoteUserID(t *testing.T) {
  248. gin.SetMode(gin.TestMode)
  249. router := gin.New()
  250. router.Use(testSyncAuthMiddleware())
  251. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  252. db := setupInternalMigrationDB(t)
  253. orig := model.DB
  254. model.DB = db
  255. defer func() {
  256. model.DB = orig
  257. }()
  258. require.NoError(t, db.Create(&model.User{
  259. Id: 10000014,
  260. Username: "ov-14",
  261. Password: "hash",
  262. AffCode: "A14",
  263. Source: common.UserSourceLocal,
  264. }).Error)
  265. body := `{"remote_user_id":0,"synced_quota":100}`
  266. req := httptest.NewRequest(http.MethodPost,
  267. "/api/internal/migration/users/10000014/convert-to-synced",
  268. strings.NewReader(body))
  269. req.Header.Set("Content-Type", "application/json")
  270. w := httptest.NewRecorder()
  271. router.ServeHTTP(w, req)
  272. require.Equal(t, http.StatusBadRequest, w.Code)
  273. }
  274. func TestConvertMigrationUserToSynced_PersistsSyncedFields(t *testing.T) {
  275. gin.SetMode(gin.TestMode)
  276. router := gin.New()
  277. router.Use(testSyncAuthMiddleware())
  278. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  279. db := setupInternalMigrationDB(t)
  280. orig := model.DB
  281. model.DB = db
  282. defer func() {
  283. model.DB = orig
  284. }()
  285. require.NoError(t, db.Create(&model.User{
  286. Id: 10000015,
  287. Username: "ov-15",
  288. Password: "hash",
  289. AffCode: "A15",
  290. Source: common.UserSourceLocal,
  291. Quota: 88,
  292. }).Error)
  293. body := `{"remote_user_id":2015,"synced_quota":166}`
  294. req := httptest.NewRequest(http.MethodPost,
  295. "/api/internal/migration/users/10000015/convert-to-synced",
  296. strings.NewReader(body))
  297. req.Header.Set("Content-Type", "application/json")
  298. w := httptest.NewRecorder()
  299. router.ServeHTTP(w, req)
  300. require.Equal(t, http.StatusOK, w.Code)
  301. var saved model.User
  302. require.NoError(t, db.First(&saved, 10000015).Error)
  303. require.Equal(t, common.UserSourceSynced, saved.Source)
  304. require.Equal(t, 2015, saved.RemoteUserId)
  305. require.Equal(t, 166, saved.SyncedQuota)
  306. require.Greater(t, saved.LastSyncAt, int64(0))
  307. }
  308. func TestCheckSyncedCopy(t *testing.T) {
  309. gin.SetMode(gin.TestMode)
  310. router := gin.New()
  311. router.Use(testSyncAuthMiddleware())
  312. router.GET("/api/internal/migration/users/check-synced-copy", CheckSyncedCopy)
  313. db := setupInternalMigrationDB(t)
  314. orig := model.DB
  315. model.DB = db
  316. defer func() {
  317. model.DB = orig
  318. }()
  319. require.NoError(t, db.Create(&model.User{
  320. Id: 10000010,
  321. Username: "synced",
  322. Password: "hash",
  323. AffCode: "S1",
  324. Source: common.UserSourceSynced,
  325. RemoteUserId: 456,
  326. }).Error)
  327. req := httptest.NewRequest(http.MethodGet,
  328. "/api/internal/migration/users/check-synced-copy?cn_user_id=456", nil)
  329. w := httptest.NewRecorder()
  330. router.ServeHTTP(w, req)
  331. require.Equal(t, http.StatusOK, w.Code)
  332. var body map[string]any
  333. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  334. require.True(t, body["has_copy"].(bool))
  335. }