Files
sub2api/backend/internal/repository/dashboard_aggregation_group_usage_test.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

320 lines
13 KiB
Go

//go:build unit
package repository
import (
"context"
"database/sql"
"errors"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func TestDashboardAggregationRepositorySyncGroupUsageRollupsNoopsAtCurrentDate(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-08-14", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectCommit()
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRebuildsWhenTimezoneChanges(t *testing.T) {
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
retainedFrom := time.Date(2026, 3, 1, 5, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-03-09", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
WithArgs("2026-03-01", "2026-03-01", "2026-03-09").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
WithArgs(retainedFrom, todayStart, "America/New_York").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs("2026-03-09", retainedFrom, "America/New_York").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositorySyncGroupUsageRollupsPublishesWatermarkLast(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
retainedFrom := time.Date(2026, 5, 1, 3, 0, 0, 0, time.UTC)
rebuildStart := time.Date(2026, 8, 12, 16, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-08-13", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
WithArgs("2026-05-01", "2026-08-13", "2026-08-14").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
WithArgs(rebuildStart, todayStart, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 2))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs("2026-08-14", retainedFrom, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRejectsFutureWatermark(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-08-15", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectRollback()
err := repo.SyncGroupUsageRollups(context.Background(), todayStart)
require.ErrorContains(t, err, "未来")
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryRecomputeRangeInvalidatesGroupRollupsBeforeDashboardRebuild(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
end := start.Add(time.Hour)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`DELETE FROM usage_dashboard_hourly`).
WillReturnError(sql.ErrConnDone)
mock.ExpectRollback()
err := repo.RecomputeRange(context.Background(), start, end)
require.ErrorIs(t, err, sql.ErrConnDone)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryRecomputeRangeRebuildsGroupRollupsBeforeCommit(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
end := start.Add(time.Hour)
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
repo.clock = func() time.Time { return fixedNow }
todayStart := service.GroupUsageTodayStart(fixedNow)
startDate := service.GroupUsageDate(start)
rebuildStart, err := service.ParseGroupUsageDate(startDate)
require.NoError(t, err)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
for _, query := range []string{
`DELETE FROM usage_dashboard_hourly WHERE`,
`DELETE FROM usage_dashboard_hourly_users WHERE`,
`DELETE FROM usage_dashboard_daily WHERE`,
`DELETE FROM usage_dashboard_daily_users WHERE`,
`INSERT INTO usage_dashboard_hourly_users`,
`INSERT INTO usage_dashboard_daily_users`,
`INSERT INTO usage_dashboard_hourly`,
`INSERT INTO usage_dashboard_daily`,
} {
mock.ExpectExec(query).WillReturnResult(sqlmock.NewResult(0, 1))
}
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow(startDate, time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(start))
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
WithArgs(startDate, startDate, service.GroupUsageDate(todayStart)).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
WithArgs(rebuildStart.UTC(), todayStart.UTC(), "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(service.GroupUsageDate(todayStart), start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
require.NoError(t, repo.RecomputeRange(context.Background(), start, end))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedInvalidatesEachBatchAndSyncs(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
earliestDeletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
repo.clock = func() time.Time { return fixedNow }
todayStart := service.GroupUsageTodayStart(fixedNow)
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
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(cutoff, usageLogsCleanupBatchSize).
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
AddRow(earliestDeletedAt.Add(time.Hour)).
AddRow(earliestDeletedAt))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(earliestDeletedAt, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectCommit()
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionedSortsAndInvalidatesEachDropBeforeSync(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
juneStart := time.Date(2026, 6, 1, 0, 0, 0, 0, time.UTC)
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
repo.clock = func() time.Time { return fixedNow }
todayStart := service.GroupUsageTodayStart(fixedNow)
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
mock.ExpectQuery(`SELECT c.relname`).
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
AddRow("usage_logs_202606").
AddRow("usage_logs_invalid").
AddRow("usage_logs_202604").
AddRow("usage_logs_202607"))
for _, partition := range []struct {
name string
start time.Time
}{
{name: "usage_logs_202604", start: aprilStart},
{name: "usage_logs_202606", start: juneStart},
} {
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(partition.start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`DROP TABLE IF EXISTS "` + partition.name + `"`).
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectCommit()
}
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectCommit()
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedFailureRollsBackWithoutSync(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
deletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery(`(?s)SELECT ctid.*ORDER BY created_at ASC, id ASC.*DELETE FROM usage_logs.*RETURNING created_at`).
WithArgs(cutoff, usageLogsCleanupBatchSize).
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.CleanupUsageLogs(context.Background(), cutoff)
require.ErrorIs(t, err, sql.ErrConnDone)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionFailureRollsBackAndStops(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
dropErr := errors.New("drop partition failed")
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
mock.ExpectQuery(`SELECT c.relname`).
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
AddRow("usage_logs_202606").
AddRow("usage_logs_202604"))
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(aprilStart, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`DROP TABLE IF EXISTS "usage_logs_202604"`).
WillReturnError(dropErr)
mock.ExpectRollback()
err := repo.CleanupUsageLogs(context.Background(), cutoff)
require.ErrorIs(t, err, dropErr)
require.NoError(t, mock.ExpectationsWereMet())
}
func setGroupUsageRollupTestTimezone(t *testing.T) {
t.Helper()
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
}