| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553 |
- //go:build unit
- package service
- import (
- "context"
- "errors"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
- "github.com/stretchr/testify/require"
- )
- type redeemCodeRepoStub struct {
- codesByCode map[string]*RedeemCode
- useCalls []struct {
- id int64
- userID int64
- }
- updateCalls []*RedeemCode
- }
- func (s *redeemCodeRepoStub) Create(context.Context, *RedeemCode) error {
- panic("unexpected Create call")
- }
- func (s *redeemCodeRepoStub) CreateBatch(context.Context, []RedeemCode) error {
- panic("unexpected CreateBatch call")
- }
- func (s *redeemCodeRepoStub) GetByID(context.Context, int64) (*RedeemCode, error) {
- panic("unexpected GetByID call")
- }
- func (s *redeemCodeRepoStub) GetByCode(_ context.Context, code string) (*RedeemCode, error) {
- if s.codesByCode == nil {
- return nil, ErrRedeemCodeNotFound
- }
- redeemCode, ok := s.codesByCode[code]
- if !ok {
- return nil, ErrRedeemCodeNotFound
- }
- cloned := *redeemCode
- return &cloned, nil
- }
- func (s *redeemCodeRepoStub) Update(_ context.Context, code *RedeemCode) error {
- if code == nil {
- return nil
- }
- cloned := *code
- s.updateCalls = append(s.updateCalls, &cloned)
- if s.codesByCode == nil {
- s.codesByCode = make(map[string]*RedeemCode)
- }
- s.codesByCode[cloned.Code] = &cloned
- return nil
- }
- func (s *redeemCodeRepoStub) BatchUpdate(context.Context, []int64, RedeemCodeBatchUpdateFields) (int64, error) {
- panic("unexpected BatchUpdate call")
- }
- func (s *redeemCodeRepoStub) Delete(context.Context, int64) error {
- panic("unexpected Delete call")
- }
- func (s *redeemCodeRepoStub) Use(_ context.Context, id, userID int64) error {
- for code, redeemCode := range s.codesByCode {
- if redeemCode.ID != id {
- continue
- }
- now := time.Now().UTC()
- redeemCode.Status = StatusUsed
- redeemCode.UsedBy = &userID
- redeemCode.UsedAt = &now
- s.codesByCode[code] = redeemCode
- s.useCalls = append(s.useCalls, struct {
- id int64
- userID int64
- }{id: id, userID: userID})
- return nil
- }
- return ErrRedeemCodeNotFound
- }
- func (s *redeemCodeRepoStub) List(context.Context, pagination.PaginationParams) ([]RedeemCode, *pagination.PaginationResult, error) {
- panic("unexpected List call")
- }
- func (s *redeemCodeRepoStub) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string) ([]RedeemCode, *pagination.PaginationResult, error) {
- panic("unexpected ListWithFilters call")
- }
- func (s *redeemCodeRepoStub) ListByUser(context.Context, int64, int) ([]RedeemCode, error) {
- panic("unexpected ListByUser call")
- }
- func (s *redeemCodeRepoStub) ListByUserPaginated(context.Context, int64, pagination.PaginationParams, string) ([]RedeemCode, *pagination.PaginationResult, error) {
- panic("unexpected ListByUserPaginated call")
- }
- func (s *redeemCodeRepoStub) SumPositiveBalanceByUser(context.Context, int64) (float64, error) {
- panic("unexpected SumPositiveBalanceByUser call")
- }
- func newOAuthEmailFlowAuthService(
- userRepo UserRepository,
- redeemRepo RedeemCodeRepository,
- refreshTokenCache RefreshTokenCache,
- settings map[string]string,
- emailCache EmailCache,
- quotaRepo UserPlatformQuotaRepository, // 新增
- ) *AuthService {
- cfg := &config.Config{
- JWT: config.JWTConfig{
- Secret: "test-secret",
- ExpireHour: 1,
- AccessTokenExpireMinutes: 60,
- RefreshTokenExpireDays: 7,
- },
- Default: config.DefaultConfig{
- UserBalance: 3.5,
- UserConcurrency: 2,
- },
- }
- settingService := NewSettingService(&settingRepoStub{values: settings}, cfg)
- emailService := NewEmailService(&settingRepoStub{values: settings}, emailCache)
- return NewAuthService(
- nil,
- userRepo,
- redeemRepo,
- refreshTokenCache,
- cfg,
- settingService,
- emailService,
- nil,
- nil,
- nil,
- nil,
- nil,
- quotaRepo, // 替换原来的 nil
- )
- }
- func TestRegisterOAuthEmailAccountRollsBackCreatedUserWhenTokenPairGenerationFails(t *testing.T) {
- userRepo := &userRepoStub{nextID: 42}
- redeemRepo := &redeemCodeRepoStub{
- codesByCode: map[string]*RedeemCode{
- "INVITE123": {
- ID: 7,
- Code: "INVITE123",
- Type: RedeemTypeInvitation,
- Status: StatusUnused,
- },
- },
- }
- emailCache := &emailCacheStub{
- data: &VerificationCodeData{
- Code: "246810",
- Attempts: 0,
- CreatedAt: time.Now().UTC(),
- ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
- },
- }
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- redeemRepo,
- nil,
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyInvitationCodeEnabled: "true",
- SettingKeyEmailVerifyEnabled: "true",
- },
- emailCache,
- nil,
- )
- tokenPair, user, err := authService.RegisterOAuthEmailAccount(
- context.Background(),
- "fresh@example.com",
- "secret-123",
- "246810",
- "INVITE123",
- "oidc",
- )
- require.Nil(t, tokenPair)
- require.Nil(t, user)
- require.Error(t, err)
- require.Contains(t, err.Error(), "generate token pair")
- require.Equal(t, []int64{42}, userRepo.deletedIDs)
- require.Len(t, userRepo.created, 1)
- require.Empty(t, redeemRepo.useCalls)
- require.Empty(t, redeemRepo.updateCalls)
- }
- func TestRegisterOAuthEmailAccount_NonWhitelistDomainLimit(t *testing.T) {
- userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 1}}
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- &redeemCodeRepoStub{},
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
- SettingKeyRegistrationEmailDomainQuotaEnabled: "true",
- },
- &emailCacheStub{data: &VerificationCodeData{
- Code: "246810",
- CreatedAt: time.Now().UTC(),
- ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
- }},
- nil,
- )
- _, _, err := authService.RegisterOAuthEmailAccount(
- context.Background(),
- "second@custom.example",
- "secret-123",
- "246810",
- "",
- "oidc",
- )
- require.ErrorIs(t, err, ErrEmailDomainRegistrationLimit)
- }
- func TestRegisterVerifiedOAuthEmailAccount_NonWhitelistDomainLimit(t *testing.T) {
- userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 1}}
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- nil,
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
- SettingKeyRegistrationEmailDomainQuotaEnabled: "true",
- },
- &emailCacheStub{},
- nil,
- )
- _, _, err := authService.RegisterVerifiedOAuthEmailAccount(
- context.Background(),
- "second@custom.example",
- "secret-123",
- "",
- "oidc",
- )
- require.ErrorIs(t, err, ErrEmailDomainRegistrationLimit)
- }
- func TestSendPendingOAuthVerifyCode_NonWhitelistDomainLimit(t *testing.T) {
- userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 1}}
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- nil,
- nil,
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
- SettingKeyRegistrationEmailDomainQuotaEnabled: "true",
- },
- &emailCacheStub{},
- nil,
- )
- _, err := authService.SendPendingOAuthVerifyCode(context.Background(), "second@custom.example")
- require.ErrorIs(t, err, ErrEmailDomainRegistrationLimit)
- }
- // 域名限量注册开关默认关闭:白名单外域名在 pending OAuth 发码阶段即被严格拒绝。
- func TestSendPendingOAuthVerifyCode_NonWhitelistDomainRejectedWhenQuotaDisabled(t *testing.T) {
- userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 0}}
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- nil,
- nil,
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
- },
- &emailCacheStub{},
- nil,
- )
- _, err := authService.SendPendingOAuthVerifyCode(context.Background(), "first@custom.example")
- require.ErrorIs(t, err, ErrEmailSuffixNotAllowed)
- }
- func TestSendPendingOAuthVerifyCode_NilServiceReturnsUnavailable(t *testing.T) {
- var authService *AuthService
- _, err := authService.SendPendingOAuthVerifyCode(context.Background(), "fresh@example.com")
- require.ErrorIs(t, err, ErrServiceUnavailable)
- }
- func TestRegisterOAuthEmailAccountSetsNormalizedSignupSourceOnCreatedUser(t *testing.T) {
- userRepo := &userRepoStub{nextID: 42}
- emailCache := &emailCacheStub{
- data: &VerificationCodeData{
- Code: "246810",
- Attempts: 0,
- CreatedAt: time.Now().UTC(),
- ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
- },
- }
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- &redeemCodeRepoStub{},
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyEmailVerifyEnabled: "true",
- },
- emailCache,
- nil,
- )
- tokenPair, user, err := authService.RegisterOAuthEmailAccount(
- context.Background(),
- "fresh@example.com",
- "secret-123",
- "246810",
- "",
- " OIDC ",
- )
- require.NoError(t, err)
- require.NotNil(t, tokenPair)
- require.NotNil(t, user)
- require.Len(t, userRepo.created, 1)
- require.Equal(t, "oidc", userRepo.created[0].SignupSource)
- }
- func TestRegisterOAuthEmailAccountKeepsGitHubAndGoogleSignupSource(t *testing.T) {
- tests := []struct {
- name string
- email string
- signupSource string
- want string
- }{
- {
- name: "github",
- email: "github@example.com",
- signupSource: " GitHub ",
- want: "github",
- },
- {
- name: "google",
- email: "google@example.com",
- signupSource: " Google ",
- want: "google",
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- userRepo := &userRepoStub{nextID: 43}
- emailCache := &emailCacheStub{
- data: &VerificationCodeData{
- Code: "246810",
- Attempts: 0,
- CreatedAt: time.Now().UTC(),
- ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
- },
- }
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- &redeemCodeRepoStub{},
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyEmailVerifyEnabled: "true",
- },
- emailCache,
- nil,
- )
- tokenPair, user, err := authService.RegisterOAuthEmailAccount(
- context.Background(),
- tt.email,
- "secret-123",
- "246810",
- "",
- tt.signupSource,
- )
- require.NoError(t, err)
- require.NotNil(t, tokenPair)
- require.NotNil(t, user)
- require.Len(t, userRepo.created, 1)
- require.Equal(t, tt.want, userRepo.created[0].SignupSource)
- })
- }
- }
- func TestRegisterOAuthEmailAccountFallsBackUnknownSignupSourceToEmail(t *testing.T) {
- userRepo := &userRepoStub{nextID: 43}
- emailCache := &emailCacheStub{
- data: &VerificationCodeData{
- Code: "246810",
- Attempts: 0,
- CreatedAt: time.Now().UTC(),
- ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
- },
- }
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- &redeemCodeRepoStub{},
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyEmailVerifyEnabled: "true",
- },
- emailCache,
- nil,
- )
- tokenPair, user, err := authService.RegisterOAuthEmailAccount(
- context.Background(),
- "fallback@example.com",
- "secret-123",
- "246810",
- "",
- "unknown-provider",
- )
- require.NoError(t, err)
- require.NotNil(t, tokenPair)
- require.NotNil(t, user)
- require.Len(t, userRepo.created, 1)
- require.Equal(t, "email", userRepo.created[0].SignupSource)
- }
- func TestRollbackOAuthEmailAccountCreationRestoresInvitationUsage(t *testing.T) {
- userRepo := &userRepoStub{}
- redeemRepo := &redeemCodeRepoStub{
- codesByCode: map[string]*RedeemCode{
- "INVITE123": {
- ID: 7,
- Code: "INVITE123",
- Type: RedeemTypeInvitation,
- Status: StatusUsed,
- UsedBy: func() *int64 {
- v := int64(42)
- return &v
- }(),
- UsedAt: func() *time.Time {
- v := time.Now().UTC()
- return &v
- }(),
- },
- },
- }
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- redeemRepo,
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyInvitationCodeEnabled: "true",
- },
- &emailCacheStub{},
- nil,
- )
- err := authService.RollbackOAuthEmailAccountCreation(context.Background(), 42, "INVITE123")
- require.NoError(t, err)
- require.Equal(t, []int64{42}, userRepo.deletedIDs)
- require.Len(t, redeemRepo.updateCalls, 1)
- require.Equal(t, StatusUnused, redeemRepo.updateCalls[0].Status)
- require.Nil(t, redeemRepo.updateCalls[0].UsedBy)
- require.Nil(t, redeemRepo.updateCalls[0].UsedAt)
- }
- func TestRollbackOAuthEmailAccountCreationPropagatesDeleteError(t *testing.T) {
- userRepo := &userRepoStub{deleteErr: errors.New("delete failed")}
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- &redeemCodeRepoStub{},
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- },
- &emailCacheStub{},
- nil,
- )
- err := authService.RollbackOAuthEmailAccountCreation(context.Background(), 42, "")
- require.Error(t, err)
- require.Contains(t, err.Error(), "delete created oauth user")
- }
- func TestFinalizeOAuthEmailAccount_SnapshotsPlatformQuotaDefaults(t *testing.T) {
- userRepo := &userRepoStub{nextID: 99}
- quotaRepo := &userPlatformQuotaRepoStub{}
- authService := newOAuthEmailFlowAuthService(
- userRepo,
- nil,
- &refreshTokenCacheStub{},
- map[string]string{
- SettingKeyRegistrationEnabled: "true",
- SettingKeyEmailVerifyEnabled: "true",
- SettingKeyDefaultPlatformQuotas: `{"anthropic": {"daily": 5.5}}`,
- },
- &emailCacheStub{},
- quotaRepo,
- )
- user := &User{
- ID: 99,
- Email: "newuser@example.com",
- Role: RoleUser,
- Status: StatusActive,
- SignupSource: "oidc",
- }
- err := authService.FinalizeOAuthEmailAccount(
- context.Background(),
- user,
- "",
- "oidc",
- "",
- )
- require.NoError(t, err)
- require.Len(t, quotaRepo.bulkInsertCalls, 1, "snapshotPlatformQuotaDefaults must call BulkInsertInitial once on successful OAuth signup")
- records := quotaRepo.bulkInsertCalls[0]
- var anthropicRecord *UserPlatformQuotaRecord
- for i := range records {
- if records[i].Platform == "anthropic" {
- anthropicRecord = &records[i]
- break
- }
- }
- require.NotNil(t, anthropicRecord, "expected anthropic platform record")
- require.Equal(t, int64(99), anthropicRecord.UserID)
- require.NotNil(t, anthropicRecord.DailyLimitUSD)
- require.InDelta(t, 5.5, *anthropicRecord.DailyLimitUSD, 0.0001)
- }
|