| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587 |
- package repository
- import (
- "context"
- "database/sql"
- "encoding/json"
- "testing"
- "time"
- "github.com/DATA-DOG/go-sqlmock"
- "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/stretchr/testify/require"
- )
- func newSQLMock(t *testing.T) (*sql.DB, sqlmock.Sqlmock) {
- t.Helper()
- db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
- require.NoError(t, err)
- t.Cleanup(func() { _ = db.Close() })
- return db, mock
- }
- func setUsageCleanupRollupTestTimezone(t *testing.T) {
- t.Helper()
- useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
- }
- func TestNewUsageCleanupRepository(t *testing.T) {
- db, _ := newSQLMock(t)
- repo := NewUsageCleanupRepository(nil, db)
- require.NotNil(t, repo)
- }
- func TestUsageCleanupRepositoryCreateTask(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2024, 1, 1, 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: 12,
- }
- now := time.Date(2024, 1, 2, 0, 0, 0, 0, time.UTC)
- mock.ExpectQuery("INSERT INTO usage_cleanup_tasks").
- WithArgs(task.Status, sqlmock.AnyArg(), task.CreatedBy, task.DeletedRows).
- WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(1), now, now))
- err := repo.CreateTask(context.Background(), task)
- require.NoError(t, err)
- require.Equal(t, int64(1), task.ID)
- require.Equal(t, now, task.CreatedAt)
- require.Equal(t, now, task.UpdatedAt)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryCreateTaskNil(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- err := repo.CreateTask(context.Background(), nil)
- require.NoError(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryCreateTaskQueryError(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- task := &service.UsageCleanupTask{
- Status: service.UsageCleanupStatusPending,
- Filters: service.UsageCleanupFilters{StartTime: time.Now(), EndTime: time.Now().Add(time.Hour)},
- CreatedBy: 1,
- }
- mock.ExpectQuery("INSERT INTO usage_cleanup_tasks").
- WithArgs(task.Status, sqlmock.AnyArg(), task.CreatedBy, task.DeletedRows).
- WillReturnError(sql.ErrConnDone)
- err := repo.CreateTask(context.Background(), task)
- require.Error(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryListTasksEmpty(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(0)))
- tasks, result, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
- require.NoError(t, err)
- require.Empty(t, tasks)
- require.Equal(t, int64(0), result.Total)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryListTasks(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(2 * time.Hour)
- filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
- filtersJSON, err := json.Marshal(filters)
- require.NoError(t, err)
- createdAt := time.Date(2024, 1, 2, 12, 0, 0, 0, time.UTC)
- updatedAt := createdAt.Add(time.Minute)
- rows := sqlmock.NewRows([]string{
- "id", "status", "filters", "created_by", "deleted_rows", "error_message",
- "canceled_by", "canceled_at",
- "started_at", "finished_at", "created_at", "updated_at",
- }).AddRow(
- int64(1),
- service.UsageCleanupStatusSucceeded,
- filtersJSON,
- int64(2),
- int64(9),
- "error",
- nil,
- nil,
- start,
- end,
- createdAt,
- updatedAt,
- )
- mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
- WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(1)))
- mock.ExpectQuery("SELECT id, status, filters, created_by, deleted_rows, error_message").
- WithArgs(20, 0).
- WillReturnRows(rows)
- tasks, result, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
- require.NoError(t, err)
- require.Len(t, tasks, 1)
- require.Equal(t, int64(1), tasks[0].ID)
- require.Equal(t, service.UsageCleanupStatusSucceeded, tasks[0].Status)
- require.Equal(t, int64(2), tasks[0].CreatedBy)
- require.Equal(t, int64(9), tasks[0].DeletedRows)
- require.NotNil(t, tasks[0].ErrorMsg)
- require.Equal(t, "error", *tasks[0].ErrorMsg)
- require.NotNil(t, tasks[0].StartedAt)
- require.NotNil(t, tasks[0].FinishedAt)
- require.Equal(t, int64(1), result.Total)
- 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}
- rows := sqlmock.NewRows([]string{
- "id", "status", "filters", "created_by", "deleted_rows", "error_message",
- "canceled_by", "canceled_at",
- "started_at", "finished_at", "created_at", "updated_at",
- }).AddRow(
- int64(1),
- service.UsageCleanupStatusSucceeded,
- []byte("not-json"),
- int64(2),
- int64(9),
- nil,
- nil,
- nil,
- nil,
- nil,
- time.Now().UTC(),
- time.Now().UTC(),
- )
- mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM usage_cleanup_tasks").
- WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(int64(1)))
- mock.ExpectQuery("SELECT id, status, filters, created_by, deleted_rows, error_message").
- WithArgs(20, 0).
- WillReturnRows(rows)
- _, _, err := repo.ListTasks(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20})
- require.Error(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryClaimNextPendingTaskNone(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- mock.ExpectQuery("UPDATE usage_cleanup_tasks").
- WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
- WillReturnRows(sqlmock.NewRows([]string{
- "id", "status", "filters", "created_by", "deleted_rows", "error_message",
- "started_at", "finished_at", "created_at", "updated_at",
- }))
- task, err := repo.ClaimNextPendingTask(context.Background(), 1800)
- require.NoError(t, err)
- require.Nil(t, task)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryClaimNextPendingTask(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
- filtersJSON, err := json.Marshal(filters)
- require.NoError(t, err)
- rows := sqlmock.NewRows([]string{
- "id", "status", "filters", "created_by", "deleted_rows", "error_message",
- "started_at", "finished_at", "created_at", "updated_at",
- }).AddRow(
- int64(4),
- service.UsageCleanupStatusRunning,
- filtersJSON,
- int64(7),
- int64(0),
- nil,
- start,
- nil,
- start,
- start,
- )
- mock.ExpectQuery("UPDATE usage_cleanup_tasks").
- WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
- WillReturnRows(rows)
- task, err := repo.ClaimNextPendingTask(context.Background(), 1800)
- require.NoError(t, err)
- require.NotNil(t, task)
- require.Equal(t, int64(4), task.ID)
- require.Equal(t, service.UsageCleanupStatusRunning, task.Status)
- require.Equal(t, int64(7), task.CreatedBy)
- require.NotNil(t, task.StartedAt)
- require.Nil(t, task.ErrorMsg)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryClaimNextPendingTaskError(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- mock.ExpectQuery("UPDATE usage_cleanup_tasks").
- WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
- WillReturnError(sql.ErrConnDone)
- _, err := repo.ClaimNextPendingTask(context.Background(), 1800)
- require.Error(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryClaimNextPendingTaskInvalidFilters(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- rows := sqlmock.NewRows([]string{
- "id", "status", "filters", "created_by", "deleted_rows", "error_message",
- "started_at", "finished_at", "created_at", "updated_at",
- }).AddRow(
- int64(4),
- service.UsageCleanupStatusRunning,
- []byte("invalid"),
- int64(7),
- int64(0),
- nil,
- nil,
- nil,
- time.Now().UTC(),
- time.Now().UTC(),
- )
- mock.ExpectQuery("UPDATE usage_cleanup_tasks").
- WithArgs(service.UsageCleanupStatusPending, service.UsageCleanupStatusRunning, int64(1800), service.UsageCleanupStatusRunning).
- WillReturnRows(rows)
- _, err := repo.ClaimNextPendingTask(context.Background(), 1800)
- require.Error(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryMarkTaskSucceeded(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- mock.ExpectExec("UPDATE usage_cleanup_tasks").
- WithArgs(service.UsageCleanupStatusSucceeded, int64(12), int64(9)).
- WillReturnResult(sqlmock.NewResult(0, 1))
- err := repo.MarkTaskSucceeded(context.Background(), 9, 12)
- require.NoError(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryMarkTaskFailed(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- mock.ExpectExec("UPDATE usage_cleanup_tasks").
- WithArgs(service.UsageCleanupStatusFailed, int64(4), "boom", int64(2)).
- WillReturnResult(sqlmock.NewResult(0, 1))
- err := repo.MarkTaskFailed(context.Background(), 2, 4, "boom")
- require.NoError(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryGetTaskStatus(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- mock.ExpectQuery("SELECT status FROM usage_cleanup_tasks").
- WithArgs(int64(9)).
- WillReturnRows(sqlmock.NewRows([]string{"status"}).AddRow(service.UsageCleanupStatusPending))
- status, err := repo.GetTaskStatus(context.Background(), 9)
- require.NoError(t, err)
- require.Equal(t, service.UsageCleanupStatusPending, status)
- 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}
- mock.ExpectExec("UPDATE usage_cleanup_tasks").
- WithArgs(int64(123), int64(8)).
- WillReturnResult(sqlmock.NewResult(0, 1))
- err := repo.UpdateTaskProgress(context.Background(), 8, 123)
- require.NoError(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryCancelTask(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"}).AddRow(int64(6)))
- ok, err := repo.CancelTask(context.Background(), 6, 9)
- require.NoError(t, err)
- require.True(t, ok)
- 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}
- _, err := repo.DeleteUsageLogsBatch(context.Background(), service.UsageCleanupFilters{}, 10)
- require.Error(t, err)
- }
- func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
- setUsageCleanupRollupTestTimezone(t)
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- userID := int64(3)
- model := " gpt-4 "
- filters := service.UsageCleanupFilters{
- StartTime: start,
- EndTime: end,
- UserID: &userID,
- Model: &model,
- }
- mock.ExpectBegin()
- mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
- WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
- mock.ExpectQuery("DELETE FROM usage_logs").
- WithArgs(start, end, userID, "gpt-4", 2).
- WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(start.Add(time.Hour)).AddRow(start.Add(2 * time.Hour)))
- mock.ExpectExec(`UPDATE usage_group_rollup_state`).
- WithArgs(start.Add(time.Hour), "Asia/Shanghai").
- WillReturnResult(sqlmock.NewResult(0, 1))
- mock.ExpectCommit()
- deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
- require.NoError(t, err)
- require.Equal(t, int64(2), deleted)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryDeleteUsageLogsBatchAtomicallyInvalidatesGroupRollups(t *testing.T) {
- setUsageCleanupRollupTestTimezone(t)
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- firstDeletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
- secondDeletedAt := firstDeletedAt.Add(time.Hour)
- filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
- mock.ExpectBegin()
- mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
- WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
- mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
- WithArgs(start, end, 2).
- WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
- AddRow(firstDeletedAt).
- AddRow(secondDeletedAt))
- mock.ExpectExec(`UPDATE usage_group_rollup_state`).
- WithArgs(firstDeletedAt, "Asia/Shanghai").
- WillReturnResult(sqlmock.NewResult(0, 1))
- mock.ExpectCommit()
- deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
- require.NoError(t, err)
- require.Equal(t, int64(2), deleted)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryDeleteUsageLogsBatchRollsBackWhenInvalidationFails(t *testing.T) {
- setUsageCleanupRollupTestTimezone(t)
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- deletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
- filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
- mock.ExpectBegin()
- mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
- WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
- mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
- WithArgs(start, end, 1).
- WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
- mock.ExpectExec(`UPDATE usage_group_rollup_state`).
- WithArgs(deletedAt, "Asia/Shanghai").
- WillReturnError(sql.ErrConnDone)
- mock.ExpectRollback()
- _, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 1)
- require.ErrorIs(t, err, sql.ErrConnDone)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
- db, mock := newSQLMock(t)
- repo := &usageCleanupRepository{sql: db}
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
- mock.ExpectBegin()
- mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
- WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
- mock.ExpectQuery("DELETE FROM usage_logs").
- WithArgs(start, end, 5).
- WillReturnError(sql.ErrConnDone)
- mock.ExpectRollback()
- _, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 5)
- require.Error(t, err)
- require.NoError(t, mock.ExpectationsWereMet())
- }
- func TestBuildUsageCleanupWhere(t *testing.T) {
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- userID := int64(1)
- apiKeyID := int64(2)
- accountID := int64(3)
- groupID := int64(4)
- model := " gpt-4 "
- stream := true
- billingType := int8(2)
- where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
- StartTime: start,
- EndTime: end,
- UserID: &userID,
- APIKeyID: &apiKeyID,
- AccountID: &accountID,
- GroupID: &groupID,
- Model: &model,
- Stream: &stream,
- BillingType: &billingType,
- })
- 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)
- require.Equal(t, []any{start, end, userID, apiKeyID, accountID, groupID, "gpt-4", stream, billingType}, args)
- }
- func TestBuildUsageCleanupWhereRequestTypePriority(t *testing.T) {
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- requestType := int16(service.RequestTypeWSV2)
- stream := false
- where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
- StartTime: start,
- EndTime: end,
- RequestType: &requestType,
- Stream: &stream,
- })
- require.Equal(t, "created_at >= $1 AND created_at <= $2 AND (request_type = $3 OR (request_type = 0 AND openai_ws_mode = TRUE))", where)
- require.Equal(t, []any{start, end, requestType}, args)
- }
- func TestBuildUsageCleanupWhereRequestTypeLegacyFallback(t *testing.T) {
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- requestType := int16(service.RequestTypeStream)
- where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
- StartTime: start,
- EndTime: end,
- RequestType: &requestType,
- })
- 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)
- require.Equal(t, []any{start, end, requestType}, args)
- }
- func TestBuildUsageCleanupWhereModelEmpty(t *testing.T) {
- start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
- end := start.Add(24 * time.Hour)
- model := " "
- where, args := buildUsageCleanupWhere(service.UsageCleanupFilters{
- StartTime: start,
- EndTime: end,
- Model: &model,
- })
- require.Equal(t, "created_at >= $1 AND created_at <= $2", where)
- require.Equal(t, []any{start, end}, args)
- }
|