usage_cleanup_repo_test.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587
  1. package repository
  2. import (
  3. "context"
  4. "database/sql"
  5. "encoding/json"
  6. "testing"
  7. "time"
  8. "github.com/DATA-DOG/go-sqlmock"
  9. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  10. "github.com/Wei-Shaw/sub2api/internal/service"
  11. "github.com/stretchr/testify/require"
  12. )
  13. func newSQLMock(t *testing.T) (*sql.DB, sqlmock.Sqlmock) {
  14. t.Helper()
  15. db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
  16. require.NoError(t, err)
  17. t.Cleanup(func() { _ = db.Close() })
  18. return db, mock
  19. }
  20. func setUsageCleanupRollupTestTimezone(t *testing.T) {
  21. t.Helper()
  22. useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
  23. }
  24. func TestNewUsageCleanupRepository(t *testing.T) {
  25. db, _ := newSQLMock(t)
  26. repo := NewUsageCleanupRepository(nil, db)
  27. require.NotNil(t, repo)
  28. }
  29. func TestUsageCleanupRepositoryCreateTask(t *testing.T) {
  30. db, mock := newSQLMock(t)
  31. repo := &usageCleanupRepository{sql: db}
  32. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  33. end := start.Add(24 * time.Hour)
  34. task := &service.UsageCleanupTask{
  35. Status: service.UsageCleanupStatusPending,
  36. Filters: service.UsageCleanupFilters{StartTime: start, EndTime: end},
  37. CreatedBy: 12,
  38. }
  39. now := time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC)
  40. mock.ExpectQuery("INSERT INTO usage_cleanup_tasks").
  41. WithArgs(task.Status, sqlmock.AnyArg(), task.CreatedBy, task.DeletedRows).
  42. WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(1), now, now))
  43. err := repo.CreateTask(context.Background(), task)
  44. require.NoError(t, err)
  45. require.Equal(t, int64(1), task.ID)
  46. require.Equal(t, now, task.CreatedAt)
  47. require.Equal(t, now, task.UpdatedAt)
  48. require.NoError(t, mock.ExpectationsWereMet())
  49. }
  50. func TestUsageCleanupRepositoryCreateTaskNil(t *testing.T) {
  51. db, mock := newSQLMock(t)
  52. repo := &usageCleanupRepository{sql: db}
  53. err := repo.CreateTask(context.Background(), nil)
  54. require.NoError(t, err)
  55. require.NoError(t, mock.ExpectationsWereMet())
  56. }
  57. func TestUsageCleanupRepositoryCreateTaskQueryError(t *testing.T) {
  58. db, mock := newSQLMock(t)
  59. repo := &usageCleanupRepository{sql: db}
  60. task := &service.UsageCleanupTask{
  61. Status: service.UsageCleanupStatusPending,
  62. Filters: service.UsageCleanupFilters{StartTime: time.Now(), EndTime: time.Now().Add(time.Hour)},
  63. CreatedBy: 1,
  64. }
  65. mock.ExpectQuery("INSERT INTO usage_cleanup_tasks").
  66. WithArgs(task.Status, sqlmock.AnyArg(), task.CreatedBy, task.DeletedRows).
  67. WillReturnError(sql.ErrConnDone)
  68. err := repo.CreateTask(context.Background(), task)
  69. require.Error(t, err)
  70. require.NoError(t, mock.ExpectationsWereMet())
  71. }
  72. func TestUsageCleanupRepositoryListTasksEmpty(t *testing.T) {
  73. db, mock := newSQLMock(t)
  74. repo := &usageCleanupRepository{sql: db}
  75. mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
  76. WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(0)))
  77. tasks, result, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  78. require.NoError(t, err)
  79. require.Empty(t, tasks)
  80. require.Equal(t, int64(0), result.Total)
  81. require.NoError(t, mock.ExpectationsWereMet())
  82. }
  83. func TestUsageCleanupRepositoryListTasks(t *testing.T) {
  84. db, mock := newSQLMock(t)
  85. repo := &usageCleanupRepository{sql: db}
  86. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  87. end := start.Add(2 * time.Hour)
  88. filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
  89. filtersJSON, err := json.Marshal(filters)
  90. require.NoError(t, err)
  91. createdAt := time.Date(2024, 1, 2, 12, 0, 0, 0, time.UTC)
  92. updatedAt := createdAt.Add(time.Minute)
  93. rows := sqlmock.NewRows([]string{
  94. "id", "status", "filters", "created_by", "deleted_rows", "error_message",
  95. "canceled_by", "canceled_at",
  96. "started_at", "finished_at", "created_at", "updated_at",
  97. }).AddRow(
  98. int64(1),
  99. service.UsageCleanupStatusSucceeded,
  100. filtersJSON,
  101. int64(2),
  102. int64(9),
  103. "error",
  104. nil,
  105. nil,
  106. start,
  107. end,
  108. createdAt,
  109. updatedAt,
  110. )
  111. mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
  112. WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(1)))
  113. mock.ExpectQuery("SELECT id, status, filters, created_by, deleted_rows, error_message").
  114. WithArgs(20, 0).
  115. WillReturnRows(rows)
  116. tasks, result, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  117. require.NoError(t, err)
  118. require.Len(t, tasks, 1)
  119. require.Equal(t, int64(1), tasks[0].ID)
  120. require.Equal(t, service.UsageCleanupStatusSucceeded, tasks[0].Status)
  121. require.Equal(t, int64(2), tasks[0].CreatedBy)
  122. require.Equal(t, int64(9), tasks[0].DeletedRows)
  123. require.NotNil(t, tasks[0].ErrorMsg)
  124. require.Equal(t, "error", *tasks[0].ErrorMsg)
  125. require.NotNil(t, tasks[0].StartedAt)
  126. require.NotNil(t, tasks[0].FinishedAt)
  127. require.Equal(t, int64(1), result.Total)
  128. require.NoError(t, mock.ExpectationsWereMet())
  129. }
  130. func TestUsageCleanupRepositoryListTasksQueryError(t *testing.T) {
  131. db, mock := newSQLMock(t)
  132. repo := &usageCleanupRepository{sql: db}
  133. mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
  134. WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(2)))
  135. mock.ExpectQuery("SELECT id, status, filters, created_by, deleted_rows, error_message").
  136. WithArgs(20, 0).
  137. WillReturnError(sql.ErrConnDone)
  138. _, _, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  139. require.Error(t, err)
  140. require.NoError(t, mock.ExpectationsWereMet())
  141. }
  142. func TestUsageCleanupRepositoryListTasksInvalidFilters(t *testing.T) {
  143. db, mock := newSQLMock(t)
  144. repo := &usageCleanupRepository{sql: db}
  145. rows := sqlmock.NewRows([]string{
  146. "id", "status", "filters", "created_by", "deleted_rows", "error_message",
  147. "canceled_by", "canceled_at",
  148. "started_at", "finished_at", "created_at", "updated_at",
  149. }).AddRow(
  150. int64(1),
  151. service.UsageCleanupStatusSucceeded,
  152. []byte("not-json"),
  153. int64(2),
  154. int64(9),
  155. nil,
  156. nil,
  157. nil,
  158. nil,
  159. nil,
  160. time.Now().UTC(),
  161. time.Now().UTC(),
  162. )
  163. mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
  164. WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(1)))
  165. mock.ExpectQuery("SELECT id, status, filters, created_by, deleted_rows, error_message").
  166. WithArgs(20, 0).
  167. WillReturnRows(rows)
  168. _, _, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
  169. require.Error(t, err)
  170. require.NoError(t, mock.ExpectationsWereMet())
  171. }
  172. func TestUsageCleanupRepositoryClaimNextPendingTaskNone(t *testing.T) {
  173. db, mock := newSQLMock(t)
  174. repo := &usageCleanupRepository{sql: db}
  175. mock.ExpectQuery("UPDATE usage_cleanup_tasks").
  176. WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
  177. WillReturnRows(sqlmock.NewRows([]string{
  178. "id", "status", "filters", "created_by", "deleted_rows", "error_message",
  179. "started_at", "finished_at", "created_at", "updated_at",
  180. }))
  181. task, err := repo.ClaimNextPendingTask(context.Background(), 1800)
  182. require.NoError(t, err)
  183. require.Nil(t, task)
  184. require.NoError(t, mock.ExpectationsWereMet())
  185. }
  186. func TestUsageCleanupRepositoryClaimNextPendingTask(t *testing.T) {
  187. db, mock := newSQLMock(t)
  188. repo := &usageCleanupRepository{sql: db}
  189. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  190. end := start.Add(24 * time.Hour)
  191. filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
  192. filtersJSON, err := json.Marshal(filters)
  193. require.NoError(t, err)
  194. rows := sqlmock.NewRows([]string{
  195. "id", "status", "filters", "created_by", "deleted_rows", "error_message",
  196. "started_at", "finished_at", "created_at", "updated_at",
  197. }).AddRow(
  198. int64(4),
  199. service.UsageCleanupStatusRunning,
  200. filtersJSON,
  201. int64(7),
  202. int64(0),
  203. nil,
  204. start,
  205. nil,
  206. start,
  207. start,
  208. )
  209. mock.ExpectQuery("UPDATE usage_cleanup_tasks").
  210. WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
  211. WillReturnRows(rows)
  212. task, err := repo.ClaimNextPendingTask(context.Background(), 1800)
  213. require.NoError(t, err)
  214. require.NotNil(t, task)
  215. require.Equal(t, int64(4), task.ID)
  216. require.Equal(t, service.UsageCleanupStatusRunning, task.Status)
  217. require.Equal(t, int64(7), task.CreatedBy)
  218. require.NotNil(t, task.StartedAt)
  219. require.Nil(t, task.ErrorMsg)
  220. require.NoError(t, mock.ExpectationsWereMet())
  221. }
  222. func TestUsageCleanupRepositoryClaimNextPendingTaskError(t *testing.T) {
  223. db, mock := newSQLMock(t)
  224. repo := &usageCleanupRepository{sql: db}
  225. mock.ExpectQuery("UPDATE usage_cleanup_tasks").
  226. WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
  227. WillReturnError(sql.ErrConnDone)
  228. _, err := repo.ClaimNextPendingTask(context.Background(), 1800)
  229. require.Error(t, err)
  230. require.NoError(t, mock.ExpectationsWereMet())
  231. }
  232. func TestUsageCleanupRepositoryClaimNextPendingTaskInvalidFilters(t *testing.T) {
  233. db, mock := newSQLMock(t)
  234. repo := &usageCleanupRepository{sql: db}
  235. rows := sqlmock.NewRows([]string{
  236. "id", "status", "filters", "created_by", "deleted_rows", "error_message",
  237. "started_at", "finished_at", "created_at", "updated_at",
  238. }).AddRow(
  239. int64(4),
  240. service.UsageCleanupStatusRunning,
  241. []byte("invalid"),
  242. int64(7),
  243. int64(0),
  244. nil,
  245. nil,
  246. nil,
  247. time.Now().UTC(),
  248. time.Now().UTC(),
  249. )
  250. mock.ExpectQuery("UPDATE usage_cleanup_tasks").
  251. WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
  252. WillReturnRows(rows)
  253. _, err := repo.ClaimNextPendingTask(context.Background(), 1800)
  254. require.Error(t, err)
  255. require.NoError(t, mock.ExpectationsWereMet())
  256. }
  257. func TestUsageCleanupRepositoryMarkTaskSucceeded(t *testing.T) {
  258. db, mock := newSQLMock(t)
  259. repo := &usageCleanupRepository{sql: db}
  260. mock.ExpectExec("UPDATE usage_cleanup_tasks").
  261. WithArgs(service.UsageCleanupStatusSucceeded, int64(12), int64(9)).
  262. WillReturnResult(sqlmock.NewResult(0, 1))
  263. err := repo.MarkTaskSucceeded(context.Background(), 9, 12)
  264. require.NoError(t, err)
  265. require.NoError(t, mock.ExpectationsWereMet())
  266. }
  267. func TestUsageCleanupRepositoryMarkTaskFailed(t *testing.T) {
  268. db, mock := newSQLMock(t)
  269. repo := &usageCleanupRepository{sql: db}
  270. mock.ExpectExec("UPDATE usage_cleanup_tasks").
  271. WithArgs(service.UsageCleanupStatusFailed, int64(4), "boom", int64(2)).
  272. WillReturnResult(sqlmock.NewResult(0, 1))
  273. err := repo.MarkTaskFailed(context.Background(), 2, 4, "boom")
  274. require.NoError(t, err)
  275. require.NoError(t, mock.ExpectationsWereMet())
  276. }
  277. func TestUsageCleanupRepositoryGetTaskStatus(t *testing.T) {
  278. db, mock := newSQLMock(t)
  279. repo := &usageCleanupRepository{sql: db}
  280. mock.ExpectQuery("SELECT status FROM usage_cleanup_tasks").
  281. WithArgs(int64(9)).
  282. WillReturnRows(sqlmock.NewRows([]string{"status"}).AddRow(service.UsageCleanupStatusPending))
  283. status, err := repo.GetTaskStatus(context.Background(), 9)
  284. require.NoError(t, err)
  285. require.Equal(t, service.UsageCleanupStatusPending, status)
  286. require.NoError(t, mock.ExpectationsWereMet())
  287. }
  288. func TestUsageCleanupRepositoryGetTaskStatusQueryError(t *testing.T) {
  289. db, mock := newSQLMock(t)
  290. repo := &usageCleanupRepository{sql: db}
  291. mock.ExpectQuery("SELECT status FROM usage_cleanup_tasks").
  292. WithArgs(int64(9)).
  293. WillReturnError(sql.ErrConnDone)
  294. _, err := repo.GetTaskStatus(context.Background(), 9)
  295. require.Error(t, err)
  296. require.NoError(t, mock.ExpectationsWereMet())
  297. }
  298. func TestUsageCleanupRepositoryUpdateTaskProgress(t *testing.T) {
  299. db, mock := newSQLMock(t)
  300. repo := &usageCleanupRepository{sql: db}
  301. mock.ExpectExec("UPDATE usage_cleanup_tasks").
  302. WithArgs(int64(123), int64(8)).
  303. WillReturnResult(sqlmock.NewResult(0, 1))
  304. err := repo.UpdateTaskProgress(context.Background(), 8, 123)
  305. require.NoError(t, err)
  306. require.NoError(t, mock.ExpectationsWereMet())
  307. }
  308. func TestUsageCleanupRepositoryCancelTask(t *testing.T) {
  309. db, mock := newSQLMock(t)
  310. repo := &usageCleanupRepository{sql: db}
  311. mock.ExpectQuery("UPDATE usage_cleanup_tasks").
  312. WithArgs(service.UsageCleanupStatusCanceled, int64(6), int64(9), service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning).
  313. WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(6)))
  314. ok, err := repo.CancelTask(context.Background(), 6, 9)
  315. require.NoError(t, err)
  316. require.True(t, ok)
  317. require.NoError(t, mock.ExpectationsWereMet())
  318. }
  319. func TestUsageCleanupRepositoryCancelTaskNoRows(t *testing.T) {
  320. db, mock := newSQLMock(t)
  321. repo := &usageCleanupRepository{sql: db}
  322. mock.ExpectQuery("UPDATE usage_cleanup_tasks").
  323. WithArgs(service.UsageCleanupStatusCanceled, int64(6), int64(9), service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning).
  324. WillReturnRows(sqlmock.NewRows([]string{"id"}))
  325. ok, err := repo.CancelTask(context.Background(), 6, 9)
  326. require.NoError(t, err)
  327. require.False(t, ok)
  328. require.NoError(t, mock.ExpectationsWereMet())
  329. }
  330. func TestUsageCleanupRepositoryDeleteUsageLogsBatchMissingRange(t *testing.T) {
  331. db, _ := newSQLMock(t)
  332. repo := &usageCleanupRepository{sql: db}
  333. _, err := repo.DeleteUsageLogsBatch(context.Background(), service.UsageCleanupFilters{}, 10)
  334. require.Error(t, err)
  335. }
  336. func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
  337. setUsageCleanupRollupTestTimezone(t)
  338. db, mock := newSQLMock(t)
  339. repo := &usageCleanupRepository{sql: db}
  340. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  341. end := start.Add(24 * time.Hour)
  342. userID := int64(3)
  343. model := " gpt-4 "
  344. filters := service.UsageCleanupFilters{
  345. StartTime: start,
  346. EndTime: end,
  347. UserID: &userID,
  348. Model: &model,
  349. }
  350. mock.ExpectBegin()
  351. mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
  352. WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
  353. mock.ExpectQuery("DELETE FROM usage_logs").
  354. WithArgs(start, end, userID, "gpt-4", 2).
  355. WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(start.Add(time.Hour)).AddRow(start.Add(2 * time.Hour)))
  356. mock.ExpectExec(`UPDATE usage_group_rollup_state`).
  357. WithArgs(start.Add(time.Hour), "Asia/Shanghai").
  358. WillReturnResult(sqlmock.NewResult(0, 1))
  359. mock.ExpectCommit()
  360. deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
  361. require.NoError(t, err)
  362. require.Equal(t, int64(2), deleted)
  363. require.NoError(t, mock.ExpectationsWereMet())
  364. }
  365. func TestUsageCleanupRepositoryDeleteUsageLogsBatchAtomicallyInvalidatesGroupRollups(t *testing.T) {
  366. setUsageCleanupRollupTestTimezone(t)
  367. db, mock := newSQLMock(t)
  368. repo := &usageCleanupRepository{sql: db}
  369. start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
  370. end := start.Add(24 * time.Hour)
  371. firstDeletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
  372. secondDeletedAt := firstDeletedAt.Add(time.Hour)
  373. filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
  374. mock.ExpectBegin()
  375. mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
  376. WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
  377. mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
  378. WithArgs(start, end, 2).
  379. WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
  380. AddRow(firstDeletedAt).
  381. AddRow(secondDeletedAt))
  382. mock.ExpectExec(`UPDATE usage_group_rollup_state`).
  383. WithArgs(firstDeletedAt, "Asia/Shanghai").
  384. WillReturnResult(sqlmock.NewResult(0, 1))
  385. mock.ExpectCommit()
  386. deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
  387. require.NoError(t, err)
  388. require.Equal(t, int64(2), deleted)
  389. require.NoError(t, mock.ExpectationsWereMet())
  390. }
  391. func TestUsageCleanupRepositoryDeleteUsageLogsBatchRollsBackWhenInvalidationFails(t *testing.T) {
  392. setUsageCleanupRollupTestTimezone(t)
  393. db, mock := newSQLMock(t)
  394. repo := &usageCleanupRepository{sql: db}
  395. start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
  396. end := start.Add(24 * time.Hour)
  397. deletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
  398. filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
  399. mock.ExpectBegin()
  400. mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
  401. WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
  402. mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
  403. WithArgs(start, end, 1).
  404. WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
  405. mock.ExpectExec(`UPDATE usage_group_rollup_state`).
  406. WithArgs(deletedAt, "Asia/Shanghai").
  407. WillReturnError(sql.ErrConnDone)
  408. mock.ExpectRollback()
  409. _, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 1)
  410. require.ErrorIs(t, err, sql.ErrConnDone)
  411. require.NoError(t, mock.ExpectationsWereMet())
  412. }
  413. func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
  414. db, mock := newSQLMock(t)
  415. repo := &usageCleanupRepository{sql: db}
  416. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  417. end := start.Add(24 * time.Hour)
  418. filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
  419. mock.ExpectBegin()
  420. mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
  421. WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
  422. mock.ExpectQuery("DELETE FROM usage_logs").
  423. WithArgs(start, end, 5).
  424. WillReturnError(sql.ErrConnDone)
  425. mock.ExpectRollback()
  426. _, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 5)
  427. require.Error(t, err)
  428. require.NoError(t, mock.ExpectationsWereMet())
  429. }
  430. func TestBuildUsageCleanupWhere(t *testing.T) {
  431. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  432. end := start.Add(24 * time.Hour)
  433. userID := int64(1)
  434. apiKeyID := int64(2)
  435. accountID := int64(3)
  436. groupID := int64(4)
  437. model := " gpt-4 "
  438. stream := true
  439. billingType := int8(2)
  440. where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
  441. StartTime: start,
  442. EndTime: end,
  443. UserID: &userID,
  444. APIKeyID: &apiKeyID,
  445. AccountID: &accountID,
  446. GroupID: &groupID,
  447. Model: &model,
  448. Stream: &stream,
  449. BillingType: &billingType,
  450. })
  451. require.Equal(t, "created_at >= $1 AND created_at <= $2 AND user_id = $3 AND api_key_id = $4 AND account_id = $5 AND group_id = $6 AND model = $7 AND stream = $8 AND billing_type = $9", where)
  452. require.Equal(t, []any{start, end, userID, apiKeyID, accountID, groupID, "gpt-4", stream, billingType}, args)
  453. }
  454. func TestBuildUsageCleanupWhereRequestTypePriority(t *testing.T) {
  455. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  456. end := start.Add(24 * time.Hour)
  457. requestType := int16(service.RequestTypeWSV2)
  458. stream := false
  459. where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
  460. StartTime: start,
  461. EndTime: end,
  462. RequestType: &requestType,
  463. Stream: &stream,
  464. })
  465. require.Equal(t, "created_at >= $1 AND created_at <= $2 AND (request_type = $3 OR (request_type = 0 AND openai_ws_mode = TRUE))", where)
  466. require.Equal(t, []any{start, end, requestType}, args)
  467. }
  468. func TestBuildUsageCleanupWhereRequestTypeLegacyFallback(t *testing.T) {
  469. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  470. end := start.Add(24 * time.Hour)
  471. requestType := int16(service.RequestTypeStream)
  472. where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
  473. StartTime: start,
  474. EndTime: end,
  475. RequestType: &requestType,
  476. })
  477. require.Equal(t, "created_at >= $1 AND created_at <= $2 AND (request_type = $3 OR (request_type = 0 AND stream = TRUE AND openai_ws_mode = FALSE))", where)
  478. require.Equal(t, []any{start, end, requestType}, args)
  479. }
  480. func TestBuildUsageCleanupWhereModelEmpty(t *testing.T) {
  481. start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
  482. end := start.Add(24 * time.Hour)
  483. model := " "
  484. where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
  485. StartTime: start,
  486. EndTime: end,
  487. Model: &model,
  488. })
  489. require.Equal(t, "created_at >= $1 AND created_at <= $2", where)
  490. require.Equal(t, []any{start, end}, args)
  491. }