239 lines
7.1 KiB
Go
239 lines
7.1 KiB
Go
//go:build unit
|
||
|
||
package service
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// resetQuotaUserSubRepoStub 支持 GetByID、ResetUsageWindows，
|
||
// 其余方法继承 userSubRepoNoop（panic）。
|
||
type resetQuotaUserSubRepoStub struct {
|
||
userSubRepoNoop
|
||
|
||
sub *UserSubscription
|
||
|
||
resetDailyCalled bool
|
||
resetWeeklyCalled bool
|
||
resetMonthlyCalled bool
|
||
resetDailyErr error
|
||
resetWeeklyErr error
|
||
resetMonthlyErr 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) ResetUsageWindows(_ context.Context, _ int64, resetDaily, resetWeekly, resetMonthly bool, windowStart time.Time) error {
|
||
r.resetDailyCalled = resetDaily
|
||
r.resetWeeklyCalled = resetWeekly
|
||
r.resetMonthlyCalled = resetMonthly
|
||
if resetDaily && r.resetDailyErr != nil {
|
||
return r.resetDailyErr
|
||
}
|
||
if resetWeekly && r.resetWeeklyErr != nil {
|
||
return r.resetWeeklyErr
|
||
}
|
||
if resetMonthly && r.resetMonthlyErr != nil {
|
||
return r.resetMonthlyErr
|
||
}
|
||
if r.sub == nil {
|
||
return nil
|
||
}
|
||
if resetDaily {
|
||
r.sub.DailyUsageUSD = 0
|
||
r.sub.DailyWindowStart = &windowStart
|
||
}
|
||
if resetWeekly {
|
||
r.sub.WeeklyUsageUSD = 0
|
||
r.sub.WeeklyWindowStart = &windowStart
|
||
}
|
||
if resetMonthly {
|
||
r.sub.MonthlyUsageUSD = 0
|
||
r.sub.MonthlyWindowStart = &windowStart
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, _ *time.Time, 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, _ time.Time) error {
|
||
r.resetWeeklyCalled = true
|
||
return r.resetWeeklyErr
|
||
}
|
||
|
||
func (r *resetQuotaUserSubRepoStub) ResetMonthlyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
|
||
r.resetMonthlyCalled = true
|
||
return r.resetMonthlyErr
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.True(t, stub.resetDailyCalled, "应调用 ResetDailyUsage")
|
||
require.True(t, stub.resetWeeklyCalled, "应调用 ResetWeeklyUsage")
|
||
require.False(t, stub.resetMonthlyCalled, "不应调用 ResetMonthlyUsage")
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.True(t, stub.resetDailyCalled, "应调用 ResetDailyUsage")
|
||
require.False(t, stub.resetWeeklyCalled, "不应调用 ResetWeeklyUsage")
|
||
require.False(t, stub.resetMonthlyCalled, "不应调用 ResetMonthlyUsage")
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.False(t, stub.resetDailyCalled, "不应调用 ResetDailyUsage")
|
||
require.True(t, stub.resetWeeklyCalled, "应调用 ResetWeeklyUsage")
|
||
require.False(t, stub.resetMonthlyCalled, "不应调用 ResetMonthlyUsage")
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.ErrorIs(t, err, ErrInvalidInput)
|
||
require.False(t, stub.resetDailyCalled)
|
||
require.False(t, stub.resetWeeklyCalled)
|
||
require.False(t, stub.resetMonthlyCalled)
|
||
}
|
||
|
||
func TestAdminResetQuota_SubscriptionNotFound(t *testing.T) {
|
||
stub := &resetQuotaUserSubRepoStub{sub: nil}
|
||
svc := newResetQuotaSvc(stub)
|
||
|
||
_, err := svc.AdminResetQuota(context.Background(), 999, true, true, true)
|
||
|
||
require.ErrorIs(t, err, ErrSubscriptionNotFound)
|
||
require.False(t, stub.resetDailyCalled)
|
||
require.False(t, stub.resetWeeklyCalled)
|
||
require.False(t, stub.resetMonthlyCalled)
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.ErrorIs(t, err, dbErr)
|
||
require.True(t, stub.resetDailyCalled)
|
||
require.True(t, stub.resetWeeklyCalled, "原子重置应在一次调用中提交所选窗口")
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.ErrorIs(t, err, dbErr)
|
||
require.True(t, stub.resetWeeklyCalled)
|
||
}
|
||
|
||
func TestAdminResetQuota_ResetMonthlyOnly(t *testing.T) {
|
||
stub := &resetQuotaUserSubRepoStub{
|
||
sub: &UserSubscription{ID: 8, UserID: 10, GroupID: 20},
|
||
}
|
||
svc := newResetQuotaSvc(stub)
|
||
|
||
result, err := svc.AdminResetQuota(context.Background(), 8, false, false, true)
|
||
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.False(t, stub.resetDailyCalled, "不应调用 ResetDailyUsage")
|
||
require.False(t, stub.resetWeeklyCalled, "不应调用 ResetWeeklyUsage")
|
||
require.True(t, stub.resetMonthlyCalled, "应调用 ResetMonthlyUsage")
|
||
}
|
||
|
||
func TestAdminResetQuota_ResetMonthlyUsageError(t *testing.T) {
|
||
dbErr := errors.New("db error")
|
||
stub := &resetQuotaUserSubRepoStub{
|
||
sub: &UserSubscription{ID: 9, UserID: 10, GroupID: 20},
|
||
resetMonthlyErr: dbErr,
|
||
}
|
||
svc := newResetQuotaSvc(stub)
|
||
|
||
_, err := svc.AdminResetQuota(context.Background(), 9, false, false, true)
|
||
|
||
require.ErrorIs(t, err, dbErr)
|
||
require.True(t, stub.resetMonthlyCalled)
|
||
}
|
||
|
||
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, false)
|
||
|
||
require.NoError(t, err)
|
||
// ResetUsageWindows stub 会将 sub.DailyUsageUSD 归零，
|
||
// 服务应返回第二次 GetByID 的刷新值而非初始的 99.9
|
||
require.Equal(t, float64(0), result.DailyUsageUSD, "返回的订阅应反映已归零的用量")
|
||
require.True(t, stub.resetDailyCalled)
|
||
}
|