bd18f4b8ef
新增 usage_cleanup_task Ent schema 与仓储实现，支持清理任务排序分页 补充清理任务全链路日志、仪表盘重算触发及 UI 过滤调整 完善 repository/service 单测并引入 sqlite 测试依赖
378 lines
13 KiB
Go
378 lines
13 KiB
Go
package admin
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
|
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type cleanupRepoStub struct {
|
|
mu sync.Mutex
|
|
created []*service.UsageCleanupTask
|
|
listTasks []service.UsageCleanupTask
|
|
listResult *pagination.PaginationResult
|
|
listErr error
|
|
statusByID map[int64]string
|
|
}
|
|
|
|
func (s *cleanupRepoStub) CreateTask(ctx context.Context, task *service.UsageCleanupTask) error {
|
|
if task == nil {
|
|
return nil
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if task.ID == 0 {
|
|
task.ID = int64(len(s.created) + 1)
|
|
}
|
|
if task.CreatedAt.IsZero() {
|
|
task.CreatedAt = time.Now().UTC()
|
|
}
|
|
task.UpdatedAt = task.CreatedAt
|
|
clone := *task
|
|
s.created = append(s.created, &clone)
|
|
return nil
|
|
}
|
|
|
|
func (s *cleanupRepoStub) ListTasks(ctx context.Context, params pagination.PaginationParams) ([]service.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) (*service.UsageCleanupTask, error) {
|
|
return nil, 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 {
|
|
return nil
|
|
}
|
|
|
|
func (s *cleanupRepoStub) CancelTask(ctx context.Context, taskID int64, canceledBy int64) (bool, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.statusByID == nil {
|
|
s.statusByID = map[int64]string{}
|
|
}
|
|
status := s.statusByID[taskID]
|
|
if status != service.UsageCleanupStatusPending && status != service.UsageCleanupStatusRunning {
|
|
return false, nil
|
|
}
|
|
s.statusByID[taskID] = service.UsageCleanupStatusCanceled
|
|
return true, nil
|
|
}
|
|
|
|
func (s *cleanupRepoStub) MarkTaskSucceeded(ctx context.Context, taskID int64, deletedRows int64) error {
|
|
return nil
|
|
}
|
|
|
|
func (s *cleanupRepoStub) MarkTaskFailed(ctx context.Context, taskID int64, deletedRows int64, errorMsg string) error {
|
|
return nil
|
|
}
|
|
|
|
func (s *cleanupRepoStub) DeleteUsageLogsBatch(ctx context.Context, filters service.UsageCleanupFilters, limit int) (int64, error) {
|
|
return 0, nil
|
|
}
|
|
|
|
var _ service.UsageCleanupRepository = (*cleanupRepoStub)(nil)
|
|
|
|
func setupCleanupRouter(cleanupService *service.UsageCleanupService, userID int64) *gin.Engine {
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
if userID > 0 {
|
|
router.Use(func(c *gin.Context) {
|
|
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: userID})
|
|
c.Next()
|
|
})
|
|
}
|
|
|
|
handler := NewUsageHandler(nil, nil, nil, cleanupService)
|
|
router.POST("/api/v1/admin/usage/cleanup-tasks", handler.CreateCleanupTask)
|
|
router.GET("/api/v1/admin/usage/cleanup-tasks", handler.ListCleanupTasks)
|
|
router.POST("/api/v1/admin/usage/cleanup-tasks/:id/cancel", handler.CancelCleanupTask)
|
|
return router
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskUnauthorized(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 0)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewBufferString(`{}`))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskUnavailable(t *testing.T) {
|
|
router := setupCleanupRouter(nil, 1)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewBufferString(`{}`))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskBindError(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 88)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewBufferString("{bad-json"))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskMissingRange(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 88)
|
|
|
|
payload := map[string]any{
|
|
"start_date": "2024-01-01",
|
|
"timezone": "UTC",
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskInvalidDate(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 88)
|
|
|
|
payload := map[string]any{
|
|
"start_date": "2024-13-01",
|
|
"end_date": "2024-01-02",
|
|
"timezone": "UTC",
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskInvalidEndDate(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 88)
|
|
|
|
payload := map[string]any{
|
|
"start_date": "2024-01-01",
|
|
"end_date": "2024-02-40",
|
|
"timezone": "UTC",
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCreateCleanupTaskSuccess(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 99)
|
|
|
|
payload := map[string]any{
|
|
"start_date": " 2024-01-01 ",
|
|
"end_date": "2024-01-02",
|
|
"timezone": "UTC",
|
|
"model": "gpt-4",
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
|
|
var resp response.Response
|
|
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &resp))
|
|
require.Equal(t, 0, resp.Code)
|
|
|
|
repo.mu.Lock()
|
|
defer repo.mu.Unlock()
|
|
require.Len(t, repo.created, 1)
|
|
created := repo.created[0]
|
|
require.Equal(t, int64(99), created.CreatedBy)
|
|
require.NotNil(t, created.Filters.Model)
|
|
require.Equal(t, "gpt-4", *created.Filters.Model)
|
|
|
|
start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
|
|
end := time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC).Add(24*time.Hour - time.Nanosecond)
|
|
require.True(t, created.Filters.StartTime.Equal(start))
|
|
require.True(t, created.Filters.EndTime.Equal(end))
|
|
}
|
|
|
|
func TestUsageHandlerListCleanupTasksUnavailable(t *testing.T) {
|
|
router := setupCleanupRouter(nil, 0)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/usage/cleanup-tasks", nil)
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerListCleanupTasksSuccess(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
repo.listTasks = []service.UsageCleanupTask{
|
|
{
|
|
ID: 7,
|
|
Status: service.UsageCleanupStatusSucceeded,
|
|
CreatedBy: 4,
|
|
},
|
|
}
|
|
repo.listResult = &pagination.PaginationResult{Total: 1, Page: 1, PageSize: 20, Pages: 1}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 1)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/usage/cleanup-tasks", nil)
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
|
|
var resp struct {
|
|
Code int `json:"code"`
|
|
Data struct {
|
|
Items []dto.UsageCleanupTask `json:"items"`
|
|
Total int64 `json:"total"`
|
|
Page int `json:"page"`
|
|
} `json:"data"`
|
|
}
|
|
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &resp))
|
|
require.Equal(t, 0, resp.Code)
|
|
require.Len(t, resp.Data.Items, 1)
|
|
require.Equal(t, int64(7), resp.Data.Items[0].ID)
|
|
require.Equal(t, int64(1), resp.Data.Total)
|
|
require.Equal(t, 1, resp.Data.Page)
|
|
}
|
|
|
|
func TestUsageHandlerListCleanupTasksError(t *testing.T) {
|
|
repo := &cleanupRepoStub{listErr: errors.New("boom")}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 1)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/usage/cleanup-tasks", nil)
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, req)
|
|
|
|
require.Equal(t, http.StatusInternalServerError, recorder.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCancelCleanupTaskUnauthorized(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 0)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks/1/cancel", nil)
|
|
rec := httptest.NewRecorder()
|
|
router.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCancelCleanupTaskNotFound(t *testing.T) {
|
|
repo := &cleanupRepoStub{}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 1)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks/999/cancel", nil)
|
|
rec := httptest.NewRecorder()
|
|
router.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusNotFound, rec.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCancelCleanupTaskConflict(t *testing.T) {
|
|
repo := &cleanupRepoStub{statusByID: map[int64]string{2: service.UsageCleanupStatusSucceeded}}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 1)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks/2/cancel", nil)
|
|
rec := httptest.NewRecorder()
|
|
router.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusConflict, rec.Code)
|
|
}
|
|
|
|
func TestUsageHandlerCancelCleanupTaskSuccess(t *testing.T) {
|
|
repo := &cleanupRepoStub{statusByID: map[int64]string{3: service.UsageCleanupStatusPending}}
|
|
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
|
|
cleanupService := service.NewUsageCleanupService(repo, nil, nil, cfg)
|
|
router := setupCleanupRouter(cleanupService, 1)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/usage/cleanup-tasks/3/cancel", nil)
|
|
rec := httptest.NewRecorder()
|
|
router.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
}
|