usage_cleanup_service_test.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892
  1. package service
  2. import (
  3. "context"
  4. "database/sql"
  5. "errors"
  6. "net/http"
  7. "strings"
  8. "sync"
  9. "testing"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/config"
  12. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  14. "github.com/stretchr/testify/require"
  15. )
  16. type cleanupDeleteResponse struct {
  17. deleted int64
  18. err error
  19. }
  20. type cleanupDeleteCall struct {
  21. filters UsageCleanupFilters
  22. limit int
  23. }
  24. type cleanupMarkCall struct {
  25. taskID int64
  26. deletedRows int64
  27. errMsg string
  28. }
  29. type cleanupRepoStub struct {
  30. mu sync.Mutex
  31. created []*UsageCleanupTask
  32. createErr error
  33. listTasks []UsageCleanupTask
  34. listResult *pagination.PaginationResult
  35. listErr error
  36. claimQueue []*UsageCleanupTask
  37. claimErr error
  38. deleteQueue []cleanupDeleteResponse
  39. deleteCalls []cleanupDeleteCall
  40. markSucceeded []cleanupMarkCall
  41. markFailed []cleanupMarkCall
  42. statusByID map[int64]string
  43. statusErr error
  44. progressCalls []cleanupMarkCall
  45. updateErr error
  46. cancelCalls []int64
  47. cancelErr error
  48. cancelResult *bool
  49. markFailedErr error
  50. }
  51. type dashboardRepoStub struct {
  52. recomputeErr error
  53. recomputeCalls int
  54. }
  55. func (s *dashboardRepoStub) AggregateRange(ctx context.Context, start, end time.Time) error {
  56. return nil
  57. }
  58. func (s *dashboardRepoStub) RecomputeRange(ctx context.Context, start, end time.Time) error {
  59. s.recomputeCalls++
  60. return s.recomputeErr
  61. }
  62. func (s *dashboardRepoStub) GetAggregationWatermark(ctx context.Context) (time.Time, error) {
  63. return time.Time{}, nil
  64. }
  65. func (s *dashboardRepoStub) UpdateAggregationWatermark(ctx context.Context, aggregatedAt time.Time) error {
  66. return nil
  67. }
  68. func (s *dashboardRepoStub) CleanupAggregates(ctx context.Context, hourlyCutoff, dailyCutoff time.Time) error {
  69. return nil
  70. }
  71. func (s *dashboardRepoStub) CleanupUsageLogs(ctx context.Context, cutoff time.Time) error {
  72. return nil
  73. }
  74. func (s *dashboardRepoStub) CleanupUsageBillingDedup(ctx context.Context, cutoff time.Time) error {
  75. return nil
  76. }
  77. func (s *dashboardRepoStub) EnsureUsageLogsPartitions(ctx context.Context, now time.Time) error {
  78. return nil
  79. }
  80. func (s *cleanupRepoStub) CreateTask(ctx context.Context, task *UsageCleanupTask) error {
  81. if task == nil {
  82. return nil
  83. }
  84. s.mu.Lock()
  85. defer s.mu.Unlock()
  86. if s.createErr != nil {
  87. return s.createErr
  88. }
  89. if task.ID == 0 {
  90. task.ID = int64(len(s.created) + 1)
  91. }
  92. if task.CreatedAt.IsZero() {
  93. task.CreatedAt = time.Now().UTC()
  94. }
  95. if task.UpdatedAt.IsZero() {
  96. task.UpdatedAt = task.CreatedAt
  97. }
  98. clone := *task
  99. s.created = append(s.created, &clone)
  100. return nil
  101. }
  102. func (s *cleanupRepoStub) ListTasks(ctx context.Context, params pagination.PaginationParams) ([]UsageCleanupTask, *pagination.PaginationResult, error) {
  103. s.mu.Lock()
  104. defer s.mu.Unlock()
  105. return s.listTasks, s.listResult, s.listErr
  106. }
  107. func (s *cleanupRepoStub) ClaimNextPendingTask(ctx context.Context, staleRunningAfterSeconds int64) (*UsageCleanupTask, error) {
  108. s.mu.Lock()
  109. defer s.mu.Unlock()
  110. if s.claimErr != nil {
  111. return nil, s.claimErr
  112. }
  113. if len(s.claimQueue) == 0 {
  114. return nil, nil
  115. }
  116. task := s.claimQueue[0]
  117. s.claimQueue = s.claimQueue[1:]
  118. if s.statusByID == nil {
  119. s.statusByID = map[int64]string{}
  120. }
  121. s.statusByID[task.ID] = UsageCleanupStatusRunning
  122. return task, nil
  123. }
  124. func (s *cleanupRepoStub) GetTaskStatus(ctx context.Context, taskID int64) (string, error) {
  125. s.mu.Lock()
  126. defer s.mu.Unlock()
  127. if s.statusErr != nil {
  128. return "", s.statusErr
  129. }
  130. if s.statusByID == nil {
  131. return "", sql.ErrNoRows
  132. }
  133. status, ok := s.statusByID[taskID]
  134. if !ok {
  135. return "", sql.ErrNoRows
  136. }
  137. return status, nil
  138. }
  139. func (s *cleanupRepoStub) UpdateTaskProgress(ctx context.Context, taskID int64, deletedRows int64) error {
  140. s.mu.Lock()
  141. defer s.mu.Unlock()
  142. s.progressCalls = append(s.progressCalls, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows})
  143. if s.updateErr != nil {
  144. return s.updateErr
  145. }
  146. return nil
  147. }
  148. func (s *cleanupRepoStub) CancelTask(ctx context.Context, taskID int64, canceledBy int64) (bool, error) {
  149. s.mu.Lock()
  150. defer s.mu.Unlock()
  151. s.cancelCalls = append(s.cancelCalls, taskID)
  152. if s.cancelErr != nil {
  153. return false, s.cancelErr
  154. }
  155. if s.cancelResult != nil {
  156. ok := *s.cancelResult
  157. if ok {
  158. if s.statusByID == nil {
  159. s.statusByID = map[int64]string{}
  160. }
  161. s.statusByID[taskID] = UsageCleanupStatusCanceled
  162. }
  163. return ok, nil
  164. }
  165. if s.statusByID == nil {
  166. s.statusByID = map[int64]string{}
  167. }
  168. status := s.statusByID[taskID]
  169. if status != UsageCleanupStatusPending && status != UsageCleanupStatusRunning {
  170. return false, nil
  171. }
  172. s.statusByID[taskID] = UsageCleanupStatusCanceled
  173. return true, nil
  174. }
  175. func (s *cleanupRepoStub) MarkTaskSucceeded(ctx context.Context, taskID int64, deletedRows int64) error {
  176. s.mu.Lock()
  177. defer s.mu.Unlock()
  178. s.markSucceeded = append(s.markSucceeded, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows})
  179. if s.statusByID == nil {
  180. s.statusByID = map[int64]string{}
  181. }
  182. s.statusByID[taskID] = UsageCleanupStatusSucceeded
  183. return nil
  184. }
  185. func (s *cleanupRepoStub) MarkTaskFailed(ctx context.Context, taskID int64, deletedRows int64, errorMsg string) error {
  186. s.mu.Lock()
  187. defer s.mu.Unlock()
  188. s.markFailed = append(s.markFailed, cleanupMarkCall{taskID: taskID, deletedRows: deletedRows, errMsg: errorMsg})
  189. if s.statusByID == nil {
  190. s.statusByID = map[int64]string{}
  191. }
  192. s.statusByID[taskID] = UsageCleanupStatusFailed
  193. if s.markFailedErr != nil {
  194. return s.markFailedErr
  195. }
  196. return nil
  197. }
  198. func (s *cleanupRepoStub) DeleteUsageLogsBatch(ctx context.Context, filters UsageCleanupFilters, limit int) (int64, error) {
  199. s.mu.Lock()
  200. defer s.mu.Unlock()
  201. s.deleteCalls = append(s.deleteCalls, cleanupDeleteCall{filters: filters, limit: limit})
  202. if len(s.deleteQueue) == 0 {
  203. return 0, nil
  204. }
  205. resp := s.deleteQueue[0]
  206. s.deleteQueue = s.deleteQueue[1:]
  207. return resp.deleted, resp.err
  208. }
  209. func TestUsageCleanupServiceCreateTaskSanitizeFilters(t *testing.T) {
  210. repo := &cleanupRepoStub{}
  211. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 31}}
  212. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  213. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  214. end := start.Add(24 * time.Hour)
  215. userID := int64(-1)
  216. apiKeyID := int64(10)
  217. model := " gpt-4 "
  218. billingType := int8(-2)
  219. filters := UsageCleanupFilters{
  220. StartTime: start,
  221. EndTime: end,
  222. UserID: &userID,
  223. APIKeyID: &apiKeyID,
  224. Model: &model,
  225. BillingType: &billingType,
  226. }
  227. task, err := svc.CreateTask(context.Background(), filters, 9)
  228. require.NoError(t, err)
  229. require.Equal(t, UsageCleanupStatusPending, task.Status)
  230. require.Nil(t, task.Filters.UserID)
  231. require.NotNil(t, task.Filters.APIKeyID)
  232. require.Equal(t, apiKeyID, *task.Filters.APIKeyID)
  233. require.NotNil(t, task.Filters.Model)
  234. require.Equal(t, "gpt-4", *task.Filters.Model)
  235. require.Nil(t, task.Filters.BillingType)
  236. require.Equal(t, int64(9), task.CreatedBy)
  237. }
  238. func TestSanitizeUsageCleanupFiltersRequestTypePriority(t *testing.T) {
  239. requestType := int16(RequestTypeWSV2)
  240. stream := false
  241. model := " gpt-5 "
  242. filters := UsageCleanupFilters{
  243. Model: &model,
  244. RequestType: &requestType,
  245. Stream: &stream,
  246. }
  247. sanitizeUsageCleanupFilters(&filters)
  248. require.NotNil(t, filters.RequestType)
  249. require.Equal(t, int16(RequestTypeWSV2), *filters.RequestType)
  250. require.Nil(t, filters.Stream)
  251. require.NotNil(t, filters.Model)
  252. require.Equal(t, "gpt-5", *filters.Model)
  253. }
  254. func TestSanitizeUsageCleanupFiltersInvalidRequestType(t *testing.T) {
  255. requestType := int16(99)
  256. stream := true
  257. filters := UsageCleanupFilters{
  258. RequestType: &requestType,
  259. Stream: &stream,
  260. }
  261. sanitizeUsageCleanupFilters(&filters)
  262. require.Nil(t, filters.RequestType)
  263. require.NotNil(t, filters.Stream)
  264. require.True(t, *filters.Stream)
  265. }
  266. func TestDescribeUsageCleanupFiltersIncludesRequestType(t *testing.T) {
  267. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  268. end := start.Add(24 * time.Hour)
  269. requestType := int16(RequestTypeWSV2)
  270. desc := describeUsageCleanupFilters(UsageCleanupFilters{
  271. StartTime: start,
  272. EndTime: end,
  273. RequestType: &requestType,
  274. })
  275. require.Contains(t, desc, "request_type=ws_v2")
  276. }
  277. func TestUsageCleanupServiceCreateTaskInvalidCreator(t *testing.T) {
  278. repo := &cleanupRepoStub{}
  279. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  280. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  281. filters := UsageCleanupFilters{
  282. StartTime: time.Now(),
  283. EndTime: time.Now().Add(24 * time.Hour),
  284. }
  285. _, err := svc.CreateTask(context.Background(), filters, 0)
  286. require.Error(t, err)
  287. require.Equal(t, "USAGE_CLEANUP_INVALID_CREATOR", infraerrors.Reason(err))
  288. }
  289. func TestUsageCleanupServiceCreateTaskDisabled(t *testing.T) {
  290. repo := &cleanupRepoStub{}
  291. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: false}}
  292. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  293. filters := UsageCleanupFilters{
  294. StartTime: time.Now(),
  295. EndTime: time.Now().Add(24 * time.Hour),
  296. }
  297. _, err := svc.CreateTask(context.Background(), filters, 1)
  298. require.Error(t, err)
  299. require.Equal(t, http.StatusServiceUnavailable, infraerrors.Code(err))
  300. require.Equal(t, "USAGE_CLEANUP_DISABLED", infraerrors.Reason(err))
  301. }
  302. func TestUsageCleanupServiceCreateTaskRangeTooLarge(t *testing.T) {
  303. repo := &cleanupRepoStub{}
  304. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, MaxRangeDays: 1}}
  305. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  306. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  307. end := start.Add(48 * time.Hour)
  308. filters := UsageCleanupFilters{StartTime: start, EndTime: end}
  309. _, err := svc.CreateTask(context.Background(), filters, 1)
  310. require.Error(t, err)
  311. require.Equal(t, "USAGE_CLEANUP_RANGE_TOO_LARGE", infraerrors.Reason(err))
  312. }
  313. func TestUsageCleanupServiceCreateTaskMissingRange(t *testing.T) {
  314. repo := &cleanupRepoStub{}
  315. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  316. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  317. _, err := svc.CreateTask(context.Background(), UsageCleanupFilters{}, 1)
  318. require.Error(t, err)
  319. require.Equal(t, "USAGE_CLEANUP_MISSING_RANGE", infraerrors.Reason(err))
  320. }
  321. func TestUsageCleanupServiceCreateTaskRepoError(t *testing.T) {
  322. repo := &cleanupRepoStub{createErr: errors.New("db down")}
  323. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  324. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  325. filters := UsageCleanupFilters{
  326. StartTime: time.Now(),
  327. EndTime: time.Now().Add(24 * time.Hour),
  328. }
  329. _, err := svc.CreateTask(context.Background(), filters, 1)
  330. require.Error(t, err)
  331. require.Contains(t, err.Error(), "create cleanup task")
  332. }
  333. func TestUsageCleanupServiceRunOnceSuccess(t *testing.T) {
  334. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  335. end := start.Add(2 * time.Hour)
  336. repo := &cleanupRepoStub{
  337. claimQueue: []*UsageCleanupTask{
  338. {ID: 5, Filters: UsageCleanupFilters{StartTime: start, EndTime: end}},
  339. },
  340. deleteQueue: []cleanupDeleteResponse{
  341. {deleted: 2},
  342. {deleted: 2},
  343. {deleted: 1},
  344. },
  345. }
  346. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2, TaskTimeoutSeconds: 30}}
  347. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  348. svc.runOnce()
  349. repo.mu.Lock()
  350. defer repo.mu.Unlock()
  351. require.Len(t, repo.deleteCalls, 3)
  352. require.Equal(t, 2, repo.deleteCalls[0].limit)
  353. require.True(t, repo.deleteCalls[0].filters.StartTime.Equal(start))
  354. require.True(t, repo.deleteCalls[0].filters.EndTime.Equal(end))
  355. require.Len(t, repo.markSucceeded, 1)
  356. require.Empty(t, repo.markFailed)
  357. require.Equal(t, int64(5), repo.markSucceeded[0].taskID)
  358. require.Equal(t, int64(5), repo.markSucceeded[0].deletedRows)
  359. require.Equal(t, 2, repo.deleteCalls[0].limit)
  360. require.Equal(t, start, repo.deleteCalls[0].filters.StartTime)
  361. require.Equal(t, end, repo.deleteCalls[0].filters.EndTime)
  362. }
  363. func TestUsageCleanupServiceRunOnceClaimError(t *testing.T) {
  364. repo := &cleanupRepoStub{claimErr: errors.New("claim failed")}
  365. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  366. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  367. svc.runOnce()
  368. repo.mu.Lock()
  369. defer repo.mu.Unlock()
  370. require.Empty(t, repo.markSucceeded)
  371. require.Empty(t, repo.markFailed)
  372. }
  373. func TestUsageCleanupServiceRunOnceAlreadyRunning(t *testing.T) {
  374. repo := &cleanupRepoStub{}
  375. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  376. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  377. svc.running = 1
  378. svc.runOnce()
  379. }
  380. func TestUsageCleanupServiceExecuteTaskFailed(t *testing.T) {
  381. longMsg := strings.Repeat("x", 600)
  382. repo := &cleanupRepoStub{
  383. deleteQueue: []cleanupDeleteResponse{
  384. {err: errors.New(longMsg)},
  385. },
  386. }
  387. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 3}}
  388. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  389. task := &UsageCleanupTask{
  390. ID: 11,
  391. Filters: UsageCleanupFilters{
  392. StartTime: time.Now(),
  393. EndTime: time.Now().Add(24 * time.Hour),
  394. },
  395. }
  396. svc.executeTask(context.Background(), task)
  397. repo.mu.Lock()
  398. defer repo.mu.Unlock()
  399. require.Len(t, repo.markFailed, 1)
  400. require.Equal(t, int64(11), repo.markFailed[0].taskID)
  401. require.Equal(t, 500, len(repo.markFailed[0].errMsg))
  402. }
  403. func TestUsageCleanupServiceExecuteTaskProgressError(t *testing.T) {
  404. repo := &cleanupRepoStub{
  405. deleteQueue: []cleanupDeleteResponse{
  406. {deleted: 2},
  407. {deleted: 0},
  408. },
  409. updateErr: errors.New("update failed"),
  410. }
  411. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  412. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  413. task := &UsageCleanupTask{
  414. ID: 8,
  415. Filters: UsageCleanupFilters{
  416. StartTime: time.Now().UTC(),
  417. EndTime: time.Now().UTC().Add(time.Hour),
  418. },
  419. }
  420. svc.executeTask(context.Background(), task)
  421. repo.mu.Lock()
  422. defer repo.mu.Unlock()
  423. require.Len(t, repo.markSucceeded, 1)
  424. require.Empty(t, repo.markFailed)
  425. require.Len(t, repo.progressCalls, 1)
  426. }
  427. func TestUsageCleanupServiceExecuteTaskDeleteCanceled(t *testing.T) {
  428. repo := &cleanupRepoStub{
  429. deleteQueue: []cleanupDeleteResponse{
  430. {err: context.Canceled},
  431. },
  432. }
  433. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  434. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  435. task := &UsageCleanupTask{
  436. ID: 12,
  437. Filters: UsageCleanupFilters{
  438. StartTime: time.Now().UTC(),
  439. EndTime: time.Now().UTC().Add(time.Hour),
  440. },
  441. }
  442. svc.executeTask(context.Background(), task)
  443. repo.mu.Lock()
  444. defer repo.mu.Unlock()
  445. require.Empty(t, repo.markSucceeded)
  446. require.Empty(t, repo.markFailed)
  447. }
  448. func TestUsageCleanupServiceExecuteTaskContextCanceled(t *testing.T) {
  449. repo := &cleanupRepoStub{}
  450. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  451. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  452. task := &UsageCleanupTask{
  453. ID: 9,
  454. Filters: UsageCleanupFilters{
  455. StartTime: time.Now().UTC(),
  456. EndTime: time.Now().UTC().Add(time.Hour),
  457. },
  458. }
  459. ctx, cancel := context.WithCancel(context.Background())
  460. cancel()
  461. svc.executeTask(ctx, task)
  462. repo.mu.Lock()
  463. defer repo.mu.Unlock()
  464. require.Empty(t, repo.markSucceeded)
  465. require.Empty(t, repo.markFailed)
  466. require.Empty(t, repo.deleteCalls)
  467. }
  468. func TestUsageCleanupServiceExecuteTaskMarkFailedUpdateError(t *testing.T) {
  469. repo := &cleanupRepoStub{
  470. deleteQueue: []cleanupDeleteResponse{
  471. {err: errors.New("boom")},
  472. },
  473. markFailedErr: errors.New("update failed"),
  474. }
  475. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  476. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  477. task := &UsageCleanupTask{
  478. ID: 13,
  479. Filters: UsageCleanupFilters{
  480. StartTime: time.Now().UTC(),
  481. EndTime: time.Now().UTC().Add(time.Hour),
  482. },
  483. }
  484. svc.executeTask(context.Background(), task)
  485. repo.mu.Lock()
  486. defer repo.mu.Unlock()
  487. require.Len(t, repo.markFailed, 1)
  488. require.Equal(t, int64(13), repo.markFailed[0].taskID)
  489. }
  490. func TestUsageCleanupServiceExecuteTaskDashboardRecomputeError(t *testing.T) {
  491. dashboardRepo := &dashboardRepoStub{recomputeErr: errors.New("recompute failed")}
  492. repo := &cleanupRepoStub{
  493. deleteQueue: []cleanupDeleteResponse{
  494. {deleted: 0},
  495. },
  496. }
  497. dashboard := NewDashboardAggregationService(dashboardRepo, nil, &config.Config{
  498. DashboardAgg: config.DashboardAggregationConfig{Enabled: true},
  499. })
  500. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  501. svc := NewUsageCleanupService(repo, nil, dashboard, cfg)
  502. task := &UsageCleanupTask{
  503. ID: 14,
  504. Filters: UsageCleanupFilters{
  505. StartTime: time.Now().UTC(),
  506. EndTime: time.Now().UTC().Add(time.Hour),
  507. },
  508. }
  509. svc.executeTask(context.Background(), task)
  510. repo.mu.Lock()
  511. defer repo.mu.Unlock()
  512. require.Len(t, repo.markSucceeded, 1)
  513. require.Eventually(t, func() bool { return dashboardRepo.recomputeCalls == 1 }, time.Second, 10*time.Millisecond)
  514. }
  515. func TestUsageCleanupServiceExecuteTaskDashboardRecomputeSuccess(t *testing.T) {
  516. dashboardRepo := &dashboardRepoStub{}
  517. repo := &cleanupRepoStub{
  518. deleteQueue: []cleanupDeleteResponse{
  519. {deleted: 0},
  520. },
  521. }
  522. dashboard := NewDashboardAggregationService(dashboardRepo, nil, &config.Config{
  523. DashboardAgg: config.DashboardAggregationConfig{Enabled: true},
  524. })
  525. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  526. svc := NewUsageCleanupService(repo, nil, dashboard, cfg)
  527. task := &UsageCleanupTask{
  528. ID: 15,
  529. Filters: UsageCleanupFilters{
  530. StartTime: time.Now().UTC(),
  531. EndTime: time.Now().UTC().Add(time.Hour),
  532. },
  533. }
  534. svc.executeTask(context.Background(), task)
  535. repo.mu.Lock()
  536. defer repo.mu.Unlock()
  537. require.Len(t, repo.markSucceeded, 1)
  538. require.Eventually(t, func() bool { return dashboardRepo.recomputeCalls == 1 }, time.Second, 10*time.Millisecond)
  539. }
  540. func TestUsageCleanupServiceExecuteTaskCanceled(t *testing.T) {
  541. repo := &cleanupRepoStub{
  542. statusByID: map[int64]string{
  543. 3: UsageCleanupStatusCanceled,
  544. },
  545. }
  546. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, BatchSize: 2}}
  547. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  548. task := &UsageCleanupTask{
  549. ID: 3,
  550. Filters: UsageCleanupFilters{
  551. StartTime: time.Now().UTC(),
  552. EndTime: time.Now().UTC().Add(time.Hour),
  553. },
  554. }
  555. svc.executeTask(context.Background(), task)
  556. repo.mu.Lock()
  557. defer repo.mu.Unlock()
  558. require.Empty(t, repo.deleteCalls)
  559. require.Empty(t, repo.markSucceeded)
  560. require.Empty(t, repo.markFailed)
  561. }
  562. func TestUsageCleanupServiceCancelTaskSuccess(t *testing.T) {
  563. repo := &cleanupRepoStub{
  564. statusByID: map[int64]string{
  565. 5: UsageCleanupStatusPending,
  566. },
  567. }
  568. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  569. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  570. err := svc.CancelTask(context.Background(), 5, 9)
  571. require.NoError(t, err)
  572. repo.mu.Lock()
  573. defer repo.mu.Unlock()
  574. require.Equal(t, UsageCleanupStatusCanceled, repo.statusByID[5])
  575. require.Len(t, repo.cancelCalls, 1)
  576. }
  577. func TestUsageCleanupServiceCancelTaskDisabled(t *testing.T) {
  578. repo := &cleanupRepoStub{}
  579. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: false}}
  580. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  581. err := svc.CancelTask(context.Background(), 1, 2)
  582. require.Error(t, err)
  583. require.Equal(t, http.StatusServiceUnavailable, infraerrors.Code(err))
  584. require.Equal(t, "USAGE_CLEANUP_DISABLED", infraerrors.Reason(err))
  585. }
  586. func TestUsageCleanupServiceCancelTaskNotFound(t *testing.T) {
  587. repo := &cleanupRepoStub{}
  588. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  589. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  590. err := svc.CancelTask(context.Background(), 999, 1)
  591. require.Error(t, err)
  592. require.Equal(t, http.StatusNotFound, infraerrors.Code(err))
  593. require.Equal(t, "USAGE_CLEANUP_TASK_NOT_FOUND", infraerrors.Reason(err))
  594. }
  595. func TestUsageCleanupServiceCancelTaskStatusError(t *testing.T) {
  596. repo := &cleanupRepoStub{statusErr: errors.New("status broken")}
  597. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  598. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  599. err := svc.CancelTask(context.Background(), 7, 1)
  600. require.Error(t, err)
  601. require.Contains(t, err.Error(), "status broken")
  602. }
  603. func TestUsageCleanupServiceCancelTaskConflict(t *testing.T) {
  604. repo := &cleanupRepoStub{
  605. statusByID: map[int64]string{
  606. 7: UsageCleanupStatusSucceeded,
  607. },
  608. }
  609. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  610. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  611. err := svc.CancelTask(context.Background(), 7, 1)
  612. require.Error(t, err)
  613. require.Equal(t, http.StatusConflict, infraerrors.Code(err))
  614. require.Equal(t, "USAGE_CLEANUP_CANCEL_CONFLICT", infraerrors.Reason(err))
  615. }
  616. func TestUsageCleanupServiceCancelTaskAlreadyCanceledIsIdempotent(t *testing.T) {
  617. repo := &cleanupRepoStub{
  618. statusByID: map[int64]string{
  619. 7: UsageCleanupStatusCanceled,
  620. },
  621. }
  622. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  623. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  624. err := svc.CancelTask(context.Background(), 7, 1)
  625. require.NoError(t, err)
  626. repo.mu.Lock()
  627. defer repo.mu.Unlock()
  628. require.Empty(t, repo.cancelCalls, "already canceled should return success without extra cancel write")
  629. }
  630. func TestUsageCleanupServiceCancelTaskRepoConflict(t *testing.T) {
  631. shouldCancel := false
  632. repo := &cleanupRepoStub{
  633. statusByID: map[int64]string{
  634. 7: UsageCleanupStatusPending,
  635. },
  636. cancelResult: &shouldCancel,
  637. }
  638. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  639. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  640. err := svc.CancelTask(context.Background(), 7, 1)
  641. require.Error(t, err)
  642. require.Equal(t, http.StatusConflict, infraerrors.Code(err))
  643. require.Equal(t, "USAGE_CLEANUP_CANCEL_CONFLICT", infraerrors.Reason(err))
  644. }
  645. func TestUsageCleanupServiceCancelTaskRepoError(t *testing.T) {
  646. repo := &cleanupRepoStub{
  647. statusByID: map[int64]string{
  648. 7: UsageCleanupStatusPending,
  649. },
  650. cancelErr: errors.New("cancel failed"),
  651. }
  652. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  653. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  654. err := svc.CancelTask(context.Background(), 7, 1)
  655. require.Error(t, err)
  656. require.Contains(t, err.Error(), "cancel failed")
  657. }
  658. func TestUsageCleanupServiceCancelTaskInvalidCanceller(t *testing.T) {
  659. repo := &cleanupRepoStub{
  660. statusByID: map[int64]string{
  661. 7: UsageCleanupStatusRunning,
  662. },
  663. }
  664. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  665. svc := NewUsageCleanupService(repo, nil, nil, cfg)
  666. err := svc.CancelTask(context.Background(), 7, 0)
  667. require.Error(t, err)
  668. require.Equal(t, "USAGE_CLEANUP_INVALID_CANCELLER", infraerrors.Reason(err))
  669. }
  670. func TestUsageCleanupServiceListTasks(t *testing.T) {
  671. repo := &cleanupRepoStub{
  672. listTasks: []UsageCleanupTask{{ID: 1}, {ID: 2}},
  673. listResult: &pagination.PaginationResult{
  674. Total: 2,
  675. Page: 1,
  676. PageSize: 20,
  677. Pages: 1,
  678. },
  679. }
  680. svc := NewUsageCleanupService(repo, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
  681. tasks, result, err := svc.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  682. require.NoError(t, err)
  683. require.Len(t, tasks, 2)
  684. require.Equal(t, int64(2), result.Total)
  685. }
  686. func TestUsageCleanupServiceListTasksNotReady(t *testing.T) {
  687. var nilSvc *UsageCleanupService
  688. _, _, err := nilSvc.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  689. require.Error(t, err)
  690. svc := NewUsageCleanupService(nil, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
  691. _, _, err = svc.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  692. require.Error(t, err)
  693. }
  694. func TestUsageCleanupServiceDefaultsAndLifecycle(t *testing.T) {
  695. var nilSvc *UsageCleanupService
  696. require.Equal(t, 31, nilSvc.maxRangeDays())
  697. require.Equal(t, 5000, nilSvc.batchSize())
  698. require.Equal(t, 10*time.Second, nilSvc.workerInterval())
  699. require.Equal(t, 30*time.Minute, nilSvc.taskTimeout())
  700. nilSvc.Start()
  701. nilSvc.Stop()
  702. repo := &cleanupRepoStub{}
  703. cfgDisabled := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: false}}
  704. svcDisabled := NewUsageCleanupService(repo, nil, nil, cfgDisabled)
  705. svcDisabled.Start()
  706. svcDisabled.Stop()
  707. timingWheel, err := NewTimingWheelService()
  708. require.NoError(t, err)
  709. cfg := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true, WorkerIntervalSeconds: 5}}
  710. svc := NewUsageCleanupService(repo, timingWheel, nil, cfg)
  711. require.Equal(t, 5*time.Second, svc.workerInterval())
  712. svc.Start()
  713. svc.Stop()
  714. cfgFallback := &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}}
  715. svcFallback := NewUsageCleanupService(repo, timingWheel, nil, cfgFallback)
  716. require.Equal(t, 31, svcFallback.maxRangeDays())
  717. require.Equal(t, 5000, svcFallback.batchSize())
  718. require.Equal(t, 10*time.Second, svcFallback.workerInterval())
  719. svcMissingDeps := NewUsageCleanupService(nil, nil, nil, cfgFallback)
  720. svcMissingDeps.Start()
  721. }
  722. func TestSanitizeUsageCleanupFiltersModelEmpty(t *testing.T) {
  723. model := " "
  724. apiKeyID := int64(-5)
  725. accountID := int64(-1)
  726. groupID := int64(-2)
  727. filters := UsageCleanupFilters{
  728. UserID: &apiKeyID,
  729. APIKeyID: &apiKeyID,
  730. AccountID: &accountID,
  731. GroupID: &groupID,
  732. Model: &model,
  733. }
  734. sanitizeUsageCleanupFilters(&filters)
  735. require.Nil(t, filters.UserID)
  736. require.Nil(t, filters.APIKeyID)
  737. require.Nil(t, filters.AccountID)
  738. require.Nil(t, filters.GroupID)
  739. require.Nil(t, filters.Model)
  740. }
  741. func TestDescribeUsageCleanupFiltersAllFields(t *testing.T) {
  742. start := time.Date(2024, 2, 1, 10, 0, 0, 0, time.UTC)
  743. end := start.Add(2 * time.Hour)
  744. userID := int64(1)
  745. apiKeyID := int64(2)
  746. accountID := int64(3)
  747. groupID := int64(4)
  748. model := " gpt-4 "
  749. stream := true
  750. billingType := int8(2)
  751. filters := UsageCleanupFilters{
  752. StartTime: start,
  753. EndTime: end,
  754. UserID: &userID,
  755. APIKeyID: &apiKeyID,
  756. AccountID: &accountID,
  757. GroupID: &groupID,
  758. Model: &model,
  759. Stream: &stream,
  760. BillingType: &billingType,
  761. }
  762. desc := describeUsageCleanupFilters(filters)
  763. 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)
  764. }
  765. func TestUsageCleanupServiceIsTaskCanceledNotFound(t *testing.T) {
  766. repo := &cleanupRepoStub{}
  767. svc := NewUsageCleanupService(repo, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
  768. canceled, err := svc.isTaskCanceled(context.Background(), 9)
  769. require.NoError(t, err)
  770. require.False(t, canceled)
  771. }
  772. func TestUsageCleanupServiceIsTaskCanceledError(t *testing.T) {
  773. repo := &cleanupRepoStub{statusErr: errors.New("status err")}
  774. svc := NewUsageCleanupService(repo, nil, nil, &config.Config{UsageCleanup: config.UsageCleanupConfig{Enabled: true}})
  775. _, err := svc.isTaskCanceled(context.Background(), 9)
  776. require.Error(t, err)
  777. require.Contains(t, err.Error(), "status err")
  778. }