320 lines
13 KiB
Go
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")
|
||
|
|
}
|