feat(usage): 添加清理任务与统计过滤
This commit is contained in:
yangjianbo
420
backend/internal/service/usage_cleanup_service_test.go
Normal file
420
backend/internal/service/usage_cleanup_service_test.go
Normal file
@@ -0,0 +1,420 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type cleanupDeleteResponse struct {
|
||||
deleted int64
|
||||
err error
|
||||
}
|
||||
|
||||
type cleanupDeleteCall struct {
|
||||
filters UsageCleanupFilters
|
||||
limit int
|
||||
}
|
||||
|
||||
type cleanupMarkCall struct {
|
||||
taskID int64
|
||||
deletedRows int64
|
||||
errMsg string
|
||||
}
|
||||
|
||||
type cleanupRepoStub struct {
|
||||
mu sync.Mutex
|
||||
created []*UsageCleanupTask
|
||||
createErr error
|
||||
listTasks []UsageCleanupTask
|
||||
listResult *pagination.PaginationResult
|
||||
listErr error
|
||||
claimQueue []*UsageCleanupTask
|
||||
claimErr error
|
||||
deleteQueue []cleanupDeleteResponse
|
||||
deleteCalls []cleanupDeleteCall
|
||||
markSucceeded []cleanupMarkCall
|
||||
markFailed []cleanupMarkCall
|
||||
statusByID map[int64]string
|
||||
progressCalls []cleanupMarkCall
|
||||
cancelCalls []int64
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) CreateTask(ctx context.Context, task *UsageCleanupTask) error {
|
||||
if task == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.createErr != nil {
|
||||
return s.createErr
|
||||
}
|
||||
if task.ID == 0 {
|
||||
task.ID = int64(len(s.created) + 1)
|
||||
}
|
||||
if task.CreatedAt.IsZero() {
|
||||
task.CreatedAt = time.Now().UTC()
|
||||
}
|
||||
if task.UpdatedAt.IsZero() {
|
||||
task.UpdatedAt = task.CreatedAt
|
||||
}
|
||||
clone := *task
|
||||
s.created = append(s.created, &clone)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) ListTasks(ctx context.Context, params pagination.PaginationParams) ([]UsageCleanupTask, *pagination.PaginationResult, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.listTasks, s.listResult, s.listErr
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) ClaimNextPendingTask(ctx context.Context, staleRunningAfterSeconds int64) (*UsageCleanupTask, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.claimErr != nil {
|
||||
return nil, s.claimErr
|
||||
}
|
||||
if len(s.claimQueue) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
task := s.claimQueue[0]
|
||||
s.claimQueue = s.claimQueue[1:]
|
||||
if s.statusByID == nil {
|
||||
s.statusByID = map[int64]string{}
|
||||
}
|
||||
s.statusByID[task.ID] = UsageCleanupStatusRunning
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) GetTaskStatus(ctx context.Context, taskID int64) (string, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.statusByID == nil {
|
||||
return "", sql.ErrNoRows
|
||||
}
|
||||
status, ok := s.statusByID[taskID]
|
||||
if !ok {
|
||||
return "", sql.ErrNoRows
|
||||
}
|
||||
return status, nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) UpdateTaskProgress(ctx context.Context, taskID int64, deletedRows int64) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.progressCalls = append(s.progressCalls, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) CancelTask(ctx context.Context, taskID int64, canceledBy int64) (bool, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.cancelCalls = append(s.cancelCalls, taskID)
|
||||
if s.statusByID == nil {
|
||||
s.statusByID = map[int64]string{}
|
||||
}
|
||||
status := s.statusByID[taskID]
|
||||
if status != UsageCleanupStatusPending && status != UsageCleanupStatusRunning {
|
||||
return false, nil
|
||||
}
|
||||
s.statusByID[taskID] = UsageCleanupStatusCanceled
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) MarkTaskSucceeded(ctx context.Context, taskID int64, deletedRows int64) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.markSucceeded = append(s.markSucceeded, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows})
|
||||
if s.statusByID == nil {
|
||||
s.statusByID = map[int64]string{}
|
||||
}
|
||||
s.statusByID[taskID] = UsageCleanupStatusSucceeded
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) MarkTaskFailed(ctx context.Context, taskID int64, deletedRows int64, errorMsg string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.markFailed = append(s.markFailed, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows, errMsg: errorMsg})
|
||||
if s.statusByID == nil {
|
||||
s.statusByID = map[int64]string{}
|
||||
}
|
||||
s.statusByID[taskID] = UsageCleanupStatusFailed
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cleanupRepoStub) DeleteUsageLogsBatch(ctx context.Context, filters UsageCleanupFilters, limit int) (int64, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.deleteCalls = append(s.deleteCalls, cleanupDeleteCall{filters: filters, limit: limit})
|
||||
if len(s.deleteQueue) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
resp := s.deleteQueue[0]
|
||||
s.deleteQueue = s.deleteQueue[1:]
|
||||
return resp.deleted, resp.err
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceCreateTaskSanitizeFilters(t *testing.T) {
|
||||
repo := &cleanupRepoStub{}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
end := start.Add(24 * time.Hour)
|
||||
userID := int64(-1)
|
||||
apiKeyID := int64(10)
|
||||
model := " gpt-4 "
|
||||
billingType := int8(-2)
|
||||
filters := UsageCleanupFilters{
|
||||
StartTime: start,
|
||||
EndTime: end,
|
||||
UserID: &userID,
|
||||
APIKeyID: &apiKeyID,
|
||||
Model: &model,
|
||||
BillingType: &billingType,
|
||||
}
|
||||
|
||||
task, err := svc.CreateTask(context.Background(), filters, 9)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, UsageCleanupStatusPending, task.Status)
|
||||
require.Nil(t, task.Filters.UserID)
|
||||
require.NotNil(t, task.Filters.APIKeyID)
|
||||
require.Equal(t, apiKeyID, *task.Filters.APIKeyID)
|
||||
require.NotNil(t, task.Filters.Model)
|
||||
require.Equal(t, "gpt-4", *task.Filters.Model)
|
||||
require.Nil(t, task.Filters.BillingType)
|
||||
require.Equal(t, int64(9), task.CreatedBy)
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceCreateTaskInvalidCreator(t *testing.T) {
|
||||
repo := &cleanupRepoStub{}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
filters := UsageCleanupFilters{
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now().Add(24 * time.Hour),
|
||||
}
|
||||
_, err := svc.CreateTask(context.Background(), filters, 0)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "USAGE_CLEANUP_INVALID_CREATOR", infraerrors.Reason(err))
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceCreateTaskDisabled(t *testing.T) {
|
||||
repo := &cleanupRepoStub{}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: false}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
filters := UsageCleanupFilters{
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now().Add(24 * time.Hour),
|
||||
}
|
||||
_, err := svc.CreateTask(context.Background(), filters, 1)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, http.StatusServiceUnavailable, infraerrors.Code(err))
|
||||
require.Equal(t, "USAGE_CLEANUP_DISABLED", infraerrors.Reason(err))
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceCreateTaskRangeTooLarge(t *testing.T) {
|
||||
repo := &cleanupRepoStub{}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 1}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
end := start.Add(48 * time.Hour)
|
||||
filters := UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
_, err := svc.CreateTask(context.Background(), filters, 1)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "USAGE_CLEANUP_RANGE_TOO_LARGE", infraerrors.Reason(err))
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceCreateTaskMissingRange(t *testing.T) {
|
||||
repo := &cleanupRepoStub{}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
_, err := svc.CreateTask(context.Background(), UsageCleanupFilters{}, 1)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "USAGE_CLEANUP_MISSING_RANGE", infraerrors.Reason(err))
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceCreateTaskRepoError(t *testing.T) {
|
||||
repo := &cleanupRepoStub{createErr: errors.New("db down")}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
filters := UsageCleanupFilters{
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now().Add(24 * time.Hour),
|
||||
}
|
||||
_, err := svc.CreateTask(context.Background(), filters, 1)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "create cleanup task")
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceRunOnceSuccess(t *testing.T) {
|
||||
repo := &cleanupRepoStub{
|
||||
claimQueue: []*UsageCleanupTask{
|
||||
{ID: 5, Filters: UsageCleanupFilters{StartTime: time.Now(), EndTime: time.Now().Add(2 * time.Hour)}},
|
||||
},
|
||||
deleteQueue: []cleanupDeleteResponse{
|
||||
{deleted: 2},
|
||||
{deleted: 2},
|
||||
{deleted: 1},
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2, TaskTimeoutSeconds: 30}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
|
||||
svc.runOnce()
|
||||
|
||||
repo.mu.Lock()
|
||||
defer repo.mu.Unlock()
|
||||
require.Len(t, repo.deleteCalls, 3)
|
||||
require.Len(t, repo.markSucceeded, 1)
|
||||
require.Empty(t, repo.markFailed)
|
||||
require.Equal(t, int64(5), repo.markSucceeded[0].taskID)
|
||||
require.Equal(t, int64(5), repo.markSucceeded[0].deletedRows)
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceRunOnceClaimError(t *testing.T) {
|
||||
repo := &cleanupRepoStub{claimErr: errors.New("claim failed")}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
svc.runOnce()
|
||||
|
||||
repo.mu.Lock()
|
||||
defer repo.mu.Unlock()
|
||||
require.Empty(t, repo.markSucceeded)
|
||||
require.Empty(t, repo.markFailed)
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceRunOnceAlreadyRunning(t *testing.T) {
|
||||
repo := &cleanupRepoStub{}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
svc.running = 1
|
||||
svc.runOnce()
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceExecuteTaskFailed(t *testing.T) {
|
||||
longMsg := strings.Repeat("x", 600)
|
||||
repo := &cleanupRepoStub{
|
||||
deleteQueue: []cleanupDeleteResponse{
|
||||
{err: errors.New(longMsg)},
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 3}}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, cfg)
|
||||
task := &UsageCleanupTask{
|
||||
ID: 11,
|
||||
Filters: UsageCleanupFilters{
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now().Add(24 * time.Hour),
|
||||
},
|
||||
}
|
||||
|
||||
svc.executeTask(context.Background(), task)
|
||||
|
||||
repo.mu.Lock()
|
||||
defer repo.mu.Unlock()
|
||||
require.Len(t, repo.markFailed, 1)
|
||||
require.Equal(t, int64(11), repo.markFailed[0].taskID)
|
||||
require.Equal(t, 500, len(repo.markFailed[0].errMsg))
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceListTasks(t *testing.T) {
|
||||
repo := &cleanupRepoStub{
|
||||
listTasks: []UsageCleanupTask{{ID: 1}, {ID: 2}},
|
||||
listResult: &pagination.PaginationResult{
|
||||
Total: 2,
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
Pages: 1,
|
||||
},
|
||||
}
|
||||
svc := NewUsageCleanupService(repo, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
|
||||
|
||||
tasks, result, err := svc.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, tasks, 2)
|
||||
require.Equal(t, int64(2), result.Total)
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceListTasksNotReady(t *testing.T) {
|
||||
var nilSvc *UsageCleanupService
|
||||
_, _, err := nilSvc.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
|
||||
require.Error(t, err)
|
||||
|
||||
svc := NewUsageCleanupService(nil, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
|
||||
_, _, err = svc.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestUsageCleanupServiceDefaultsAndLifecycle(t *testing.T) {
|
||||
var nilSvc *UsageCleanupService
|
||||
require.Equal(t, 31, nilSvc.maxRangeDays())
|
||||
require.Equal(t, 5000, nilSvc.batchSize())
|
||||
require.Equal(t, 10*time.Second, nilSvc.workerInterval())
|
||||
require.Equal(t, 30*time.Minute, nilSvc.taskTimeout())
|
||||
nilSvc.Start()
|
||||
nilSvc.Stop()
|
||||
|
||||
repo := &cleanupRepoStub{}
|
||||
cfgDisabled := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: false}}
|
||||
svcDisabled := NewUsageCleanupService(repo, nil, nil, cfgDisabled)
|
||||
svcDisabled.Start()
|
||||
svcDisabled.Stop()
|
||||
|
||||
timingWheel, err := NewTimingWheelService()
|
||||
require.NoError(t, err)
|
||||
|
||||
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, WorkerIntervalSeconds: 5}}
|
||||
svc := NewUsageCleanupService(repo, timingWheel, nil, cfg)
|
||||
require.Equal(t, 5*time.Second, svc.workerInterval())
|
||||
svc.Start()
|
||||
svc.Stop()
|
||||
|
||||
cfgFallback := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
||||
svcFallback := NewUsageCleanupService(repo, timingWheel, nil, cfgFallback)
|
||||
require.Equal(t, 31, svcFallback.maxRangeDays())
|
||||
require.Equal(t, 5000, svcFallback.batchSize())
|
||||
require.Equal(t, 10*time.Second, svcFallback.workerInterval())
|
||||
|
||||
svcMissingDeps := NewUsageCleanupService(nil, nil, nil, cfgFallback)
|
||||
svcMissingDeps.Start()
|
||||
}
|
||||
|
||||
func TestSanitizeUsageCleanupFiltersModelEmpty(t *testing.T) {
|
||||
model := " "
|
||||
apiKeyID := int64(-5)
|
||||
accountID := int64(-1)
|
||||
groupID := int64(-2)
|
||||
filters := UsageCleanupFilters{
|
||||
UserID: &apiKeyID,
|
||||
APIKeyID: &apiKeyID,
|
||||
AccountID: &accountID,
|
||||
GroupID: &groupID,
|
||||
Model: &model,
|
||||
}
|
||||
|
||||
sanitizeUsageCleanupFilters(&filters)
|
||||
require.Nil(t, filters.UserID)
|
||||
require.Nil(t, filters.APIKeyID)
|
||||
require.Nil(t, filters.AccountID)
|
||||
require.Nil(t, filters.GroupID)
|
||||
require.Nil(t, filters.Model)
|
||||
}
|
||||
Reference in New Issue
Block a user