Commit bd18f4b8 authored by yangjianbo's avatar yangjianbo
Browse files

feat(清理任务): 引入Ent存储并补充日志与测试

新增 usage_cleanup_task Ent schema 与仓储实现,支持清理任务排序分页
补充清理任务全链路日志、仪表盘重算触发及 UI 过滤调整
完善 repository/service 单测并引入 sqlite 测试依赖
parent bf7b79f2
......@@ -200,6 +200,8 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
......@@ -225,6 +227,8 @@ github.com/redis/go-redis/v9 v9.17.2 h1:P2EGsA4qVIM3Pp+aPocCJ7DguDHhqrXNhVcEp4Vi
github.com/redis/go-redis/v9 v9.17.2/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370=
github.com/refraction-networking/utls v1.8.1 h1:yNY1kapmQU8JeM1sSw2H2asfTIwWxIkrMJI0pRUOCAo=
github.com/refraction-networking/utls v1.8.1/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
......@@ -339,6 +343,8 @@ golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
golang.org/x/exp v0.0.0-20230905200255-921286631fa9 h1:GoHiUyI/Tp2nVkLI2mCxVkOjsbSXD66ic0XW0js0R9g=
golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqRShND87VCbxmc6bL1Yd2oYrm6k=
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY=
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70=
golang.org/x/mod v0.30.0 h1:fDEXFVZ/fmCKProc/yAXXUijritrDzahmwwefnjoPFk=
golang.org/x/mod v0.30.0/go.mod h1:lAsf5O2EvJeSFMiBxXDki7sCgAxEUcZHXoXMKT4GJKc=
golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
......@@ -366,6 +372,7 @@ golang.org/x/tools v0.39.0 h1:ik4ho21kwuQln40uelmciQPp9SipgNDdrafrYA4TmQQ=
golang.org/x/tools v0.39.0/go.mod h1:JnefbkDPyD8UU2kI5fuf8ZX4/yUeh9W877ZeBONxUqQ=
golang.org/x/tools/go/expect v0.1.0-deprecated h1:jY2C5HGYR5lqex3gEniOQL0r7Dq5+VGVgY1nudX5lXY=
golang.org/x/tools/go/expect v0.1.0-deprecated/go.mod h1:eihoPOH+FgIqa3FpoTwguz/bVUSGBlGQU67vpBeOrBY=
golang.org/x/tools/go/expect v0.1.1-deprecated h1:jpBZDwmgPhXsKZC6WhL20P4b/wmnpsEAGHaNy0n/rJM=
golang.org/x/tools/go/packages/packagestest v0.1.1-deprecated h1:1h2MnaIAIXISqTFKdENegdpAgUXz6NrPEsbIeWaBRvM=
golang.org/x/tools/go/packages/packagestest v0.1.1-deprecated/go.mod h1:RVAQXBGNv1ib0J382/DPCRS/BPnsGebyM1Gj5VSDpG8=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
......@@ -388,4 +395,12 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q=
gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA=
modernc.org/libc v1.67.6 h1:eVOQvpModVLKOdT+LvBPjdQqfrZq+pC39BygcT+E7OI=
modernc.org/libc v1.67.6/go.mod h1:JAhxUVlolfYDErnwiqaLvUqc8nfb2r6S6slAgZOnaiE=
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
modernc.org/sqlite v1.44.1 h1:qybx/rNpfQipX/t47OxbHmkkJuv2JWifCMH8SVUiDas=
modernc.org/sqlite v1.44.1/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA=
rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
......@@ -3,8 +3,8 @@ package admin
import (
"bytes"
"context"
"encoding/json"
"database/sql"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
......
......@@ -7,43 +7,41 @@ import (
"errors"
"fmt"
"strings"
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
dbusagecleanuptask "github.com/Wei-Shaw/sub2api/ent/usagecleanuptask"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
)
type usageCleanupRepository struct {
client *dbent.Client
sql sqlExecutor
}
func NewUsageCleanupRepository(sqlDB *sql.DB) service.UsageCleanupRepository {
return &usageCleanupRepository{sql: sqlDB}
func NewUsageCleanupRepository(client *dbent.Client, sqlDB *sql.DB) service.UsageCleanupRepository {
return newUsageCleanupRepositoryWithSQL(client, sqlDB)
}
func newUsageCleanupRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *usageCleanupRepository {
return &usageCleanupRepository{client: client, sql: sqlq}
}
func (r *usageCleanupRepository) CreateTask(ctx context.Context, task *service.UsageCleanupTask) error {
if task == nil {
return nil
}
filtersJSON, err := json.Marshal(task.Filters)
if err != nil {
return fmt.Errorf("marshal cleanup filters: %w", err)
if r.client != nil {
return r.createTaskWithEnt(ctx, task)
}
query := `
INSERT INTO usage_cleanup_tasks (
status,
filters,
created_by,
deleted_rows
) VALUES ($1, $2, $3, $4)
RETURNING id, created_at, updated_at
`
if err := scanSingleRow(ctx, r.sql, query, []any{task.Status, filtersJSON, task.CreatedBy, task.DeletedRows}, &task.ID, &task.CreatedAt, &task.UpdatedAt); err != nil {
return err
}
return nil
return r.createTaskWithSQL(ctx, task)
}
func (r *usageCleanupRepository) ListTasks(ctx context.Context, params pagination.PaginationParams) ([]service.UsageCleanupTask, *pagination.PaginationResult, error) {
if r.client != nil {
return r.listTasksWithEnt(ctx, params)
}
var total int64
if err := scanSingleRow(ctx, r.sql, "SELECT COUNT(*) FROM usage_cleanup_tasks", nil, &total); err != nil {
return nil, nil, err
......@@ -57,14 +55,14 @@ func (r *usageCleanupRepository) ListTasks(ctx context.Context, params paginatio
canceled_by, canceled_at,
started_at, finished_at, created_at, updated_at
FROM usage_cleanup_tasks
ORDER BY created_at DESC
ORDER BY created_at DESC, id DESC
LIMIT $1 OFFSET $2
`
rows, err := r.sql.QueryContext(ctx, query, params.Limit(), params.Offset())
if err != nil {
return nil, nil, err
}
defer rows.Close()
defer func() { _ = rows.Close() }()
tasks := make([]service.UsageCleanupTask, 0)
for rows.Next() {
......@@ -194,6 +192,9 @@ func (r *usageCleanupRepository) ClaimNextPendingTask(ctx context.Context, stale
}
func (r *usageCleanupRepository) GetTaskStatus(ctx context.Context, taskID int64) (string, error) {
if r.client != nil {
return r.getTaskStatusWithEnt(ctx, taskID)
}
var status string
if err := scanSingleRow(ctx, r.sql, "SELECT status FROM usage_cleanup_tasks WHERE id = $1", []any{taskID}, &status); err != nil {
return "", err
......@@ -202,6 +203,9 @@ func (r *usageCleanupRepository) GetTaskStatus(ctx context.Context, taskID int64
}
func (r *usageCleanupRepository) UpdateTaskProgress(ctx context.Context, taskID int64, deletedRows int64) error {
if r.client != nil {
return r.updateTaskProgressWithEnt(ctx, taskID, deletedRows)
}
query := `
UPDATE usage_cleanup_tasks
SET deleted_rows = $1,
......@@ -213,6 +217,9 @@ func (r *usageCleanupRepository) UpdateTaskProgress(ctx context.Context, taskID
}
func (r *usageCleanupRepository) CancelTask(ctx context.Context, taskID int64, canceledBy int64) (bool, error) {
if r.client != nil {
return r.cancelTaskWithEnt(ctx, taskID, canceledBy)
}
query := `
UPDATE usage_cleanup_tasks
SET status = $1,
......@@ -243,6 +250,9 @@ func (r *usageCleanupRepository) CancelTask(ctx context.Context, taskID int64, c
}
func (r *usageCleanupRepository) MarkTaskSucceeded(ctx context.Context, taskID int64, deletedRows int64) error {
if r.client != nil {
return r.markTaskSucceededWithEnt(ctx, taskID, deletedRows)
}
query := `
UPDATE usage_cleanup_tasks
SET status = $1,
......@@ -256,6 +266,9 @@ func (r *usageCleanupRepository) MarkTaskSucceeded(ctx context.Context, taskID i
}
func (r *usageCleanupRepository) MarkTaskFailed(ctx context.Context, taskID int64, deletedRows int64, errorMsg string) error {
if r.client != nil {
return r.markTaskFailedWithEnt(ctx, taskID, deletedRows, errorMsg)
}
query := `
UPDATE usage_cleanup_tasks
SET status = $1,
......@@ -295,7 +308,7 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
if err != nil {
return 0, err
}
defer rows.Close()
defer func() { _ = rows.Close() }()
var deleted int64
for rows.Next() {
......@@ -357,7 +370,182 @@ func buildUsageCleanupWhere(filters service.UsageCleanupFilters) (string, []any)
if filters.BillingType != nil {
conditions = append(conditions, fmt.Sprintf("billing_type = $%d", idx))
args = append(args, *filters.BillingType)
idx++
}
return strings.Join(conditions, " AND "), args
}
func (r *usageCleanupRepository) createTaskWithEnt(ctx context.Context, task *service.UsageCleanupTask) error {
client := clientFromContext(ctx, r.client)
filtersJSON, err := json.Marshal(task.Filters)
if err != nil {
return fmt.Errorf("marshal cleanup filters: %w", err)
}
created, err := client.UsageCleanupTask.
Create().
SetStatus(task.Status).
SetFilters(json.RawMessage(filtersJSON)).
SetCreatedBy(task.CreatedBy).
SetDeletedRows(task.DeletedRows).
Save(ctx)
if err != nil {
return err
}
task.ID = created.ID
task.CreatedAt = created.CreatedAt
task.UpdatedAt = created.UpdatedAt
return nil
}
func (r *usageCleanupRepository) createTaskWithSQL(ctx context.Context, task *service.UsageCleanupTask) error {
filtersJSON, err := json.Marshal(task.Filters)
if err != nil {
return fmt.Errorf("marshal cleanup filters: %w", err)
}
query := `
INSERT INTO usage_cleanup_tasks (
status,
filters,
created_by,
deleted_rows
) VALUES ($1, $2, $3, $4)
RETURNING id, created_at, updated_at
`
if err := scanSingleRow(ctx, r.sql, query, []any{task.Status, filtersJSON, task.CreatedBy, task.DeletedRows}, &task.ID, &task.CreatedAt, &task.UpdatedAt); err != nil {
return err
}
return nil
}
func (r *usageCleanupRepository) listTasksWithEnt(ctx context.Context, params pagination.PaginationParams) ([]service.UsageCleanupTask, *pagination.PaginationResult, error) {
client := clientFromContext(ctx, r.client)
query := client.UsageCleanupTask.Query()
total, err := query.Clone().Count(ctx)
if err != nil {
return nil, nil, err
}
if total == 0 {
return []service.UsageCleanupTask{}, paginationResultFromTotal(0, params), nil
}
rows, err := query.
Order(dbent.Desc(dbusagecleanuptask.FieldCreatedAt), dbent.Desc(dbusagecleanuptask.FieldID)).
Offset(params.Offset()).
Limit(params.Limit()).
All(ctx)
if err != nil {
return nil, nil, err
}
tasks := make([]service.UsageCleanupTask, 0, len(rows))
for _, row := range rows {
task, err := usageCleanupTaskFromEnt(row)
if err != nil {
return nil, nil, err
}
tasks = append(tasks, task)
}
return tasks, paginationResultFromTotal(int64(total), params), nil
}
func (r *usageCleanupRepository) getTaskStatusWithEnt(ctx context.Context, taskID int64) (string, error) {
client := clientFromContext(ctx, r.client)
task, err := client.UsageCleanupTask.Query().
Where(dbusagecleanuptask.IDEQ(taskID)).
Only(ctx)
if err != nil {
if dbent.IsNotFound(err) {
return "", sql.ErrNoRows
}
return "", err
}
return task.Status, nil
}
func (r *usageCleanupRepository) updateTaskProgressWithEnt(ctx context.Context, taskID int64, deletedRows int64) error {
client := clientFromContext(ctx, r.client)
now := time.Now()
_, err := client.UsageCleanupTask.Update().
Where(dbusagecleanuptask.IDEQ(taskID)).
SetDeletedRows(deletedRows).
SetUpdatedAt(now).
Save(ctx)
return err
}
func (r *usageCleanupRepository) cancelTaskWithEnt(ctx context.Context, taskID int64, canceledBy int64) (bool, error) {
client := clientFromContext(ctx, r.client)
now := time.Now()
affected, err := client.UsageCleanupTask.Update().
Where(
dbusagecleanuptask.IDEQ(taskID),
dbusagecleanuptask.StatusIn(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning),
).
SetStatus(service.UsageCleanupStatusCanceled).
SetCanceledBy(canceledBy).
SetCanceledAt(now).
SetFinishedAt(now).
ClearErrorMessage().
SetUpdatedAt(now).
Save(ctx)
if err != nil {
return false, err
}
return affected > 0, nil
}
func (r *usageCleanupRepository) markTaskSucceededWithEnt(ctx context.Context, taskID int64, deletedRows int64) error {
client := clientFromContext(ctx, r.client)
now := time.Now()
_, err := client.UsageCleanupTask.Update().
Where(dbusagecleanuptask.IDEQ(taskID)).
SetStatus(service.UsageCleanupStatusSucceeded).
SetDeletedRows(deletedRows).
SetFinishedAt(now).
SetUpdatedAt(now).
Save(ctx)
return err
}
func (r *usageCleanupRepository) markTaskFailedWithEnt(ctx context.Context, taskID int64, deletedRows int64, errorMsg string) error {
client := clientFromContext(ctx, r.client)
now := time.Now()
_, err := client.UsageCleanupTask.Update().
Where(dbusagecleanuptask.IDEQ(taskID)).
SetStatus(service.UsageCleanupStatusFailed).
SetDeletedRows(deletedRows).
SetErrorMessage(errorMsg).
SetFinishedAt(now).
SetUpdatedAt(now).
Save(ctx)
return err
}
func usageCleanupTaskFromEnt(row *dbent.UsageCleanupTask) (service.UsageCleanupTask, error) {
task := service.UsageCleanupTask{
ID: row.ID,
Status: row.Status,
CreatedBy: row.CreatedBy,
DeletedRows: row.DeletedRows,
CreatedAt: row.CreatedAt,
UpdatedAt: row.UpdatedAt,
}
if len(row.Filters) > 0 {
if err := json.Unmarshal(row.Filters, &task.Filters); err != nil {
return service.UsageCleanupTask{}, fmt.Errorf("parse cleanup filters: %w", err)
}
}
if row.ErrorMessage != nil {
task.ErrorMsg = row.ErrorMessage
}
if row.CanceledBy != nil {
task.CanceledBy = row.CanceledBy
}
if row.CanceledAt != nil {
task.CanceledAt = row.CanceledAt
}
if row.StartedAt != nil {
task.StartedAt = row.StartedAt
}
if row.FinishedAt != nil {
task.FinishedAt = row.FinishedAt
}
return task, nil
}
package repository
import (
"context"
"database/sql"
"encoding/json"
"testing"
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/ent/enttest"
dbusagecleanuptask "github.com/Wei-Shaw/sub2api/ent/usagecleanuptask"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
"entgo.io/ent/dialect"
entsql "entgo.io/ent/dialect/sql"
_ "modernc.org/sqlite"
)
func newUsageCleanupEntRepo(t *testing.T) (*usageCleanupRepository, *dbent.Client) {
t.Helper()
db, err := sql.Open("sqlite", "file:usage_cleanup?mode=memory&cache=shared")
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
_, err = db.Exec("PRAGMA foreign_keys = ON")
require.NoError(t, err)
drv := entsql.OpenDB(dialect.SQLite, db)
client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(drv)))
t.Cleanup(func() { _ = client.Close() })
repo := &usageCleanupRepository{client: client, sql: db}
return repo, client
}
func TestUsageCleanupRepositoryEntCreateAndList(t *testing.T) {
repo, _ := newUsageCleanupEntRepo(t)
start := time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC)
end := start.Add(24 * time.Hour)
task := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusPending,
Filters: service.UsageCleanupFilters{StartTime: start, EndTime: end},
CreatedBy: 9,
}
require.NoError(t, repo.CreateTask(context.Background(), task))
require.NotZero(t, task.ID)
task2 := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusRunning,
Filters: service.UsageCleanupFilters{StartTime: start.Add(-24 * time.Hour), EndTime: end.Add(-24 * time.Hour)},
CreatedBy: 10,
}
require.NoError(t, repo.CreateTask(context.Background(), task2))
tasks, result, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 10})
require.NoError(t, err)
require.Len(t, tasks, 2)
require.Equal(t, int64(2), result.Total)
require.Greater(t, tasks[0].ID, tasks[1].ID)
require.Equal(t, start, tasks[1].Filters.StartTime)
require.Equal(t, end, tasks[1].Filters.EndTime)
}
func TestUsageCleanupRepositoryEntListEmpty(t *testing.T) {
repo, _ := newUsageCleanupEntRepo(t)
tasks, result, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 10})
require.NoError(t, err)
require.Empty(t, tasks)
require.Equal(t, int64(0), result.Total)
}
func TestUsageCleanupRepositoryEntGetStatusAndProgress(t *testing.T) {
repo, client := newUsageCleanupEntRepo(t)
task := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusPending,
Filters: service.UsageCleanupFilters{StartTime: time.Now().UTC(), EndTime: time.Now().UTC().Add(time.Hour)},
CreatedBy: 3,
}
require.NoError(t, repo.CreateTask(context.Background(), task))
status, err := repo.GetTaskStatus(context.Background(), task.ID)
require.NoError(t, err)
require.Equal(t, service.UsageCleanupStatusPending, status)
_, err = repo.GetTaskStatus(context.Background(), task.ID+99)
require.ErrorIs(t, err, sql.ErrNoRows)
require.NoError(t, repo.UpdateTaskProgress(context.Background(), task.ID, 42))
loaded, err := client.UsageCleanupTask.Get(context.Background(), task.ID)
require.NoError(t, err)
require.Equal(t, int64(42), loaded.DeletedRows)
}
func TestUsageCleanupRepositoryEntCancelAndFinish(t *testing.T) {
repo, client := newUsageCleanupEntRepo(t)
task := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusPending,
Filters: service.UsageCleanupFilters{StartTime: time.Now().UTC(), EndTime: time.Now().UTC().Add(time.Hour)},
CreatedBy: 5,
}
require.NoError(t, repo.CreateTask(context.Background(), task))
ok, err := repo.CancelTask(context.Background(), task.ID, 7)
require.NoError(t, err)
require.True(t, ok)
loaded, err := client.UsageCleanupTask.Get(context.Background(), task.ID)
require.NoError(t, err)
require.Equal(t, service.UsageCleanupStatusCanceled, loaded.Status)
require.NotNil(t, loaded.CanceledBy)
require.NotNil(t, loaded.CanceledAt)
require.NotNil(t, loaded.FinishedAt)
loaded.Status = service.UsageCleanupStatusSucceeded
_, err = client.UsageCleanupTask.Update().Where(dbusagecleanuptask.IDEQ(task.ID)).SetStatus(loaded.Status).Save(context.Background())
require.NoError(t, err)
ok, err = repo.CancelTask(context.Background(), task.ID, 7)
require.NoError(t, err)
require.False(t, ok)
}
func TestUsageCleanupRepositoryEntCancelError(t *testing.T) {
repo, client := newUsageCleanupEntRepo(t)
task := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusPending,
Filters: service.UsageCleanupFilters{StartTime: time.Now().UTC(), EndTime: time.Now().UTC().Add(time.Hour)},
CreatedBy: 5,
}
require.NoError(t, repo.CreateTask(context.Background(), task))
require.NoError(t, client.Close())
_, err := repo.CancelTask(context.Background(), task.ID, 7)
require.Error(t, err)
}
func TestUsageCleanupRepositoryEntMarkResults(t *testing.T) {
repo, client := newUsageCleanupEntRepo(t)
task := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusRunning,
Filters: service.UsageCleanupFilters{StartTime: time.Now().UTC(), EndTime: time.Now().UTC().Add(time.Hour)},
CreatedBy: 12,
}
require.NoError(t, repo.CreateTask(context.Background(), task))
require.NoError(t, repo.MarkTaskSucceeded(context.Background(), task.ID, 6))
loaded, err := client.UsageCleanupTask.Get(context.Background(), task.ID)
require.NoError(t, err)
require.Equal(t, service.UsageCleanupStatusSucceeded, loaded.Status)
require.Equal(t, int64(6), loaded.DeletedRows)
require.NotNil(t, loaded.FinishedAt)
task2 := &service.UsageCleanupTask{
Status: service.UsageCleanupStatusRunning,
Filters: service.UsageCleanupFilters{StartTime: time.Now().UTC(), EndTime: time.Now().UTC().Add(time.Hour)},
CreatedBy: 12,
}
require.NoError(t, repo.CreateTask(context.Background(), task2))
require.NoError(t, repo.MarkTaskFailed(context.Background(), task2.ID, 4, "boom"))
loaded2, err := client.UsageCleanupTask.Get(context.Background(), task2.ID)
require.NoError(t, err)
require.Equal(t, service.UsageCleanupStatusFailed, loaded2.Status)
require.Equal(t, "boom", *loaded2.ErrorMessage)
}
func TestUsageCleanupRepositoryEntInvalidStatus(t *testing.T) {
repo, _ := newUsageCleanupEntRepo(t)
task := &service.UsageCleanupTask{
Status: "invalid",
Filters: service.UsageCleanupFilters{StartTime: time.Now().UTC(), EndTime: time.Now().UTC().Add(time.Hour)},
CreatedBy: 1,
}
require.Error(t, repo.CreateTask(context.Background(), task))
}
func TestUsageCleanupRepositoryEntListInvalidFilters(t *testing.T) {
repo, client := newUsageCleanupEntRepo(t)
now := time.Now().UTC()
driver, ok := client.Driver().(*entsql.Driver)
require.True(t, ok)
_, err := driver.DB().ExecContext(
context.Background(),
`INSERT INTO usage_cleanup_tasks (status, filters, created_by, deleted_rows, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?)`,
service.UsageCleanupStatusPending,
[]byte("invalid-json"),
int64(1),
int64(0),
now,
now,
)
require.NoError(t, err)
_, _, err = repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 10})
require.Error(t, err)
}
func TestUsageCleanupTaskFromEntFull(t *testing.T) {
start := time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC)
end := start.Add(24 * time.Hour)
errMsg := "failed"
canceledBy := int64(2)
canceledAt := start.Add(time.Minute)
startedAt := start.Add(2 * time.Minute)
finishedAt := start.Add(3 * time.Minute)
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
filtersJSON, err := json.Marshal(filters)
require.NoError(t, err)
task, err := usageCleanupTaskFromEnt(&dbent.UsageCleanupTask{
ID: 10,
Status: service.UsageCleanupStatusFailed,
Filters: filtersJSON,
CreatedBy: 11,
DeletedRows: 7,
ErrorMessage: &errMsg,
CanceledBy: &canceledBy,
CanceledAt: &canceledAt,
StartedAt: &startedAt,
FinishedAt: &finishedAt,
CreatedAt: start,
UpdatedAt: end,
})
require.NoError(t, err)
require.Equal(t, int64(10), task.ID)
require.Equal(t, service.UsageCleanupStatusFailed, task.Status)
require.NotNil(t, task.ErrorMsg)
require.NotNil(t, task.CanceledBy)
require.NotNil(t, task.CanceledAt)
require.NotNil(t, task.StartedAt)
require.NotNil(t, task.FinishedAt)
}
func TestUsageCleanupTaskFromEntInvalidFilters(t *testing.T) {
task, err := usageCleanupTaskFromEnt(&dbent.UsageCleanupTask{
Filters: json.RawMessage("invalid-json"),
})
require.Error(t, err)
require.Empty(t, task)
}
......@@ -23,7 +23,7 @@ func newSQLMock(t *testing.T) (*sql.DB, sqlmock.Sqlmock) {
func TestNewUsageCleanupRepository(t *testing.T) {
db, _ := newSQLMock(t)
repo := NewUsageCleanupRepository(db)
repo := NewUsageCleanupRepository(nil, db)
require.NotNil(t, repo)
}
......@@ -146,6 +146,21 @@ func TestUsageCleanupRepositoryListTasks(t *testing.T) {
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryListTasksQueryError(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(2)))
mock.ExpectQuery("SELECT id, status, filters, created_by, deleted_rows, error_message").
WithArgs(20, 0).
WillReturnError(sql.ErrConnDone)
_, _, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
require.Error(t, err)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryListTasksInvalidFilters(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
......@@ -320,6 +335,19 @@ func TestUsageCleanupRepositoryGetTaskStatus(t *testing.T) {
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryGetTaskStatusQueryError(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
mock.ExpectQuery("SELECT status FROM usage_cleanup_tasks").
WithArgs(int64(9)).
WillReturnError(sql.ErrConnDone)
_, err := repo.GetTaskStatus(context.Background(), 9)
require.Error(t, err)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryUpdateTaskProgress(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
......@@ -347,6 +375,20 @@ func TestUsageCleanupRepositoryCancelTask(t *testing.T) {
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryCancelTaskNoRows(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
mock.ExpectQuery("UPDATE usage_cleanup_tasks").
WithArgs(service.UsageCleanupStatusCanceled, int64(6), int64(9), service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning).
WillReturnRows(sqlmock.NewRows([]string{"id"}))
ok, err := repo.CancelTask(context.Background(), 6, 9)
require.NoError(t, err)
require.False(t, ok)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryDeleteUsageLogsBatchMissingRange(t *testing.T) {
db, _ := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
......
......@@ -151,20 +151,24 @@ func (s *UsageCleanupService) CreateTask(ctx context.Context, filters UsageClean
}
func (s *UsageCleanupService) runOnce() {
if !atomic.CompareAndSwapInt32(&s.running, 0, 1) {
svc := s
if svc == nil {
return
}
if !atomic.CompareAndSwapInt32(&svc.running, 0, 1) {
log.Printf("[UsageCleanup] run_once skipped: already_running=true")
return
}
defer atomic.StoreInt32(&s.running, 0)
defer atomic.StoreInt32(&svc.running, 0)
parent := context.Background()
if s != nil && s.workerCtx != nil {
parent = s.workerCtx
if svc.workerCtx != nil {
parent = svc.workerCtx
}
ctx, cancel := context.WithTimeout(parent, s.taskTimeout())
ctx, cancel := context.WithTimeout(parent, svc.taskTimeout())
defer cancel()
task, err := s.repo.ClaimNextPendingTask(ctx, int64(s.taskTimeout().Seconds()))
task, err := svc.repo.ClaimNextPendingTask(ctx, int64(svc.taskTimeout().Seconds()))
if err != nil {
log.Printf("[UsageCleanup] claim pending task failed: %v", err)
return
......@@ -175,7 +179,7 @@ func (s *UsageCleanupService) runOnce() {
}
log.Printf("[UsageCleanup] task claimed: task=%d status=%s created_by=%d deleted_rows=%d %s", task.ID, task.Status, task.CreatedBy, task.DeletedRows, describeUsageCleanupFilters(task.Filters))
s.executeTask(ctx, task)
svc.executeTask(ctx, task)
}
func (s *UsageCleanupService) executeTask(ctx context.Context, task *UsageCleanupTask) {
......
......@@ -46,8 +46,45 @@ type cleanupRepoStub struct {
markSucceeded []cleanupMarkCall
markFailed []cleanupMarkCall
statusByID map[int64]string
statusErr error
progressCalls []cleanupMarkCall
updateErr error
cancelCalls []int64
cancelErr error
cancelResult *bool
markFailedErr error
}
type dashboardRepoStub struct {
recomputeErr error
}
func (s *dashboardRepoStub) AggregateRange(ctx context.Context, start, end time.Time) error {
return nil
}
func (s *dashboardRepoStub) RecomputeRange(ctx context.Context, start, end time.Time) error {
return s.recomputeErr
}
func (s *dashboardRepoStub) GetAggregationWatermark(ctx context.Context) (time.Time, error) {
return time.Time{}, nil
}
func (s *dashboardRepoStub) UpdateAggregationWatermark(ctx context.Context, aggregatedAt time.Time) error {
return nil
}
func (s *dashboardRepoStub) CleanupAggregates(ctx context.Context, hourlyCutoff, dailyCutoff time.Time) error {
return nil
}
func (s *dashboardRepoStub) CleanupUsageLogs(ctx context.Context, cutoff time.Time) error {
return nil
}
func (s *dashboardRepoStub) EnsureUsageLogsPartitions(ctx context.Context, now time.Time) error {
return nil
}
func (s *cleanupRepoStub) CreateTask(ctx context.Context, task *UsageCleanupTask) error {
......@@ -100,6 +137,9 @@ func (s *cleanupRepoStub) ClaimNextPendingTask(ctx context.Context, staleRunning
func (s *cleanupRepoStub) GetTaskStatus(ctx context.Context, taskID int64) (string, error) {
s.mu.Lock()
defer s.mu.Unlock()
if s.statusErr != nil {
return "", s.statusErr
}
if s.statusByID == nil {
return "", sql.ErrNoRows
}
......@@ -114,6 +154,9 @@ func (s *cleanupRepoStub) UpdateTaskProgress(ctx context.Context, taskID int64,
s.mu.Lock()
defer s.mu.Unlock()
s.progressCalls = append(s.progressCalls, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows})
if s.updateErr != nil {
return s.updateErr
}
return nil
}
......@@ -121,6 +164,19 @@ func (s *cleanupRepoStub) CancelTask(ctx context.Context, taskID int64, canceled
s.mu.Lock()
defer s.mu.Unlock()
s.cancelCalls = append(s.cancelCalls, taskID)
if s.cancelErr != nil {
return false, s.cancelErr
}
if s.cancelResult != nil {
ok := *s.cancelResult
if ok {
if s.statusByID == nil {
s.statusByID = map[int64]string{}
}
s.statusByID[taskID] = UsageCleanupStatusCanceled
}
return ok, nil
}
if s.statusByID == nil {
s.statusByID = map[int64]string{}
}
......@@ -151,6 +207,9 @@ func (s *cleanupRepoStub) MarkTaskFailed(ctx context.Context, taskID int64, dele
s.statusByID = map[int64]string{}
}
s.statusByID[taskID] = UsageCleanupStatusFailed
if s.markFailedErr != nil {
return s.markFailedErr
}
return nil
}
......@@ -266,9 +325,11 @@ func TestUsageCleanupServiceCreateTaskRepoError(t *testing.T) {
}
func TestUsageCleanupServiceRunOnceSuccess(t *testing.T) {
start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
end := start.Add(2 * time.Hour)
repo := &cleanupRepoStub{
claimQueue: []*UsageCleanupTask{
{ID: 5, Filters: UsageCleanupFilters{StartTime: time.Now(), EndTime: time.Now().Add(2 * time.Hour)}},
{ID: 5, Filters: UsageCleanupFilters{StartTime: start, EndTime: end}},
},
deleteQueue: []cleanupDeleteResponse{
{deleted: 2},
......@@ -288,6 +349,9 @@ func TestUsageCleanupServiceRunOnceSuccess(t *testing.T) {
require.Empty(t, repo.markFailed)
require.Equal(t, int64(5), repo.markSucceeded[0].taskID)
require.Equal(t, int64(5), repo.markSucceeded[0].deletedRows)
require.Equal(t, 2, repo.deleteCalls[0].limit)
require.Equal(t, start, repo.deleteCalls[0].filters.StartTime)
require.Equal(t, end, repo.deleteCalls[0].filters.EndTime)
}
func TestUsageCleanupServiceRunOnceClaimError(t *testing.T) {
......@@ -336,6 +400,293 @@ func TestUsageCleanupServiceExecuteTaskFailed(t *testing.T) {
require.Equal(t, 500, len(repo.markFailed[0].errMsg))
}
func TestUsageCleanupServiceExecuteTaskProgressError(t *testing.T) {
repo := &cleanupRepoStub{
deleteQueue: []cleanupDeleteResponse{
{deleted: 2},
{deleted: 0},
},
updateErr: errors.New("update failed"),
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
task := &UsageCleanupTask{
ID: 8,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
svc.executeTask(context.Background(), task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Len(t, repo.markSucceeded, 1)
require.Empty(t, repo.markFailed)
require.Len(t, repo.progressCalls, 1)
}
func TestUsageCleanupServiceExecuteTaskDeleteCanceled(t *testing.T) {
repo := &cleanupRepoStub{
deleteQueue: []cleanupDeleteResponse{
{err: context.Canceled},
},
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
task := &UsageCleanupTask{
ID: 12,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
svc.executeTask(context.Background(), task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Empty(t, repo.markSucceeded)
require.Empty(t, repo.markFailed)
}
func TestUsageCleanupServiceExecuteTaskContextCanceled(t *testing.T) {
repo := &cleanupRepoStub{}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
task := &UsageCleanupTask{
ID: 9,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
svc.executeTask(ctx, task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Empty(t, repo.markSucceeded)
require.Empty(t, repo.markFailed)
require.Empty(t, repo.deleteCalls)
}
func TestUsageCleanupServiceExecuteTaskMarkFailedUpdateError(t *testing.T) {
repo := &cleanupRepoStub{
deleteQueue: []cleanupDeleteResponse{
{err: errors.New("boom")},
},
markFailedErr: errors.New("update failed"),
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
task := &UsageCleanupTask{
ID: 13,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
svc.executeTask(context.Background(), task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Len(t, repo.markFailed, 1)
require.Equal(t, int64(13), repo.markFailed[0].taskID)
}
func TestUsageCleanupServiceExecuteTaskDashboardRecomputeError(t *testing.T) {
repo := &cleanupRepoStub{
deleteQueue: []cleanupDeleteResponse{
{deleted: 0},
},
}
dashboard := NewDashboardAggregationService(&dashboardRepoStub{}, nil, &config.Config{
DashboardAgg: config.DashboardAggregationConfig{Enabled: false},
})
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, dashboard, cfg)
task := &UsageCleanupTask{
ID: 14,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
svc.executeTask(context.Background(), task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Len(t, repo.markSucceeded, 1)
}
func TestUsageCleanupServiceExecuteTaskDashboardRecomputeSuccess(t *testing.T) {
repo := &cleanupRepoStub{
deleteQueue: []cleanupDeleteResponse{
{deleted: 0},
},
}
dashboard := NewDashboardAggregationService(&dashboardRepoStub{}, nil, &config.Config{
DashboardAgg: config.DashboardAggregationConfig{Enabled: true},
})
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, dashboard, cfg)
task := &UsageCleanupTask{
ID: 15,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
svc.executeTask(context.Background(), task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Len(t, repo.markSucceeded, 1)
}
func TestUsageCleanupServiceExecuteTaskCanceled(t *testing.T) {
repo := &cleanupRepoStub{
statusByID: map[int64]string{
3: UsageCleanupStatusCanceled,
},
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
task := &UsageCleanupTask{
ID: 3,
Filters: UsageCleanupFilters{
StartTime: time.Now().UTC(),
EndTime: time.Now().UTC().Add(time.Hour),
},
}
svc.executeTask(context.Background(), task)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Empty(t, repo.deleteCalls)
require.Empty(t, repo.markSucceeded)
require.Empty(t, repo.markFailed)
}
func TestUsageCleanupServiceCancelTaskSuccess(t *testing.T) {
repo := &cleanupRepoStub{
statusByID: map[int64]string{
5: UsageCleanupStatusPending,
},
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 5, 9)
require.NoError(t, err)
repo.mu.Lock()
defer repo.mu.Unlock()
require.Equal(t, UsageCleanupStatusCanceled, repo.statusByID[5])
require.Len(t, repo.cancelCalls, 1)
}
func TestUsageCleanupServiceCancelTaskDisabled(t *testing.T) {
repo := &cleanupRepoStub{}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: false}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 1, 2)
require.Error(t, err)
require.Equal(t, http.StatusServiceUnavailable, infraerrors.Code(err))
require.Equal(t, "USAGE_CLEANUP_DISABLED", infraerrors.Reason(err))
}
func TestUsageCleanupServiceCancelTaskNotFound(t *testing.T) {
repo := &cleanupRepoStub{}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 999, 1)
require.Error(t, err)
require.Equal(t, http.StatusNotFound, infraerrors.Code(err))
require.Equal(t, "USAGE_CLEANUP_TASK_NOT_FOUND", infraerrors.Reason(err))
}
func TestUsageCleanupServiceCancelTaskStatusError(t *testing.T) {
repo := &cleanupRepoStub{statusErr: errors.New("status broken")}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 7, 1)
require.Error(t, err)
require.Contains(t, err.Error(), "status broken")
}
func TestUsageCleanupServiceCancelTaskConflict(t *testing.T) {
repo := &cleanupRepoStub{
statusByID: map[int64]string{
7: UsageCleanupStatusSucceeded,
},
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 7, 1)
require.Error(t, err)
require.Equal(t, http.StatusConflict, infraerrors.Code(err))
require.Equal(t, "USAGE_CLEANUP_CANCEL_CONFLICT", infraerrors.Reason(err))
}
func TestUsageCleanupServiceCancelTaskRepoConflict(t *testing.T) {
shouldCancel := false
repo := &cleanupRepoStub{
statusByID: map[int64]string{
7: UsageCleanupStatusPending,
},
cancelResult: &shouldCancel,
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 7, 1)
require.Error(t, err)
require.Equal(t, http.StatusConflict, infraerrors.Code(err))
require.Equal(t, "USAGE_CLEANUP_CANCEL_CONFLICT", infraerrors.Reason(err))
}
func TestUsageCleanupServiceCancelTaskRepoError(t *testing.T) {
repo := &cleanupRepoStub{
statusByID: map[int64]string{
7: UsageCleanupStatusPending,
},
cancelErr: errors.New("cancel failed"),
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 7, 1)
require.Error(t, err)
require.Contains(t, err.Error(), "cancel failed")
}
func TestUsageCleanupServiceCancelTaskInvalidCanceller(t *testing.T) {
repo := &cleanupRepoStub{
statusByID: map[int64]string{
7: UsageCleanupStatusRunning,
},
}
cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
svc := NewUsageCleanupService(repo, nil, nil, cfg)
err := svc.CancelTask(context.Background(), 7, 0)
require.Error(t, err)
require.Equal(t, "USAGE_CLEANUP_INVALID_CANCELLER", infraerrors.Reason(err))
}
func TestUsageCleanupServiceListTasks(t *testing.T) {
repo := &cleanupRepoStub{
listTasks: []UsageCleanupTask{{ID: 1}, {ID: 2}},
......@@ -418,3 +769,47 @@ func TestSanitizeUsageCleanupFiltersModelEmpty(t *testing.T) {
require.Nil(t, filters.GroupID)
require.Nil(t, filters.Model)
}
func TestDescribeUsageCleanupFiltersAllFields(t *testing.T) {
start := time.Date(2024, 2, 1, 10, 0, 0, 0, time.UTC)
end := start.Add(2 * time.Hour)
userID := int64(1)
apiKeyID := int64(2)
accountID := int64(3)
groupID := int64(4)
model := " gpt-4 "
stream := true
billingType := int8(2)
filters := UsageCleanupFilters{
StartTime: start,
EndTime: end,
UserID: &userID,
APIKeyID: &apiKeyID,
AccountID: &accountID,
GroupID: &groupID,
Model: &model,
Stream: &stream,
BillingType: &billingType,
}
desc := describeUsageCleanupFilters(filters)
require.Equal(t, "start=2024-02-01T10:00:00Z end=2024-02-01T12:00:00Z user_id=1 api_key_id=2 account_id=3 group_id=4 model=gpt-4 stream=true billing_type=2", desc)
}
func TestUsageCleanupServiceIsTaskCanceledNotFound(t *testing.T) {
repo := &cleanupRepoStub{}
svc := NewUsageCleanupService(repo, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
canceled, err := svc.isTaskCanceled(context.Background(), 9)
require.NoError(t, err)
require.False(t, canceled)
}
func TestUsageCleanupServiceIsTaskCanceledError(t *testing.T) {
repo := &cleanupRepoStub{statusErr: errors.New("status err")}
svc := NewUsageCleanupService(repo, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
_, err := svc.isTaskCanceled(context.Background(), 9)
require.Error(t, err)
require.Contains(t, err.Error(), "status err")
}
......@@ -219,7 +219,7 @@ const loadTasks = async () => {
if (!props.show) return
tasksLoading.value = true
try {
const res = await adminUsageAPI.listCleanupTasks({ page: 1, page_size: 10 })
const res = await adminUsageAPI.listCleanupTasks({ page: 1, page_size: 5 })
tasks.value = res.items || []
} catch (error) {
console.error('Failed to load cleanup tasks:', error)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment