|
|
|
|
@@ -23,6 +23,15 @@ type fakeAPIKeyRepo struct {
|
|
|
|
|
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
type fakeGoogleSubscriptionRepo struct {
|
|
|
|
|
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
|
|
|
|
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
|
|
|
|
activateWindow func(ctx context.Context, id int64, start time.Time) error
|
|
|
|
|
resetDaily func(ctx context.Context, id int64, start time.Time) error
|
|
|
|
|
resetWeekly func(ctx context.Context, id int64, start time.Time) error
|
|
|
|
|
resetMonthly func(ctx context.Context, id int64, start time.Time) error
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (f fakeAPIKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
@@ -87,6 +96,85 @@ func (f fakeAPIKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt tim
|
|
|
|
|
return nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
|
|
|
|
return nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
|
|
|
|
return nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
|
|
|
|
if f.getActive != nil {
|
|
|
|
|
return f.getActive(ctx, userID, groupID)
|
|
|
|
|
}
|
|
|
|
|
return nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) Delete(ctx context.Context, id int64) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
|
|
|
|
return nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
|
|
|
|
return nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
|
|
|
|
return nil, nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
|
|
|
|
return nil, nil, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
|
|
|
|
return false, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
|
|
|
|
|
if f.updateStatus != nil {
|
|
|
|
|
return f.updateStatus(ctx, subscriptionID, status)
|
|
|
|
|
}
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
|
|
|
|
if f.activateWindow != nil {
|
|
|
|
|
return f.activateWindow(ctx, id, start)
|
|
|
|
|
}
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, start time.Time) error {
|
|
|
|
|
if f.resetDaily != nil {
|
|
|
|
|
return f.resetDaily(ctx, id, start)
|
|
|
|
|
}
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, start time.Time) error {
|
|
|
|
|
if f.resetWeekly != nil {
|
|
|
|
|
return f.resetWeekly(ctx, id, start)
|
|
|
|
|
}
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, start time.Time) error {
|
|
|
|
|
if f.resetMonthly != nil {
|
|
|
|
|
return f.resetMonthly(ctx, id, start)
|
|
|
|
|
}
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
|
|
|
|
|
return errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
func (f fakeGoogleSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
|
|
|
|
|
return 0, errors.New("not implemented")
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
type googleErrorResponse struct {
|
|
|
|
|
Error struct {
|
|
|
|
|
Code int `json:"code"`
|
|
|
|
|
@@ -505,3 +593,85 @@ func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedInStandardMode(t *testi
|
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
|
require.Equal(t, 1, touchCalls)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_SubscriptionLimitExceededReturns429(t *testing.T) {
|
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
|
|
|
|
|
|
limit := 1.0
|
|
|
|
|
group := &service.Group{
|
|
|
|
|
ID: 77,
|
|
|
|
|
Name: "gemini-sub",
|
|
|
|
|
Status: service.StatusActive,
|
|
|
|
|
Platform: service.PlatformGemini,
|
|
|
|
|
Hydrated: true,
|
|
|
|
|
SubscriptionType: service.SubscriptionTypeSubscription,
|
|
|
|
|
DailyLimitUSD: &limit,
|
|
|
|
|
}
|
|
|
|
|
user := &service.User{
|
|
|
|
|
ID: 999,
|
|
|
|
|
Role: service.RoleUser,
|
|
|
|
|
Status: service.StatusActive,
|
|
|
|
|
Balance: 10,
|
|
|
|
|
Concurrency: 3,
|
|
|
|
|
}
|
|
|
|
|
apiKey := &service.APIKey{
|
|
|
|
|
ID: 501,
|
|
|
|
|
UserID: user.ID,
|
|
|
|
|
Key: "google-sub-limit",
|
|
|
|
|
Status: service.StatusActive,
|
|
|
|
|
User: user,
|
|
|
|
|
Group: group,
|
|
|
|
|
}
|
|
|
|
|
apiKey.GroupID = &group.ID
|
|
|
|
|
|
|
|
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
|
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
|
|
|
if key != apiKey.Key {
|
|
|
|
|
return nil, service.ErrAPIKeyNotFound
|
|
|
|
|
}
|
|
|
|
|
clone := *apiKey
|
|
|
|
|
return &clone, nil
|
|
|
|
|
},
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
now := time.Now()
|
|
|
|
|
sub := &service.UserSubscription{
|
|
|
|
|
ID: 601,
|
|
|
|
|
UserID: user.ID,
|
|
|
|
|
GroupID: group.ID,
|
|
|
|
|
Status: service.SubscriptionStatusActive,
|
|
|
|
|
ExpiresAt: now.Add(24 * time.Hour),
|
|
|
|
|
DailyWindowStart: &now,
|
|
|
|
|
DailyUsageUSD: 10,
|
|
|
|
|
}
|
|
|
|
|
subscriptionService := service.NewSubscriptionService(nil, fakeGoogleSubscriptionRepo{
|
|
|
|
|
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
|
|
|
|
if userID != user.ID || groupID != group.ID {
|
|
|
|
|
return nil, service.ErrSubscriptionNotFound
|
|
|
|
|
}
|
|
|
|
|
clone := *sub
|
|
|
|
|
return &clone, nil
|
|
|
|
|
},
|
|
|
|
|
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
|
|
|
|
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
|
|
|
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
|
|
|
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
|
|
|
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
|
|
|
}, nil, nil, &config.Config{RunMode: config.RunModeStandard})
|
|
|
|
|
|
|
|
|
|
r := gin.New()
|
|
|
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, &config.Config{RunMode: config.RunModeStandard}))
|
|
|
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
|
|
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
|
|
|
req.Header.Set("x-goog-api-key", apiKey.Key)
|
|
|
|
|
rec := httptest.NewRecorder()
|
|
|
|
|
r.ServeHTTP(rec, req)
|
|
|
|
|
|
|
|
|
|
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
|
|
|
|
var resp googleErrorResponse
|
|
|
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
|
|
|
require.Equal(t, http.StatusTooManyRequests, resp.Error.Code)
|
|
|
|
|
require.Equal(t, "RESOURCE_EXHAUSTED", resp.Error.Status)
|
|
|
|
|
require.Contains(t, resp.Error.Message, "daily usage limit exceeded")
|
|
|
|
|
}
|
|
|
|
|
|