| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694 |
- //go:build unit
- package service
- import (
- "context"
- "errors"
- "strings"
- "sync"
- "sync/atomic"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
- "github.com/redis/go-redis/v9"
- "github.com/stretchr/testify/require"
- )
- type authRepoStub struct {
- getByKeyForAuth func(ctx context.Context, key string) (*APIKey, error)
- listKeysByUserID func(ctx context.Context, userID int64) ([]string, error)
- listKeysByGroupID func(ctx context.Context, groupID int64) ([]string, error)
- }
- func (s *authRepoStub) Create(ctx context.Context, key *APIKey) error {
- panic("unexpected Create call")
- }
- func (s *authRepoStub) GetByID(ctx context.Context, id int64) (*APIKey, error) {
- panic("unexpected GetByID call")
- }
- func (s *authRepoStub) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
- panic("unexpected GetKeyAndOwnerID call")
- }
- func (s *authRepoStub) GetByKey(ctx context.Context, key string) (*APIKey, error) {
- panic("unexpected GetByKey call")
- }
- func (s *authRepoStub) GetByKeyForAuth(ctx context.Context, key string) (*APIKey, error) {
- if s.getByKeyForAuth == nil {
- panic("unexpected GetByKeyForAuth call")
- }
- return s.getByKeyForAuth(ctx, key)
- }
- func (s *authRepoStub) Update(ctx context.Context, key *APIKey, _ APIKeyUpdateFields) error {
- panic("unexpected Update call")
- }
- func (s *authRepoStub) Delete(ctx context.Context, id int64) error {
- panic("unexpected Delete call")
- }
- func (s *authRepoStub) DeleteWithAudit(ctx context.Context, id int64) error {
- panic("unexpected DeleteWithAudit call")
- }
- func (s *authRepoStub) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
- panic("unexpected ListByUserID call")
- }
- func (s *authRepoStub) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
- panic("unexpected VerifyOwnership call")
- }
- func (s *authRepoStub) CountByUserID(ctx context.Context, userID int64) (int64, error) {
- panic("unexpected CountByUserID call")
- }
- func (s *authRepoStub) ExistsByKey(ctx context.Context, key string) (bool, error) {
- panic("unexpected ExistsByKey call")
- }
- func (s *authRepoStub) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]APIKey, *pagination.PaginationResult, error) {
- panic("unexpected ListByGroupID call")
- }
- func (s *authRepoStub) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error) {
- panic("unexpected SearchAPIKeys call")
- }
- func (s *authRepoStub) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
- panic("unexpected ClearGroupIDByGroupID call")
- }
- func (s *authRepoStub) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
- panic("unexpected UpdateGroupIDByUserAndGroup call")
- }
- func (s *authRepoStub) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
- panic("unexpected CountByGroupID call")
- }
- func (s *authRepoStub) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
- if s.listKeysByUserID == nil {
- panic("unexpected ListKeysByUserID call")
- }
- return s.listKeysByUserID(ctx, userID)
- }
- func (s *authRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
- if s.listKeysByGroupID == nil {
- panic("unexpected ListKeysByGroupID call")
- }
- return s.listKeysByGroupID(ctx, groupID)
- }
- func (s *authRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
- panic("unexpected IncrementQuotaUsed call")
- }
- func (s *authRepoStub) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
- panic("unexpected UpdateLastUsed call")
- }
- func (s *authRepoStub) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
- panic("unexpected IncrementRateLimitUsage call")
- }
- func (s *authRepoStub) ResetRateLimitWindows(ctx context.Context, id int64) error {
- panic("unexpected ResetRateLimitWindows call")
- }
- func (s *authRepoStub) GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error) {
- panic("unexpected GetRateLimitData call")
- }
- type authCacheStub struct {
- getAuthCache func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
- setAuthKeys []string
- deleteAuthKeys []string
- }
- func (s *authCacheStub) GetCreateAttemptCount(ctx context.Context, userID int64) (int, error) {
- return 0, nil
- }
- func (s *authCacheStub) IncrementCreateAttemptCount(ctx context.Context, userID int64) error {
- return nil
- }
- func (s *authCacheStub) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
- return nil
- }
- func (s *authCacheStub) IncrementDailyUsage(ctx context.Context, apiKey string) error {
- return nil
- }
- func (s *authCacheStub) SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error {
- return nil
- }
- func (s *authCacheStub) GetAuthCache(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
- if s.getAuthCache == nil {
- return nil, redis.Nil
- }
- return s.getAuthCache(ctx, key)
- }
- func (s *authCacheStub) SetAuthCache(ctx context.Context, key string, entry *APIKeyAuthCacheEntry, ttl time.Duration) error {
- s.setAuthKeys = append(s.setAuthKeys, key)
- return nil
- }
- func (s *authCacheStub) DeleteAuthCache(ctx context.Context, key string) error {
- s.deleteAuthKeys = append(s.deleteAuthKeys, key)
- return nil
- }
- func (s *authCacheStub) PublishAuthCacheInvalidation(ctx context.Context, cacheKey string) error {
- return nil
- }
- func (s *authCacheStub) SubscribeAuthCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error {
- return nil
- }
- func TestAPIKeyService_GetByKey_UsesL2Cache(t *testing.T) {
- cache := &authCacheStub{}
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- return nil, errors.New("unexpected repo call")
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- NegativeTTLSeconds: 30,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- groupID := int64(9)
- cacheEntry := &APIKeyAuthCacheEntry{
- Snapshot: &APIKeyAuthSnapshot{
- Version: apiKeyAuthSnapshotVersion,
- APIKeyID: 1,
- UserID: 2,
- GroupID: &groupID,
- Status: StatusActive,
- User: APIKeyAuthUserSnapshot{
- ID: 2,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 10,
- Concurrency: 3,
- },
- Group: &APIKeyAuthGroupSnapshot{
- ID: groupID,
- Name: "g",
- Platform: PlatformAnthropic,
- Status: StatusActive,
- SubscriptionType: SubscriptionTypeStandard,
- RateMultiplier: 1,
- ModelRoutingEnabled: true,
- ModelRouting: map[string][]int64{
- "claude-opus-*": {1, 2},
- },
- },
- },
- }
- cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
- return cacheEntry, nil
- }
- apiKey, err := svc.GetByKey(context.Background(), "k1")
- require.NoError(t, err)
- require.Equal(t, int64(1), apiKey.ID)
- require.Equal(t, int64(2), apiKey.User.ID)
- require.Equal(t, groupID, apiKey.Group.ID)
- require.True(t, apiKey.Group.ModelRoutingEnabled)
- require.Equal(t, map[string][]int64{"claude-opus-*": {1, 2}}, apiKey.Group.ModelRouting)
- }
- func TestAPIKeyService_SnapshotRoundTrip_PreservesMessagesDispatchModelConfig(t *testing.T) {
- svc := NewAPIKeyService(nil, nil, nil, nil, nil, nil, &config.Config{})
- groupID := int64(9)
- apiKey := &APIKey{
- ID: 1,
- UserID: 2,
- GroupID: &groupID,
- Key: "k-roundtrip",
- Name: "Audit Key",
- Status: StatusActive,
- User: &User{
- ID: 2,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 10,
- Concurrency: 3,
- },
- Group: &Group{
- ID: groupID,
- Name: "openai",
- Platform: PlatformOpenAI,
- Status: StatusActive,
- SubscriptionType: SubscriptionTypeStandard,
- RateMultiplier: 1,
- AllowMessagesDispatch: true,
- DefaultMappedModel: "gpt-5.4",
- MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
- OpusMappedModel: "gpt-5.4-nano",
- SonnetMappedModel: "gpt-5.3-codex",
- HaikuMappedModel: "gpt-5.4-mini",
- ExactModelMappings: map[string]string{
- "claude-sonnet-4.5": "gpt-5.4-nano",
- },
- },
- },
- }
- snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
- roundTrip := svc.snapshotToAPIKey(apiKey.Key, snapshot)
- require.NotNil(t, roundTrip)
- require.Equal(t, apiKey.Name, roundTrip.Name)
- require.NotNil(t, roundTrip.Group)
- require.Equal(t, apiKey.Group.MessagesDispatchModelConfig, roundTrip.Group.MessagesDispatchModelConfig)
- }
- func TestAPIKeyService_SnapshotRoundTrip_PreservesReasoningEffortPolicy(t *testing.T) {
- svc := NewAPIKeyService(nil, nil, nil, nil, nil, nil, &config.Config{})
- groupID := int64(9)
- apiKey := &APIKey{
- ID: 1,
- UserID: 2,
- GroupID: &groupID,
- Key: "k-reasoning-policy",
- Status: StatusActive,
- User: &User{
- ID: 2,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 10,
- Concurrency: 3,
- },
- Group: &Group{
- ID: groupID,
- Name: "composite",
- Platform: PlatformComposite,
- Status: StatusActive,
- SubscriptionType: SubscriptionTypeStandard,
- RateMultiplier: 1,
- MaxReasoningEffort: "medium",
- MaxReasoningEffortOverLimit: ReasoningEffortOverLimitDeny,
- ReasoningEffortMappings: []ReasoningEffortMapping{
- {From: "max", To: "xhigh"},
- },
- },
- }
- snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
- roundTrip := svc.snapshotToAPIKey(apiKey.Key, snapshot)
- require.NotNil(t, roundTrip)
- require.NotNil(t, roundTrip.Group)
- require.Equal(t, PlatformComposite, roundTrip.Group.Platform)
- require.Equal(t, "medium", roundTrip.Group.MaxReasoningEffort)
- require.Equal(t, ReasoningEffortOverLimitDeny, roundTrip.Group.MaxReasoningEffortOverLimit)
- require.Equal(t, apiKey.Group.ReasoningEffortMappings, roundTrip.Group.ReasoningEffortMappings)
- }
- func TestAPIKeyService_GetByKey_IgnoresLegacyAuthCacheSnapshotWithoutMessagesDispatchConfig(t *testing.T) {
- cache := &authCacheStub{}
- var repoCalls int32
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- atomic.AddInt32(&repoCalls, 1)
- groupID := int64(9)
- return &APIKey{
- ID: 1,
- UserID: 2,
- GroupID: &groupID,
- Status: StatusActive,
- User: &User{
- ID: 2,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 10,
- Concurrency: 3,
- },
- Group: &Group{
- ID: groupID,
- Name: "openai",
- Platform: PlatformOpenAI,
- Status: StatusActive,
- Hydrated: true,
- SubscriptionType: SubscriptionTypeStandard,
- RateMultiplier: 1,
- AllowMessagesDispatch: true,
- DefaultMappedModel: "gpt-5.4",
- MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
- OpusMappedModel: "gpt-5.4-nano",
- },
- },
- }, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- groupID := int64(9)
- cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
- return &APIKeyAuthCacheEntry{
- Snapshot: &APIKeyAuthSnapshot{
- APIKeyID: 1,
- UserID: 2,
- GroupID: &groupID,
- Status: StatusActive,
- User: APIKeyAuthUserSnapshot{
- ID: 2,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 10,
- Concurrency: 3,
- },
- Group: &APIKeyAuthGroupSnapshot{
- ID: groupID,
- Name: "openai",
- Platform: PlatformOpenAI,
- Status: StatusActive,
- SubscriptionType: SubscriptionTypeStandard,
- RateMultiplier: 1,
- AllowMessagesDispatch: true,
- DefaultMappedModel: "gpt-5.4",
- },
- },
- }, nil
- }
- apiKey, err := svc.GetByKey(context.Background(), "k-legacy")
- require.NoError(t, err)
- require.Equal(t, int32(1), atomic.LoadInt32(&repoCalls))
- require.NotNil(t, apiKey.Group)
- require.Equal(t, "gpt-5.4-nano", apiKey.Group.MessagesDispatchModelConfig.OpusMappedModel)
- }
- func TestAPIKeyService_GetByKey_NegativeCache(t *testing.T) {
- cache := &authCacheStub{}
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- return nil, errors.New("unexpected repo call")
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- NegativeTTLSeconds: 30,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
- return &APIKeyAuthCacheEntry{NotFound: true}, nil
- }
- _, err := svc.GetByKey(context.Background(), "missing")
- require.ErrorIs(t, err, ErrAPIKeyNotFound)
- }
- func TestAPIKeyService_GetByKey_CacheMissStoresL2(t *testing.T) {
- cache := &authCacheStub{}
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- return &APIKey{
- ID: 5,
- UserID: 7,
- Status: StatusActive,
- User: &User{
- ID: 7,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 12,
- Concurrency: 2,
- },
- }, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- NegativeTTLSeconds: 30,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
- return nil, redis.Nil
- }
- apiKey, err := svc.GetByKey(context.Background(), "k2")
- require.NoError(t, err)
- require.Equal(t, int64(5), apiKey.ID)
- require.Len(t, cache.setAuthKeys, 1)
- }
- func TestAPIKeyService_GetByKey_UsesL1Cache(t *testing.T) {
- var calls int32
- cache := &authCacheStub{}
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- atomic.AddInt32(&calls, 1)
- return &APIKey{
- ID: 21,
- UserID: 3,
- Status: StatusActive,
- User: &User{
- ID: 3,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 5,
- Concurrency: 2,
- },
- }, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L1Size: 1000,
- L1TTLSeconds: 60,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- require.NotNil(t, svc.authCacheL1)
- _, err := svc.GetByKey(context.Background(), "k-l1")
- require.NoError(t, err)
- svc.authCacheL1.Wait()
- cacheKey := svc.authCacheKey("k-l1")
- _, ok := svc.authCacheL1.Get(cacheKey)
- require.True(t, ok)
- _, err = svc.GetByKey(context.Background(), "k-l1")
- require.NoError(t, err)
- require.Equal(t, int32(1), atomic.LoadInt32(&calls))
- }
- func TestAPIKeyService_InvalidateAuthCacheByUserID(t *testing.T) {
- cache := &authCacheStub{}
- repo := &authRepoStub{
- listKeysByUserID: func(ctx context.Context, userID int64) ([]string, error) {
- return []string{"k1", "k2"}, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- NegativeTTLSeconds: 30,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- svc.InvalidateAuthCacheByUserID(context.Background(), 7)
- require.Len(t, cache.deleteAuthKeys, 2)
- }
- func TestAPIKeyService_InvalidateAuthCacheByGroupID(t *testing.T) {
- cache := &authCacheStub{}
- repo := &authRepoStub{
- listKeysByGroupID: func(ctx context.Context, groupID int64) ([]string, error) {
- return []string{"k1", "k2"}, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- svc.InvalidateAuthCacheByGroupID(context.Background(), 9)
- require.Len(t, cache.deleteAuthKeys, 2)
- }
- func TestAPIKeyService_InvalidateAuthCacheByKey(t *testing.T) {
- cache := &authCacheStub{}
- repo := &authRepoStub{
- listKeysByUserID: func(ctx context.Context, userID int64) ([]string, error) {
- return nil, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L2TTLSeconds: 60,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- svc.InvalidateAuthCacheByKey(context.Background(), "k1")
- require.Len(t, cache.deleteAuthKeys, 1)
- }
- func TestAPIKeyService_GetByKey_CachesNegativeOnRepoMiss(t *testing.T) {
- var repoCalls atomic.Int32
- cache := &authCacheStub{}
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- repoCalls.Add(1)
- return nil, ErrAPIKeyNotFound
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- L1Size: 100,
- L1TTLSeconds: 60,
- L2TTLSeconds: 60,
- NegativeTTLSeconds: 30,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
- return nil, redis.Nil
- }
- _, err := svc.GetByKey(context.Background(), "missing")
- require.ErrorIs(t, err, ErrAPIKeyNotFound)
- require.Empty(t, cache.setAuthKeys, "attacker-controlled misses must not be written to Redis")
- svc.authNegativeCacheL1.Wait()
- _, err = svc.GetByKey(context.Background(), "missing")
- require.ErrorIs(t, err, ErrAPIKeyNotFound)
- require.Equal(t, int32(1), repoCalls.Load())
- }
- func TestAPIKeyService_GetByKeyRejectsInvalidLengthBeforeCaches(t *testing.T) {
- var cacheCalls atomic.Int32
- cache := &authCacheStub{getAuthCache: func(context.Context, string) (*APIKeyAuthCacheEntry, error) {
- cacheCalls.Add(1)
- return nil, redis.Nil
- }}
- repo := &authRepoStub{getByKeyForAuth: func(context.Context, string) (*APIKey, error) {
- t.Fatal("invalid credential reached repository")
- return nil, nil
- }}
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{L2TTLSeconds: 60}})
- for _, key := range []string{"", strings.Repeat("x", MaxAPIKeyCredentialBytes+1)} {
- _, err := svc.GetByKey(context.Background(), key)
- require.ErrorIs(t, err, ErrAPIKeyNotFound)
- }
- require.Zero(t, cacheCalls.Load())
- }
- func TestAPIKeyService_GetByKeyAllowsMaximumLength(t *testing.T) {
- key := strings.Repeat("x", MaxAPIKeyCredentialBytes)
- var repoCalls atomic.Int32
- repo := &authRepoStub{getByKeyForAuth: func(_ context.Context, got string) (*APIKey, error) {
- repoCalls.Add(1)
- require.Equal(t, key, got)
- return nil, ErrAPIKeyNotFound
- }}
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{})
- _, err := svc.GetByKey(context.Background(), key)
- require.ErrorIs(t, err, ErrAPIKeyNotFound)
- require.Equal(t, int32(1), repoCalls.Load())
- }
- func TestAPIKeyService_AuthLookupBulkheadRejectsExcessMisses(t *testing.T) {
- entered := make(chan struct{})
- release := make(chan struct{})
- repo := &authRepoStub{getByKeyForAuth: func(context.Context, string) (*APIKey, error) {
- close(entered)
- <-release
- return nil, ErrAPIKeyNotFound
- }}
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{LookupConcurrency: 1}})
- done := make(chan error, 1)
- go func() {
- _, err := svc.GetByKey(context.Background(), "first")
- done <- err
- }()
- <-entered
- _, err := svc.GetByKey(context.Background(), "second")
- require.ErrorIs(t, err, ErrAPIKeyAuthOverloaded)
- metrics := svc.AuthLookupMetrics()
- require.Equal(t, uint64(2), metrics.Total)
- require.Equal(t, uint64(1), metrics.Rejected)
- require.Equal(t, int64(1), metrics.InFlight)
- require.Equal(t, 1, metrics.Capacity)
- close(release)
- require.ErrorIs(t, <-done, ErrAPIKeyNotFound)
- }
- func TestAPIKeyService_GetByKey_SingleflightCollapses(t *testing.T) {
- var calls int32
- cache := &authCacheStub{}
- repo := &authRepoStub{
- getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
- atomic.AddInt32(&calls, 1)
- time.Sleep(50 * time.Millisecond)
- return &APIKey{
- ID: 11,
- UserID: 2,
- Status: StatusActive,
- User: &User{
- ID: 2,
- Status: StatusActive,
- Role: RoleUser,
- Balance: 1,
- Concurrency: 1,
- },
- }, nil
- },
- }
- cfg := &config.Config{
- APIKeyAuth: config.APIKeyAuthCacheConfig{
- Singleflight: true,
- },
- }
- svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
- start := make(chan struct{})
- wg := sync.WaitGroup{}
- errs := make([]error, 5)
- for i := 0; i < 5; i++ {
- wg.Add(1)
- go func(idx int) {
- defer wg.Done()
- <-start
- _, err := svc.GetByKey(context.Background(), "k1")
- errs[idx] = err
- }(i)
- }
- close(start)
- wg.Wait()
- for _, err := range errs {
- require.NoError(t, err)
- }
- require.Equal(t, int32(1), atomic.LoadInt32(&calls))
- }
|