mirror of
https://gitee.com/wanwujie/sub2api
synced 2026-04-18 13:54:46 +08:00
Merge pull request #909 from StarryKira/feature/admin-reset-subscription-quota
Feature/管理员可以重置账号额度
This commit is contained in:
Wesley Liddick
166
backend/internal/service/subscription_reset_quota_test.go
Normal file
166
backend/internal/service/subscription_reset_quota_test.go
Normal file
@@ -0,0 +1,166 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// resetQuotaUserSubRepoStub 支持 GetByID、ResetDailyUsage、ResetWeeklyUsage，
|
||||
// 其余方法继承 userSubRepoNoop（panic）。
|
||||
type resetQuotaUserSubRepoStub struct {
|
||||
userSubRepoNoop
|
||||
|
||||
sub *UserSubscription
|
||||
|
||||
resetDailyCalled bool
|
||||
resetWeeklyCalled bool
|
||||
resetDailyErr error
|
||||
resetWeeklyErr error
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) GetByID(_ context.Context, id int64) (*UserSubscription, error) {
|
||||
if r.sub == nil || r.sub.ID != id {
|
||||
return nil, ErrSubscriptionNotFound
|
||||
}
|
||||
cp := *r.sub
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, windowStart time.Time) error {
|
||||
r.resetDailyCalled = true
|
||||
if r.resetDailyErr == nil && r.sub != nil {
|
||||
r.sub.DailyUsageUSD = 0
|
||||
r.sub.DailyWindowStart = &windowStart
|
||||
}
|
||||
return r.resetDailyErr
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) ResetWeeklyUsage(_ context.Context, _ int64, _ time.Time) error {
|
||||
r.resetWeeklyCalled = true
|
||||
return r.resetWeeklyErr
|
||||
}
|
||||
|
||||
func newResetQuotaSvc(stub *resetQuotaUserSubRepoStub) *SubscriptionService {
|
||||
return NewSubscriptionService(groupRepoNoop{}, stub, nil, nil, nil)
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ResetBoth(t *testing.T) {
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{ID: 1, UserID: 10, GroupID: 20},
|
||||
}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
result, err := svc.AdminResetQuota(context.Background(), 1, true, true)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, stub.resetDailyCalled, "应调用 ResetDailyUsage")
|
||||
require.True(t, stub.resetWeeklyCalled, "应调用 ResetWeeklyUsage")
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ResetDailyOnly(t *testing.T) {
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{ID: 2, UserID: 10, GroupID: 20},
|
||||
}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
result, err := svc.AdminResetQuota(context.Background(), 2, true, false)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, stub.resetDailyCalled, "应调用 ResetDailyUsage")
|
||||
require.False(t, stub.resetWeeklyCalled, "不应调用 ResetWeeklyUsage")
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ResetWeeklyOnly(t *testing.T) {
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{ID: 3, UserID: 10, GroupID: 20},
|
||||
}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
result, err := svc.AdminResetQuota(context.Background(), 3, false, true)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.False(t, stub.resetDailyCalled, "不应调用 ResetDailyUsage")
|
||||
require.True(t, stub.resetWeeklyCalled, "应调用 ResetWeeklyUsage")
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_BothFalseReturnsError(t *testing.T) {
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{ID: 7, UserID: 10, GroupID: 20},
|
||||
}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
_, err := svc.AdminResetQuota(context.Background(), 7, false, false)
|
||||
|
||||
require.ErrorIs(t, err, ErrInvalidInput)
|
||||
require.False(t, stub.resetDailyCalled)
|
||||
require.False(t, stub.resetWeeklyCalled)
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_SubscriptionNotFound(t *testing.T) {
|
||||
stub := &resetQuotaUserSubRepoStub{sub: nil}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
_, err := svc.AdminResetQuota(context.Background(), 999, true, true)
|
||||
|
||||
require.ErrorIs(t, err, ErrSubscriptionNotFound)
|
||||
require.False(t, stub.resetDailyCalled)
|
||||
require.False(t, stub.resetWeeklyCalled)
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ResetDailyUsageError(t *testing.T) {
|
||||
dbErr := errors.New("db error")
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{ID: 4, UserID: 10, GroupID: 20},
|
||||
resetDailyErr: dbErr,
|
||||
}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
_, err := svc.AdminResetQuota(context.Background(), 4, true, true)
|
||||
|
||||
require.ErrorIs(t, err, dbErr)
|
||||
require.True(t, stub.resetDailyCalled)
|
||||
require.False(t, stub.resetWeeklyCalled, "daily 失败后不应继续调用 weekly")
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ResetWeeklyUsageError(t *testing.T) {
|
||||
dbErr := errors.New("db error")
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{ID: 5, UserID: 10, GroupID: 20},
|
||||
resetWeeklyErr: dbErr,
|
||||
}
|
||||
svc := newResetQuotaSvc(stub)
|
||||
|
||||
_, err := svc.AdminResetQuota(context.Background(), 5, false, true)
|
||||
|
||||
require.ErrorIs(t, err, dbErr)
|
||||
require.True(t, stub.resetWeeklyCalled)
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ReturnsRefreshedSub(t *testing.T) {
|
||||
stub := &resetQuotaUserSubRepoStub{
|
||||
sub: &UserSubscription{
|
||||
ID: 6,
|
||||
UserID: 10,
|
||||
GroupID: 20,
|
||||
DailyUsageUSD: 99.9,
|
||||
},
|
||||
}
|
||||
|
||||
svc := newResetQuotaSvc(stub)
|
||||
result, err := svc.AdminResetQuota(context.Background(), 6, true, false)
|
||||
|
||||
require.NoError(t, err)
|
||||
// ResetDailyUsage stub 会将 sub.DailyUsageUSD 归零，
|
||||
// 服务应返回第二次 GetByID 的刷新值而非初始的 99.9
|
||||
require.Equal(t, float64(0), result.DailyUsageUSD, "返回的订阅应反映已归零的用量")
|
||||
require.True(t, stub.resetDailyCalled)
|
||||
}
|
||||
@@ -31,6 +31,7 @@ var (
|
||||
ErrSubscriptionAlreadyExists = infraerrors.Conflict("SUBSCRIPTION_ALREADY_EXISTS", "subscription already exists for this user and group")
|
||||
ErrSubscriptionAssignConflict = infraerrors.Conflict("SUBSCRIPTION_ASSIGN_CONFLICT", "subscription exists but request conflicts with existing assignment semantics")
|
||||
ErrGroupNotSubscriptionType = infraerrors.BadRequest("GROUP_NOT_SUBSCRIPTION_TYPE", "group is not a subscription type")
|
||||
ErrInvalidInput = infraerrors.BadRequest("INVALID_INPUT", "at least one of resetDaily or resetWeekly must be true")
|
||||
ErrDailyLimitExceeded = infraerrors.TooManyRequests("DAILY_LIMIT_EXCEEDED", "daily usage limit exceeded")
|
||||
ErrWeeklyLimitExceeded = infraerrors.TooManyRequests("WEEKLY_LIMIT_EXCEEDED", "weekly usage limit exceeded")
|
||||
ErrMonthlyLimitExceeded = infraerrors.TooManyRequests("MONTHLY_LIMIT_EXCEEDED", "monthly usage limit exceeded")
|
||||
@@ -695,6 +696,36 @@ func (s *SubscriptionService) CheckAndActivateWindow(ctx context.Context, sub *U
|
||||
return s.userSubRepo.ActivateWindows(ctx, sub.ID, windowStart)
|
||||
}
|
||||
|
||||
// AdminResetQuota manually resets the daily and/or weekly usage windows.
|
||||
// Uses startOfDay(now) as the new window start, matching automatic resets.
|
||||
func (s *SubscriptionService) AdminResetQuota(ctx context.Context, subscriptionID int64, resetDaily, resetWeekly bool) (*UserSubscription, error) {
|
||||
if !resetDaily && !resetWeekly {
|
||||
return nil, ErrInvalidInput
|
||||
}
|
||||
sub, err := s.userSubRepo.GetByID(ctx, subscriptionID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
windowStart := startOfDay(time.Now())
|
||||
if resetDaily {
|
||||
if err := s.userSubRepo.ResetDailyUsage(ctx, sub.ID, windowStart); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if resetWeekly {
|
||||
if err := s.userSubRepo.ResetWeeklyUsage(ctx, sub.ID, windowStart); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
// Invalidate caches, same as CheckAndResetWindows
|
||||
s.InvalidateSubCache(sub.UserID, sub.GroupID)
|
||||
if s.billingCacheService != nil {
|
||||
_ = s.billingCacheService.InvalidateSubscription(ctx, sub.UserID, sub.GroupID)
|
||||
}
|
||||
// Return the refreshed subscription from DB
|
||||
return s.userSubRepo.GetByID(ctx, subscriptionID)
|
||||
}
|
||||
|
||||
// CheckAndResetWindows 检查并重置过期的窗口
|
||||
func (s *SubscriptionService) CheckAndResetWindows(ctx context.Context, sub *UserSubscription) error {
|
||||
// 使用当天零点作为新窗口起始时间
|
||||
|
||||
Reference in New Issue
Block a user