api_key_repo_integration_test.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "sync"
  6. "testing"
  7. "time"
  8. dbent "github.com/Wei-Shaw/sub2api/ent"
  9. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  10. "github.com/Wei-Shaw/sub2api/internal/service"
  11. "github.com/stretchr/testify/require"
  12. "github.com/stretchr/testify/suite"
  13. )
  14. type APIKeyRepoSuite struct {
  15. suite.Suite
  16. ctx context.Context
  17. client *dbent.Client
  18. repo *apiKeyRepository
  19. }
  20. func (s *APIKeyRepoSuite) SetupTest() {
  21. s.ctx = context.Background()
  22. tx := testEntTx(s.T())
  23. s.client = tx.Client()
  24. s.repo = newAPIKeyRepositoryWithSQL(s.client, tx)
  25. }
  26. func TestAPIKeyRepoSuite(t *testing.T) {
  27. suite.Run(t, new(APIKeyRepoSuite))
  28. }
  29. // --- Create / GetByID / GetByKey ---
  30. func (s *APIKeyRepoSuite) TestCreate() {
  31. user := s.mustCreateUser("create@test.com")
  32. key := &service.APIKey{
  33. UserID: user.ID,
  34. Key: "sk-create-test",
  35. Name: "Test Key",
  36. Status: service.StatusActive,
  37. }
  38. err := s.repo.Create(s.ctx, key)
  39. s.Require().NoError(err, "Create")
  40. s.Require().NotZero(key.ID, "expected ID to be set")
  41. got, err := s.repo.GetByID(s.ctx, key.ID)
  42. s.Require().NoError(err, "GetByID")
  43. s.Require().Equal("sk-create-test", got.Key)
  44. }
  45. func (s *APIKeyRepoSuite) TestGetByID_NotFound() {
  46. _, err := s.repo.GetByID(s.ctx, 999999)
  47. s.Require().Error(err, "expected error for non-existent ID")
  48. }
  49. func (s *APIKeyRepoSuite) TestGetByKey() {
  50. user := s.mustCreateUser("getbykey@test.com")
  51. group := s.mustCreateGroup("g-key")
  52. key := &service.APIKey{
  53. UserID: user.ID,
  54. Key: "sk-getbykey",
  55. Name: "My Key",
  56. GroupID: &group.ID,
  57. Status: service.StatusActive,
  58. }
  59. s.Require().NoError(s.repo.Create(s.ctx, key))
  60. got, err := s.repo.GetByKey(s.ctx, key.Key)
  61. s.Require().NoError(err, "GetByKey")
  62. s.Require().Equal(key.ID, got.ID)
  63. s.Require().NotNil(got.User, "expected User preload")
  64. s.Require().Equal(user.ID, got.User.ID)
  65. s.Require().NotNil(got.Group, "expected Group preload")
  66. s.Require().Equal(group.ID, got.Group.ID)
  67. }
  68. func (s *APIKeyRepoSuite) TestGetByKey_NotFound() {
  69. _, err := s.repo.GetByKey(s.ctx, "non-existent-key")
  70. s.Require().Error(err, "expected error for non-existent key")
  71. }
  72. func (s *APIKeyRepoSuite) TestGetByKeyForAuth_PreservesMessagesDispatchModelConfig() {
  73. user := s.mustCreateUser("getbykey-auth-dispatch@test.com")
  74. group, err := s.client.Group.Create().
  75. SetName("g-auth-dispatch").
  76. SetPlatform(service.PlatformOpenAI).
  77. SetStatus(service.StatusActive).
  78. SetSubscriptionType(service.SubscriptionTypeStandard).
  79. SetRateMultiplier(1).
  80. SetAllowMessagesDispatch(true).
  81. SetDefaultMappedModel("gpt-5.4").
  82. SetMessagesDispatchModelConfig(service.OpenAIMessagesDispatchModelConfig{
  83. OpusMappedModel: "gpt-5.4-nano",
  84. SonnetMappedModel: "gpt-5.3-codex",
  85. HaikuMappedModel: "gpt-5.4-mini",
  86. ExactModelMappings: map[string]string{
  87. "claude-sonnet-4.5": "gpt-5.4-nano",
  88. },
  89. }).
  90. Save(s.ctx)
  91. s.Require().NoError(err)
  92. key := &service.APIKey{
  93. UserID: user.ID,
  94. Key: "sk-getbykey-auth-dispatch",
  95. Name: "Dispatch Key",
  96. GroupID: &group.ID,
  97. Status: service.StatusActive,
  98. }
  99. s.Require().NoError(s.repo.Create(s.ctx, key))
  100. got, err := s.repo.GetByKeyForAuth(s.ctx, key.Key)
  101. s.Require().NoError(err)
  102. s.Require().NotNil(got.Group)
  103. s.Require().True(got.Group.AllowMessagesDispatch)
  104. s.Require().Equal("gpt-5.4", got.Group.DefaultMappedModel)
  105. s.Require().Equal("gpt-5.4-nano", got.Group.MessagesDispatchModelConfig.OpusMappedModel)
  106. s.Require().Equal("gpt-5.4-nano", got.Group.MessagesDispatchModelConfig.ExactModelMappings["claude-sonnet-4.5"])
  107. }
  108. // --- Update ---
  109. func (s *APIKeyRepoSuite) TestUpdate() {
  110. user := s.mustCreateUser("update@test.com")
  111. key := &service.APIKey{
  112. UserID: user.ID,
  113. Key: "sk-update",
  114. Name: "Original",
  115. Status: service.StatusActive,
  116. }
  117. s.Require().NoError(s.repo.Create(s.ctx, key))
  118. key.Name = "Renamed"
  119. key.Status = service.StatusDisabled
  120. err := s.repo.Update(s.ctx, key, service.APIKeyUpdateFields{Name: true, Status: true})
  121. s.Require().NoError(err, "Update")
  122. got, err := s.repo.GetByID(s.ctx, key.ID)
  123. s.Require().NoError(err, "GetByID after update")
  124. s.Require().Equal("sk-update", got.Key, "Update should not change key")
  125. s.Require().Equal(user.ID, got.UserID, "Update should not change user_id")
  126. s.Require().Equal("Renamed", got.Name)
  127. s.Require().Equal(service.StatusDisabled, got.Status)
  128. }
  129. func (s *APIKeyRepoSuite) TestUpdate_ClearGroupID() {
  130. user := s.mustCreateUser("cleargroup@test.com")
  131. group := s.mustCreateGroup("g-clear")
  132. key := &service.APIKey{
  133. UserID: user.ID,
  134. Key: "sk-clear-group",
  135. Name: "Group Key",
  136. GroupID: &group.ID,
  137. Status: service.StatusActive,
  138. }
  139. s.Require().NoError(s.repo.Create(s.ctx, key))
  140. key.GroupID = nil
  141. err := s.repo.Update(s.ctx, key, service.APIKeyUpdateFields{GroupID: true})
  142. s.Require().NoError(err, "Update")
  143. got, err := s.repo.GetByID(s.ctx, key.ID)
  144. s.Require().NoError(err)
  145. s.Require().Nil(got.GroupID, "expected GroupID to be cleared")
  146. }
  147. // --- Delete ---
  148. func (s *APIKeyRepoSuite) TestDelete() {
  149. user := s.mustCreateUser("delete@test.com")
  150. key := &service.APIKey{
  151. UserID: user.ID,
  152. Key: "sk-delete",
  153. Name: "Delete Me",
  154. Status: service.StatusActive,
  155. }
  156. s.Require().NoError(s.repo.Create(s.ctx, key))
  157. err := s.repo.Delete(s.ctx, key.ID)
  158. s.Require().NoError(err, "Delete")
  159. _, err = s.repo.GetByID(s.ctx, key.ID)
  160. s.Require().Error(err, "expected error after delete")
  161. }
  162. func (s *APIKeyRepoSuite) TestCreate_AfterSoftDelete_AllowsSameKey() {
  163. user := s.mustCreateUser("recreate-after-soft-delete@test.com")
  164. const reusedKey = "sk-reuse-after-soft-delete"
  165. first := &service.APIKey{
  166. UserID: user.ID,
  167. Key: reusedKey,
  168. Name: "First Key",
  169. Status: service.StatusActive,
  170. }
  171. s.Require().NoError(s.repo.Create(s.ctx, first), "create first key")
  172. s.Require().NoError(s.repo.Delete(s.ctx, first.ID), "soft delete first key")
  173. second := &service.APIKey{
  174. UserID: user.ID,
  175. Key: reusedKey,
  176. Name: "Second Key",
  177. Status: service.StatusActive,
  178. }
  179. s.Require().NoError(s.repo.Create(s.ctx, second), "create second key with same key")
  180. s.Require().NotZero(second.ID)
  181. s.Require().NotEqual(first.ID, second.ID, "recreated key should be a new row")
  182. }
  183. // --- ListByUserID / CountByUserID ---
  184. func (s *APIKeyRepoSuite) TestListByUserID() {
  185. user := s.mustCreateUser("listbyuser@test.com")
  186. s.mustCreateApiKey(user.ID, "sk-list-1", "Key 1", nil)
  187. s.mustCreateApiKey(user.ID, "sk-list-2", "Key 2", nil)
  188. keys, page, err := s.repo.ListByUserID(s.ctx, user.ID, pagination.PaginationParams{Page: 1, PageSize: 10}, service.APIKeyListFilters{})
  189. s.Require().NoError(err, "ListByUserID")
  190. s.Require().Len(keys, 2)
  191. s.Require().Equal(int64(2), page.Total)
  192. }
  193. func (s *APIKeyRepoSuite) TestListByUserID_Pagination() {
  194. user := s.mustCreateUser("paging@test.com")
  195. for i := 0; i < 5; i++ {
  196. s.mustCreateApiKey(user.ID, "sk-page-"+string(rune('a'+i)), "Key", nil)
  197. }
  198. keys, page, err := s.repo.ListByUserID(s.ctx, user.ID, pagination.PaginationParams{Page: 1, PageSize: 2}, service.APIKeyListFilters{})
  199. s.Require().NoError(err)
  200. s.Require().Len(keys, 2)
  201. s.Require().Equal(int64(5), page.Total)
  202. s.Require().Equal(3, page.Pages)
  203. }
  204. func (s *APIKeyRepoSuite) TestCountByUserID() {
  205. user := s.mustCreateUser("count@test.com")
  206. s.mustCreateApiKey(user.ID, "sk-count-1", "K1", nil)
  207. s.mustCreateApiKey(user.ID, "sk-count-2", "K2", nil)
  208. count, err := s.repo.CountByUserID(s.ctx, user.ID)
  209. s.Require().NoError(err, "CountByUserID")
  210. s.Require().Equal(int64(2), count)
  211. }
  212. // --- ListByGroupID / CountByGroupID ---
  213. func (s *APIKeyRepoSuite) TestListByGroupID() {
  214. user := s.mustCreateUser("listbygroup@test.com")
  215. group := s.mustCreateGroup("g-list")
  216. s.mustCreateApiKey(user.ID, "sk-grp-1", "K1", &group.ID)
  217. s.mustCreateApiKey(user.ID, "sk-grp-2", "K2", &group.ID)
  218. s.mustCreateApiKey(user.ID, "sk-grp-3", "K3", nil) // no group
  219. keys, page, err := s.repo.ListByGroupID(s.ctx, group.ID, pagination.PaginationParams{Page: 1, PageSize: 10})
  220. s.Require().NoError(err, "ListByGroupID")
  221. s.Require().Len(keys, 2)
  222. s.Require().Equal(int64(2), page.Total)
  223. // User preloaded
  224. s.Require().NotNil(keys[0].User)
  225. }
  226. func (s *APIKeyRepoSuite) TestCountByGroupID() {
  227. user := s.mustCreateUser("countgroup@test.com")
  228. group := s.mustCreateGroup("g-count")
  229. s.mustCreateApiKey(user.ID, "sk-gc-1", "K1", &group.ID)
  230. count, err := s.repo.CountByGroupID(s.ctx, group.ID)
  231. s.Require().NoError(err, "CountByGroupID")
  232. s.Require().Equal(int64(1), count)
  233. }
  234. // --- ExistsByKey ---
  235. func (s *APIKeyRepoSuite) TestExistsByKey() {
  236. user := s.mustCreateUser("exists@test.com")
  237. s.mustCreateApiKey(user.ID, "sk-exists", "K", nil)
  238. exists, err := s.repo.ExistsByKey(s.ctx, "sk-exists")
  239. s.Require().NoError(err, "ExistsByKey")
  240. s.Require().True(exists)
  241. notExists, err := s.repo.ExistsByKey(s.ctx, "sk-not-exists")
  242. s.Require().NoError(err)
  243. s.Require().False(notExists)
  244. }
  245. // --- SearchAPIKeys ---
  246. func (s *APIKeyRepoSuite) TestSearchAPIKeys() {
  247. user := s.mustCreateUser("search@test.com")
  248. s.mustCreateApiKey(user.ID, "sk-search-1", "Production Key", nil)
  249. s.mustCreateApiKey(user.ID, "sk-search-2", "Development Key", nil)
  250. found, err := s.repo.SearchAPIKeys(s.ctx, user.ID, "prod", 10)
  251. s.Require().NoError(err, "SearchAPIKeys")
  252. s.Require().Len(found, 1)
  253. s.Require().Contains(found[0].Name, "Production")
  254. }
  255. func (s *APIKeyRepoSuite) TestSearchAPIKeys_NoKeyword() {
  256. user := s.mustCreateUser("searchnokw@test.com")
  257. s.mustCreateApiKey(user.ID, "sk-nk-1", "K1", nil)
  258. s.mustCreateApiKey(user.ID, "sk-nk-2", "K2", nil)
  259. found, err := s.repo.SearchAPIKeys(s.ctx, user.ID, "", 10)
  260. s.Require().NoError(err)
  261. s.Require().Len(found, 2)
  262. }
  263. func (s *APIKeyRepoSuite) TestSearchAPIKeys_NoUserID() {
  264. user := s.mustCreateUser("searchnouid@test.com")
  265. s.mustCreateApiKey(user.ID, "sk-nu-1", "TestKey", nil)
  266. found, err := s.repo.SearchAPIKeys(s.ctx, 0, "testkey", 10)
  267. s.Require().NoError(err)
  268. s.Require().Len(found, 1)
  269. }
  270. // --- ClearGroupIDByGroupID ---
  271. func (s *APIKeyRepoSuite) TestClearGroupIDByGroupID() {
  272. user := s.mustCreateUser("cleargrp@test.com")
  273. group := s.mustCreateGroup("g-clear-bulk")
  274. k1 := s.mustCreateApiKey(user.ID, "sk-clr-1", "K1", &group.ID)
  275. k2 := s.mustCreateApiKey(user.ID, "sk-clr-2", "K2", &group.ID)
  276. s.mustCreateApiKey(user.ID, "sk-clr-3", "K3", nil) // no group
  277. affected, err := s.repo.ClearGroupIDByGroupID(s.ctx, group.ID)
  278. s.Require().NoError(err, "ClearGroupIDByGroupID")
  279. s.Require().Equal(int64(2), affected)
  280. got1, _ := s.repo.GetByID(s.ctx, k1.ID)
  281. got2, _ := s.repo.GetByID(s.ctx, k2.ID)
  282. s.Require().Nil(got1.GroupID)
  283. s.Require().Nil(got2.GroupID)
  284. count, _ := s.repo.CountByGroupID(s.ctx, group.ID)
  285. s.Require().Zero(count)
  286. }
  287. // --- Combined CRUD/Search/ClearGroupID (original test preserved as integration) ---
  288. func (s *APIKeyRepoSuite) TestCRUD_Search_ClearGroupID() {
  289. user := s.mustCreateUser("k@example.com")
  290. group := s.mustCreateGroup("g-k")
  291. key := s.mustCreateApiKey(user.ID, "sk-test-1", "My Key", &group.ID)
  292. key.GroupID = &group.ID
  293. got, err := s.repo.GetByKey(s.ctx, key.Key)
  294. s.Require().NoError(err, "GetByKey")
  295. s.Require().Equal(key.ID, got.ID)
  296. s.Require().NotNil(got.User)
  297. s.Require().Equal(user.ID, got.User.ID)
  298. s.Require().NotNil(got.Group)
  299. s.Require().Equal(group.ID, got.Group.ID)
  300. key.Name = "Renamed"
  301. key.Status = service.StatusDisabled
  302. key.GroupID = nil
  303. s.Require().NoError(s.repo.Update(s.ctx, key, service.APIKeyUpdateFields{Name: true, Status: true, GroupID: true}), "Update")
  304. got2, err := s.repo.GetByID(s.ctx, key.ID)
  305. s.Require().NoError(err, "GetByID")
  306. s.Require().Equal("sk-test-1", got2.Key, "Update should not change key")
  307. s.Require().Equal(user.ID, got2.UserID, "Update should not change user_id")
  308. s.Require().Equal("Renamed", got2.Name)
  309. s.Require().Equal(service.StatusDisabled, got2.Status)
  310. s.Require().Nil(got2.GroupID)
  311. keys, page, err := s.repo.ListByUserID(s.ctx, user.ID, pagination.PaginationParams{Page: 1, PageSize: 10}, service.APIKeyListFilters{})
  312. s.Require().NoError(err, "ListByUserID")
  313. s.Require().Equal(int64(1), page.Total)
  314. s.Require().Len(keys, 1)
  315. exists, err := s.repo.ExistsByKey(s.ctx, "sk-test-1")
  316. s.Require().NoError(err, "ExistsByKey")
  317. s.Require().True(exists, "expected key to exist")
  318. found, err := s.repo.SearchAPIKeys(s.ctx, user.ID, "renam", 10)
  319. s.Require().NoError(err, "SearchAPIKeys")
  320. s.Require().Len(found, 1)
  321. s.Require().Equal(key.ID, found[0].ID)
  322. // ClearGroupIDByGroupID
  323. k2 := s.mustCreateApiKey(user.ID, "sk-test-2", "Group Key", &group.ID)
  324. k2.GroupID = &group.ID
  325. countBefore, err := s.repo.CountByGroupID(s.ctx, group.ID)
  326. s.Require().NoError(err, "CountByGroupID")
  327. s.Require().Equal(int64(1), countBefore, "expected 1 key in group before clear")
  328. affected, err := s.repo.ClearGroupIDByGroupID(s.ctx, group.ID)
  329. s.Require().NoError(err, "ClearGroupIDByGroupID")
  330. s.Require().Equal(int64(1), affected, "expected 1 affected row")
  331. got3, err := s.repo.GetByID(s.ctx, k2.ID)
  332. s.Require().NoError(err, "GetByID")
  333. s.Require().Nil(got3.GroupID, "expected GroupID cleared")
  334. countAfter, err := s.repo.CountByGroupID(s.ctx, group.ID)
  335. s.Require().NoError(err, "CountByGroupID after clear")
  336. s.Require().Equal(int64(0), countAfter, "expected 0 keys in group after clear")
  337. }
  338. func (s *APIKeyRepoSuite) mustCreateUser(email string) *service.User {
  339. s.T().Helper()
  340. u, err := s.client.User.Create().
  341. SetEmail(email).
  342. SetPasswordHash("test-password-hash").
  343. SetStatus(service.StatusActive).
  344. SetRole(service.RoleUser).
  345. Save(s.ctx)
  346. s.Require().NoError(err, "create user")
  347. return userEntityToService(u)
  348. }
  349. func (s *APIKeyRepoSuite) mustCreateGroup(name string) *service.Group {
  350. s.T().Helper()
  351. g, err := s.client.Group.Create().
  352. SetName(name).
  353. SetStatus(service.StatusActive).
  354. Save(s.ctx)
  355. s.Require().NoError(err, "create group")
  356. return groupEntityToService(g)
  357. }
  358. func (s *APIKeyRepoSuite) mustCreateApiKey(userID int64, key, name string, groupID *int64) *service.APIKey {
  359. s.T().Helper()
  360. k := &service.APIKey{
  361. UserID: userID,
  362. Key: key,
  363. Name: name,
  364. GroupID: groupID,
  365. Status: service.StatusActive,
  366. }
  367. s.Require().NoError(s.repo.Create(s.ctx, k), "create api key")
  368. return k
  369. }
  370. // --- IncrementQuotaUsed ---
  371. func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_Basic() {
  372. user := s.mustCreateUser("incr-basic@test.com")
  373. key := s.mustCreateApiKey(user.ID, "sk-incr-basic", "Incr", nil)
  374. newQuota, err := s.repo.IncrementQuotaUsed(s.ctx, key.ID, 1.5)
  375. s.Require().NoError(err, "IncrementQuotaUsed")
  376. s.Require().Equal(1.5, newQuota, "第一次递增后应为 1.5")
  377. newQuota, err = s.repo.IncrementQuotaUsed(s.ctx, key.ID, 2.5)
  378. s.Require().NoError(err, "IncrementQuotaUsed second")
  379. s.Require().Equal(4.0, newQuota, "第二次递增后应为 4.0")
  380. }
  381. func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_NotFound() {
  382. _, err := s.repo.IncrementQuotaUsed(s.ctx, 999999, 1.0)
  383. s.Require().ErrorIs(err, service.ErrAPIKeyNotFound, "不存在的 key 应返回 ErrAPIKeyNotFound")
  384. }
  385. func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_DeletedKey() {
  386. user := s.mustCreateUser("incr-deleted@test.com")
  387. key := s.mustCreateApiKey(user.ID, "sk-incr-del", "Deleted", nil)
  388. s.Require().NoError(s.repo.Delete(s.ctx, key.ID), "Delete")
  389. _, err := s.repo.IncrementQuotaUsed(s.ctx, key.ID, 1.0)
  390. s.Require().ErrorIs(err, service.ErrAPIKeyNotFound, "已删除的 key 应返回 ErrAPIKeyNotFound")
  391. }
  392. func (s *APIKeyRepoSuite) TestIncrementQuotaUsedAndGetState() {
  393. user := s.mustCreateUser("quota-state@test.com")
  394. key := s.mustCreateApiKey(user.ID, "sk-quota-state", "QuotaState", nil)
  395. key.Quota = 3
  396. key.QuotaUsed = 1
  397. s.Require().NoError(s.repo.Update(s.ctx, key, service.APIKeyUpdateFields{Quota: true, QuotaUsed: true}), "Update quota")
  398. state, err := s.repo.IncrementQuotaUsedAndGetState(s.ctx, key.ID, 2.5)
  399. s.Require().NoError(err, "IncrementQuotaUsedAndGetState")
  400. s.Require().NotNil(state)
  401. s.Require().Equal(3.5, state.QuotaUsed)
  402. s.Require().Equal(3.0, state.Quota)
  403. s.Require().Equal(service.StatusAPIKeyQuotaExhausted, state.Status)
  404. s.Require().Equal(key.Key, state.Key)
  405. got, err := s.repo.GetByID(s.ctx, key.ID)
  406. s.Require().NoError(err, "GetByID")
  407. s.Require().Equal(3.5, got.QuotaUsed)
  408. s.Require().Equal(service.StatusAPIKeyQuotaExhausted, got.Status)
  409. }
  410. // TestIncrementQuotaUsed_Concurrent 使用真实数据库验证并发原子性。
  411. // 注意:此测试使用 testEntClient(非事务隔离),数据会真正写入数据库。
  412. func TestIncrementQuotaUsed_Concurrent(t *testing.T) {
  413. client := testEntClient(t)
  414. repo := NewAPIKeyRepository(client, integrationDB).(*apiKeyRepository)
  415. ctx := context.Background()
  416. // 创建测试用户和 API Key
  417. u, err := client.User.Create().
  418. SetEmail("concurrent-incr-" + time.Now().Format(time.RFC3339Nano) + "@test.com").
  419. SetPasswordHash("hash").
  420. SetStatus(service.StatusActive).
  421. SetRole(service.RoleUser).
  422. Save(ctx)
  423. require.NoError(t, err, "create user")
  424. k := &service.APIKey{
  425. UserID: u.ID,
  426. Key: "sk-concurrent-" + time.Now().Format(time.RFC3339Nano),
  427. Name: "Concurrent",
  428. Status: service.StatusActive,
  429. }
  430. require.NoError(t, repo.Create(ctx, k), "create api key")
  431. t.Cleanup(func() {
  432. _ = client.APIKey.DeleteOneID(k.ID).Exec(ctx)
  433. _ = client.User.DeleteOneID(u.ID).Exec(ctx)
  434. })
  435. // 10 个 goroutine 各递增 1.0,总计应为 10.0
  436. const goroutines = 10
  437. const increment = 1.0
  438. var wg sync.WaitGroup
  439. errs := make([]error, goroutines)
  440. for i := 0; i < goroutines; i++ {
  441. wg.Add(1)
  442. go func(idx int) {
  443. defer wg.Done()
  444. _, errs[idx] = repo.IncrementQuotaUsed(ctx, k.ID, increment)
  445. }(i)
  446. }
  447. wg.Wait()
  448. for i, e := range errs {
  449. require.NoError(t, e, "goroutine %d failed", i)
  450. }
  451. // 验证最终结果
  452. got, err := repo.GetByID(ctx, k.ID)
  453. require.NoError(t, err, "GetByID")
  454. require.Equal(t, float64(goroutines)*increment, got.QuotaUsed,
  455. "并发递增后总和应为 %v,实际为 %v", float64(goroutines)*increment, got.QuotaUsed)
  456. }
  457. func (s *APIKeyRepoSuite) TestDeleteWithAudit_TombstonesWithoutRetainingCredential() {
  458. user := s.mustCreateUser("delwithaudit@test.com")
  459. key := &service.APIKey{
  460. UserID: user.ID,
  461. Key: "sk-del-audit-1",
  462. Name: "Audit Me",
  463. Status: service.StatusActive,
  464. }
  465. s.Require().NoError(s.repo.Create(s.ctx, key))
  466. s.Require().NoError(s.repo.DeleteWithAudit(s.ctx, key.ID))
  467. _, err := s.repo.GetByID(s.ctx, key.ID)
  468. s.Require().Error(err)
  469. var tombstone string
  470. var deletedAt time.Time
  471. rows, err := s.repo.sql.QueryContext(s.ctx, `SELECT key, deleted_at FROM api_keys WHERE id = $1`, key.ID)
  472. s.Require().NoError(err)
  473. s.Require().True(rows.Next())
  474. s.Require().NoError(rows.Scan(&tombstone, &deletedAt))
  475. s.Require().NoError(rows.Close())
  476. s.Require().NotEqual("sk-del-audit-1", tombstone)
  477. s.Require().Contains(tombstone, "__deleted__")
  478. var auditCount int
  479. auditRows, err := s.repo.sql.QueryContext(s.ctx,
  480. `SELECT COUNT(*) FROM deleted_api_key_audits WHERE api_key_id = $1`, key.ID)
  481. s.Require().NoError(err)
  482. s.Require().True(auditRows.Next())
  483. s.Require().NoError(auditRows.Scan(&auditCount))
  484. s.Require().NoError(auditRows.Close())
  485. s.Require().Zero(auditCount, "deleted credentials must not be retained")
  486. }
  487. func (s *APIKeyRepoSuite) TestDeleteWithAudit_RepeatIsIdempotent() {
  488. user := s.mustCreateUser("delwithaudit-idem@test.com")
  489. key := &service.APIKey{UserID: user.ID, Key: "sk-del-audit-2", Name: "K", Status: service.StatusActive}
  490. s.Require().NoError(s.repo.Create(s.ctx, key))
  491. s.Require().NoError(s.repo.DeleteWithAudit(s.ctx, key.ID))
  492. s.Require().NoError(s.repo.DeleteWithAudit(s.ctx, key.ID))
  493. }
  494. func (s *APIKeyRepoSuite) TestDeleteWithAudit_NotFound() {
  495. err := s.repo.DeleteWithAudit(s.ctx, 999999)
  496. s.Require().ErrorIs(err, service.ErrAPIKeyNotFound)
  497. }