api_key_service_update_fields_test.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "testing"
  6. "github.com/stretchr/testify/require"
  7. )
  8. // api_keys 的 quota_used / usage_5h|1d|7d 由计费热路径原子递增。
  9. // 编辑 Key(改名、换分组……)若整行回写,并发累计的用量就会被旧快照覆盖。
  10. // 这些用例锁死"只声明请求真正要改的列"。
  11. type updateFieldsAPIKeyRepoStub struct {
  12. quotaBaseAPIKeyRepoStub
  13. key *APIKey
  14. updateFields []APIKeyUpdateFields
  15. }
  16. // IncrementQuotaUsed 模拟计费热路径上的原子递增:只动 quota_used。
  17. func (s *updateFieldsAPIKeyRepoStub) IncrementQuotaUsed(_ context.Context, _ int64, amount float64) (float64, error) {
  18. s.key.QuotaUsed += amount
  19. return s.key.QuotaUsed, nil
  20. }
  21. func (s *updateFieldsAPIKeyRepoStub) GetByID(context.Context, int64) (*APIKey, error) {
  22. clone := *s.key
  23. return &clone, nil
  24. }
  25. func (s *updateFieldsAPIKeyRepoStub) Update(_ context.Context, _ *APIKey, fields APIKeyUpdateFields) error {
  26. s.updateFields = append(s.updateFields, fields)
  27. return nil
  28. }
  29. func newUpdateFieldsAPIKeyService(key *APIKey) (*APIKeyService, *updateFieldsAPIKeyRepoStub) {
  30. repo := &updateFieldsAPIKeyRepoStub{key: key}
  31. return &APIKeyService{apiKeyRepo: repo}, repo
  32. }
  33. func TestAPIKeyUpdate_OnlyDeclaresRequestedColumns(t *testing.T) {
  34. name := "renamed"
  35. quota := 500.0
  36. rateLimit := 42.0
  37. whitelist := []string{"10.0.0.1"}
  38. tests := []struct {
  39. name string
  40. req UpdateAPIKeyRequest
  41. want APIKeyUpdateFields
  42. }{
  43. {
  44. name: "name only",
  45. req: UpdateAPIKeyRequest{Name: &name},
  46. want: APIKeyUpdateFields{Name: true},
  47. },
  48. {
  49. name: "quota only",
  50. req: UpdateAPIKeyRequest{Quota: &quota},
  51. want: APIKeyUpdateFields{Quota: true},
  52. },
  53. {
  54. name: "rate limit threshold only",
  55. req: UpdateAPIKeyRequest{RateLimit5h: &rateLimit},
  56. want: APIKeyUpdateFields{RateLimits: true},
  57. },
  58. {
  59. name: "ip whitelist only",
  60. req: UpdateAPIKeyRequest{IPWhitelist: &whitelist},
  61. want: APIKeyUpdateFields{IPRules: true},
  62. },
  63. }
  64. for _, tt := range tests {
  65. t.Run(tt.name, func(t *testing.T) {
  66. svc, repo := newUpdateFieldsAPIKeyService(&APIKey{
  67. ID: 1,
  68. UserID: 7,
  69. Key: "sk-test",
  70. Name: "before",
  71. Status: StatusActive,
  72. Quota: 100,
  73. QuotaUsed: 30,
  74. Usage5h: 12,
  75. })
  76. _, err := svc.Update(context.Background(), 1, 7, tt.req)
  77. require.NoError(t, err)
  78. require.Equal(t, []APIKeyUpdateFields{tt.want}, repo.updateFields)
  79. })
  80. }
  81. }
  82. // 显式重置仍需声明对应的列,避免收窄写入列时把功能改坏。
  83. func TestAPIKeyUpdate_DeclaresUsageColumnsOnExplicitReset(t *testing.T) {
  84. reset := true
  85. svc, repo := newUpdateFieldsAPIKeyService(&APIKey{
  86. ID: 1, UserID: 7, Key: "sk-test", Status: StatusActive, Quota: 100, QuotaUsed: 30, Usage5h: 12,
  87. })
  88. _, err := svc.Update(context.Background(), 1, 7, UpdateAPIKeyRequest{
  89. ResetQuota: &reset,
  90. ResetRateLimitUsage: &reset,
  91. })
  92. require.NoError(t, err)
  93. require.Equal(t, []APIKeyUpdateFields{{QuotaUsed: true, RateLimitUsage: true}}, repo.updateFields)
  94. }
  95. // 配额扩容会顺带把 quota_exhausted 复活为 active,此时必须声明 status。
  96. func TestAPIKeyUpdate_DeclaresStatusWhenReactivated(t *testing.T) {
  97. quota := 500.0
  98. svc, repo := newUpdateFieldsAPIKeyService(&APIKey{
  99. ID: 1, UserID: 7, Key: "sk-test", Status: StatusAPIKeyQuotaExhausted, Quota: 100, QuotaUsed: 100,
  100. })
  101. _, err := svc.Update(context.Background(), 1, 7, UpdateAPIKeyRequest{Quota: &quota})
  102. require.NoError(t, err)
  103. require.Equal(t, []APIKeyUpdateFields{{Quota: true, Status: true}}, repo.updateFields)
  104. }
  105. // 计费热路径把 Key 标记为配额耗尽时只写 status,
  106. // 否则会把刚原子递增的 quota_used 按快照覆盖掉。
  107. func TestUpdateQuotaUsed_ExhaustedMarkOnlyDeclaresStatus(t *testing.T) {
  108. repo := &updateFieldsAPIKeyRepoStub{key: &APIKey{
  109. ID: 1, UserID: 7, Key: "sk-test", Status: StatusActive, Quota: 10, QuotaUsed: 10,
  110. }}
  111. svc := &APIKeyService{apiKeyRepo: repo}
  112. require.NoError(t, svc.UpdateQuotaUsed(context.Background(), 1, 5))
  113. require.Equal(t, []APIKeyUpdateFields{{Status: true}}, repo.updateFields)
  114. }