package service import ( "context" "errors" "testing" "time" "github.com/Wei-Shaw/sub2api/internal/config" ) // --------------------------------------------------------------------------- // Mock: quotaDirtyCache // --------------------------------------------------------------------------- type mockQuotaDirtyCache struct { // popSequence: 第 0 次 Pop 返回 popSequence[0],之后返回 nil(空集) popSequence [][]UserPlatformQuotaKey popCallIdx int // getEntries: BatchGetUserPlatformQuotaCache 返回的 entries(与 keys 对齐) getEntries []*UserPlatformQuotaCacheEntry getErr error // readdCalled: 记录 Readd 收到的 keys(累积所有次调用) readdCalled [][]UserPlatformQuotaKey readdErr error } func (m *mockQuotaDirtyCache) PopDirtyUserPlatformQuotaKeys(_ context.Context, _ int) ([]UserPlatformQuotaKey, error) { if m.popCallIdx < len(m.popSequence) { keys := m.popSequence[m.popCallIdx] m.popCallIdx++ return keys, nil } // 超出序列 → 空集(模拟脏集已清空) return nil, nil } func (m *mockQuotaDirtyCache) ReaddDirtyUserPlatformQuotaKeys(_ context.Context, keys []UserPlatformQuotaKey) error { m.readdCalled = append(m.readdCalled, keys) return m.readdErr } func (m *mockQuotaDirtyCache) BatchGetUserPlatformQuotaCache(_ context.Context, _ []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { if m.getErr != nil { return nil, m.getErr } return m.getEntries, nil } // --------------------------------------------------------------------------- // Mock: quotaSnapshotWriter // --------------------------------------------------------------------------- type mockQuotaSnapshotWriter struct { receivedSnaps []UserPlatformQuotaSnapshot returnErr error } func (m *mockQuotaSnapshotWriter) BatchSnapshotUsage(_ context.Context, snaps []UserPlatformQuotaSnapshot, _ time.Time) error { m.receivedSnaps = append(m.receivedSnaps, snaps...) return m.returnErr } // --------------------------------------------------------------------------- // Helper: 构造窗口起始时间(非 nil) // --------------------------------------------------------------------------- func flusherPtrTime(t time.Time) *time.Time { return &t } func makeEntry(daily, weekly, monthly float64) *UserPlatformQuotaCacheEntry { now := time.Now().UTC() return &UserPlatformQuotaCacheEntry{ DailyUsageUSD: daily, WeeklyUsageUSD: weekly, MonthlyUsageUSD: monthly, DailyWindowStart: flusherPtrTime(now), WeeklyWindowStart: flusherPtrTime(now), MonthlyWindowStart: flusherPtrTime(now), } } // --------------------------------------------------------------------------- // newTestFlusher: 直接构造 struct(跳过构造函数,B7 才注入) // --------------------------------------------------------------------------- func newTestFlusher(cache quotaDirtyCache, writer quotaSnapshotWriter) *UserPlatformQuotaUsageFlusher { return &UserPlatformQuotaUsageFlusher{ cache: cache, quotaRepo: writer, timingWheel: nil, // 单测不启动 TimingWheel interval: 5 * time.Second, batchSize: 100, flushTimeout: 5 * time.Second, metrics: &FlusherMetrics{}, } } // --------------------------------------------------------------------------- // 场景 1: PopSnapshotUpsert — 2 key + 2 个含 window 的 entry → writer 收 2 行 // --------------------------------------------------------------------------- func TestFlusher_PopSnapshotUpsert(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 1, Platform: "anthropic"}, {UserID: 2, Platform: "openai"}, } cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, // 第 1 次返回 keys,之后空 getEntries: []*UserPlatformQuotaCacheEntry{ makeEntry(1.0, 2.0, 3.0), makeEntry(4.0, 5.0, 6.0), }, } writer := &mockQuotaSnapshotWriter{} f := newTestFlusher(cache, writer) f.flush() if len(writer.receivedSnaps) != 2 { t.Fatalf("expected 2 snaps, got %d", len(writer.receivedSnaps)) } if f.metrics.FlushBatchSizeTotal.Load() != 2 { t.Errorf("FlushBatchSizeTotal = %d, want 2", f.metrics.FlushBatchSizeTotal.Load()) } if f.metrics.FlushSuccessTotal.Load() != 1 { t.Errorf("FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) } if f.metrics.FlushErrorTotal.Load() != 0 { t.Errorf("FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) } } // --------------------------------------------------------------------------- // 场景 2: MissKeySkipped — 2 key,BatchGet 返回 [entry, nil] → 只刷 1 行,nil 跳过,不 Readd // --------------------------------------------------------------------------- func TestFlusher_MissKeySkipped(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 1, Platform: "anthropic"}, {UserID: 2, Platform: "openai"}, } cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, getEntries: []*UserPlatformQuotaCacheEntry{ makeEntry(1.0, 2.0, 3.0), nil, // MISS }, } writer := &mockQuotaSnapshotWriter{} f := newTestFlusher(cache, writer) f.flush() if len(writer.receivedSnaps) != 1 { t.Fatalf("expected 1 snap, got %d", len(writer.receivedSnaps)) } if writer.receivedSnaps[0].UserID != 1 { t.Errorf("expected snap for UserID=1, got %d", writer.receivedSnaps[0].UserID) } if len(cache.readdCalled) != 0 { t.Errorf("Readd should NOT be called on MISS, got %d calls", len(cache.readdCalled)) } if f.metrics.FlushSuccessTotal.Load() != 1 { t.Errorf("FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) } } // --------------------------------------------------------------------------- // 场景 3: UpsertFailReadds — writer 返普通 error → keys 被 Readd,FlushErrorTotal=1,DirtyReaddTotal=len // --------------------------------------------------------------------------- func TestFlusher_UpsertFailReadds(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 1, Platform: "anthropic"}, {UserID: 2, Platform: "openai"}, } cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, getEntries: []*UserPlatformQuotaCacheEntry{ makeEntry(1.0, 2.0, 3.0), makeEntry(4.0, 5.0, 6.0), }, } writeErr := errors.New("db connection timeout") writer := &mockQuotaSnapshotWriter{returnErr: writeErr} f := newTestFlusher(cache, writer) f.flush() if f.metrics.FlushErrorTotal.Load() != 1 { t.Errorf("FlushErrorTotal = %d, want 1", f.metrics.FlushErrorTotal.Load()) } if len(cache.readdCalled) == 0 { t.Fatal("Readd should be called after write error") } totalReadd := 0 for _, rk := range cache.readdCalled { totalReadd += len(rk) } if totalReadd != len(keys) { t.Errorf("DirtyReaddTotal (from Readd calls) = %d, want %d", totalReadd, len(keys)) } if f.metrics.DirtyReaddTotal.Load() != int64(len(keys)) { t.Errorf("DirtyReaddTotal metric = %d, want %d", f.metrics.DirtyReaddTotal.Load(), len(keys)) } if f.metrics.FlushSuccessTotal.Load() != 0 { t.Errorf("FlushSuccessTotal = %d, want 0", f.metrics.FlushSuccessTotal.Load()) } } // --------------------------------------------------------------------------- // 场景 4: FKViolationDropsNoReadd — writer 返 ErrUserPlatformQuotaFKViolation → 不 Readd,FlushFKViolationTotal=1 // --------------------------------------------------------------------------- func TestFlusher_FKViolationDropsNoReadd(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 999, Platform: "anthropic"}, } cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, getEntries: []*UserPlatformQuotaCacheEntry{ makeEntry(1.0, 2.0, 3.0), }, } writer := &mockQuotaSnapshotWriter{returnErr: ErrUserPlatformQuotaFKViolation} f := newTestFlusher(cache, writer) f.flush() if f.metrics.FlushFKViolationTotal.Load() != 1 { t.Errorf("FlushFKViolationTotal = %d, want 1", f.metrics.FlushFKViolationTotal.Load()) } if f.metrics.FlushErrorTotal.Load() != 1 { t.Errorf("FlushErrorTotal = %d, want 1", f.metrics.FlushErrorTotal.Load()) } if len(cache.readdCalled) != 0 { t.Errorf("Readd should NOT be called for FK violation (drop), got %d calls", len(cache.readdCalled)) } if f.metrics.DirtyReaddTotal.Load() != 0 { t.Errorf("DirtyReaddTotal = %d, want 0 (FK violation drops)", f.metrics.DirtyReaddTotal.Load()) } } // --------------------------------------------------------------------------- // 场景 5: NilSafe — var f *UserPlatformQuotaUsageFlusher; f.flush(); f.Stop() 不 panic // --------------------------------------------------------------------------- func TestFlusher_NilSafe(t *testing.T) { var f *UserPlatformQuotaUsageFlusher // 下面两行不应 panic f.flush() f.Stop() } // --------------------------------------------------------------------------- // 场景 6: StopPreventsFlush — stopped=true 后 tick() 不调 flush(writer 没收到 snaps) // --------------------------------------------------------------------------- func TestFlusher_StopPreventsFlush(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 1, Platform: "anthropic"}, } cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, getEntries: []*UserPlatformQuotaCacheEntry{ makeEntry(1.0, 2.0, 3.0), }, } writer := &mockQuotaSnapshotWriter{} f := newTestFlusher(cache, writer) // 标记为已停止 f.stopped.Store(true) // tick 应该直接返回,不触发 flush f.tick() if len(writer.receivedSnaps) != 0 { t.Errorf("expected 0 snaps after stop, got %d", len(writer.receivedSnaps)) } if cache.popCallIdx != 0 { t.Errorf("Pop should not be called after stop, popCallIdx = %d", cache.popCallIdx) } } // --------------------------------------------------------------------------- // 场景 B13-1: ZeroPercentCompany — 0% 公司脏集恒空,flusher 空跑无 DB 写 // // 模拟几乎没有用户配置 quota limit 的公司:脏集始终为空(popSequence 为空切片), // Pop 每次返回空集。flush() 应早退,不写 DB、不计成功、不 Readd。 // --------------------------------------------------------------------------- func TestScenario_ZeroPercentCompany(t *testing.T) { cache := &mockQuotaDirtyCache{ // popSequence 为空 → Pop 超出序列 → 始终返回 nil(空集) popSequence: [][]UserPlatformQuotaKey{}, } writer := &mockQuotaSnapshotWriter{} f := newTestFlusher(cache, writer) f.flush() if len(writer.receivedSnaps) != 0 { t.Errorf("0%% company: expected 0 snaps, got %d", len(writer.receivedSnaps)) } if f.metrics.FlushBatchSizeTotal.Load() != 0 { t.Errorf("0%% company: FlushBatchSizeTotal = %d, want 0", f.metrics.FlushBatchSizeTotal.Load()) } if f.metrics.FlushSuccessTotal.Load() != 0 { t.Errorf("0%% company: FlushSuccessTotal = %d, want 0 (empty-set early return)", f.metrics.FlushSuccessTotal.Load()) } if f.metrics.FlushErrorTotal.Load() != 0 { t.Errorf("0%% company: FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) } if len(cache.readdCalled) != 0 { t.Errorf("0%% company: Readd should never be called, got %d calls", len(cache.readdCalled)) } } // --------------------------------------------------------------------------- // P1: IntervalFallback — flush_interval_ms ≤0 时回退 2s;正常值保留 // --------------------------------------------------------------------------- func TestNewUserPlatformQuotaUsageFlusher_IntervalFallback(t *testing.T) { cases := []struct { name string inMs int wantDu time.Duration }{ {"零值回退 2s", 0, 2 * time.Second}, {"负数回退 2s", -100, 2 * time.Second}, {"正常 2000ms 保留", 2000, 2 * time.Second}, {"正常 500ms 保留", 500, 500 * time.Millisecond}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { cfg := &config.Config{} cfg.Database.UserPlatformQuotaFlushIntervalMs = tc.inMs f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) if f.interval != tc.wantDu { t.Fatalf("interval = %v, want %v", f.interval, tc.wantDu) } }) } } // --------------------------------------------------------------------------- // P1: EnabledField — flusher_enabled 配置正确写入 f.enabled // --------------------------------------------------------------------------- func TestNewUserPlatformQuotaUsageFlusher_EnabledField(t *testing.T) { for _, enabled := range []bool{true, false} { cfg := &config.Config{} cfg.Database.UserPlatformQuotaFlusherEnabled = enabled cfg.Database.UserPlatformQuotaFlushIntervalMs = 500 f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) if f.enabled != enabled { t.Errorf("enabled = %v, want %v", f.enabled, enabled) } } } // --------------------------------------------------------------------------- // P2: ReaddFailCounts — BatchGet 失败 + Readd 失败 → DirtyLostTotal 增、DirtyReaddTotal 不变 // BatchGet 失败 + Readd 成功 → DirtyReaddTotal 增、DirtyLostTotal 不变 // --------------------------------------------------------------------------- func TestFlusher_ReaddFailCounts(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 10, Platform: "anthropic"}, {UserID: 11, Platform: "openai"}, } t.Run("Readd 失败计 DirtyLostTotal", func(t *testing.T) { cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, getErr: errors.New("redis timeout"), // 触发 BatchGet 失败路径 readdErr: errors.New("redis connection refused"), // Readd 也失败 } f := newTestFlusher(cache, &mockQuotaSnapshotWriter{}) f.flush() if f.metrics.DirtyLostTotal.Load() != int64(len(keys)) { t.Errorf("DirtyLostTotal = %d, want %d", f.metrics.DirtyLostTotal.Load(), len(keys)) } if f.metrics.DirtyReaddTotal.Load() != 0 { t.Errorf("DirtyReaddTotal = %d, want 0 (Readd 失败不应计入)", f.metrics.DirtyReaddTotal.Load()) } }) t.Run("Readd 成功计 DirtyReaddTotal", func(t *testing.T) { cache := &mockQuotaDirtyCache{ popSequence: [][]UserPlatformQuotaKey{keys}, getErr: errors.New("redis timeout"), // 触发 BatchGet 失败路径 readdErr: nil, // Readd 成功 } f := newTestFlusher(cache, &mockQuotaSnapshotWriter{}) f.flush() if f.metrics.DirtyReaddTotal.Load() != int64(len(keys)) { t.Errorf("DirtyReaddTotal = %d, want %d", f.metrics.DirtyReaddTotal.Load(), len(keys)) } if f.metrics.DirtyLostTotal.Load() != 0 { t.Errorf("DirtyLostTotal = %d, want 0 (Readd 成功不应计 lost)", f.metrics.DirtyLostTotal.Load()) } }) } // --------------------------------------------------------------------------- // ClampsBatchSize — NewUserPlatformQuotaUsageFlusher 构造时按 // [defaultFlushBatchSize, maxFlushBatchSize] 区间 clamp batchSize // --------------------------------------------------------------------------- func TestNewUserPlatformQuotaUsageFlusher_ClampsBatchSize(t *testing.T) { cases := []struct { name string in int want int }{ {"超上限被 clamp", 7000, maxFlushBatchSize}, {"恰好上限保留", maxFlushBatchSize, maxFlushBatchSize}, {"零回退默认", 0, defaultFlushBatchSize}, {"负数回退默认", -5, defaultFlushBatchSize}, {"正常值保留", 500, 500}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { cfg := &config.Config{} cfg.Database.UserPlatformQuotaFlushBatchSize = tc.in f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) if f.batchSize != tc.want { t.Fatalf("batchSize = %d, want %d", f.batchSize, tc.want) } }) } } // --------------------------------------------------------------------------- // 场景 B13-2: NinetyPercentCompany — 90% 公司大量用户配 limit,一批 5 key 批量刷库 // // 模拟大量用户配置了 quota limit 的公司:脏集第一次 Pop 返回 5 个不同用户的 key, // 之后返回空集(避免 flush 循环)。flush() 应构造 5 条 snapshot 写入 DB, // 断言绝对值语义(snap 的 DailyUsageUSD 等于 entry 的值)、metrics 正确、不 Readd。 // --------------------------------------------------------------------------- func TestScenario_NinetyPercentCompany(t *testing.T) { keys := []UserPlatformQuotaKey{ {UserID: 101, Platform: "anthropic"}, {UserID: 102, Platform: "anthropic"}, {UserID: 103, Platform: "openai"}, {UserID: 104, Platform: "openai"}, {UserID: 105, Platform: "anthropic"}, } entries := []*UserPlatformQuotaCacheEntry{ makeEntry(1.1, 2.2, 3.3), makeEntry(4.4, 5.5, 6.6), makeEntry(7.7, 8.8, 9.9), makeEntry(0.5, 1.0, 1.5), makeEntry(10.0, 20.0, 30.0), } cache := &mockQuotaDirtyCache{ // 第 1 次 Pop 返回 5 keys,之后返回空集(防止 flush 无限循环) popSequence: [][]UserPlatformQuotaKey{keys}, getEntries: entries, } writer := &mockQuotaSnapshotWriter{} f := newTestFlusher(cache, writer) f.flush() // 应收到 5 条 snapshot if len(writer.receivedSnaps) != 5 { t.Fatalf("90%% company: expected 5 snaps, got %d", len(writer.receivedSnaps)) } // 验证绝对值语义:第 1 条 snap 的各窗口 usage 应等于 entries[0] 的值 snap0 := writer.receivedSnaps[0] entry0 := entries[0] if snap0.DailyUsageUSD != entry0.DailyUsageUSD { t.Errorf("snap[0].DailyUsageUSD = %v, want %v", snap0.DailyUsageUSD, entry0.DailyUsageUSD) } if snap0.WeeklyUsageUSD != entry0.WeeklyUsageUSD { t.Errorf("snap[0].WeeklyUsageUSD = %v, want %v", snap0.WeeklyUsageUSD, entry0.WeeklyUsageUSD) } if snap0.MonthlyUsageUSD != entry0.MonthlyUsageUSD { t.Errorf("snap[0].MonthlyUsageUSD = %v, want %v", snap0.MonthlyUsageUSD, entry0.MonthlyUsageUSD) } // FlushBatchSizeTotal 应为 5(本批 keys 数量) if f.metrics.FlushBatchSizeTotal.Load() != 5 { t.Errorf("90%% company: FlushBatchSizeTotal = %d, want 5", f.metrics.FlushBatchSizeTotal.Load()) } // FlushSuccessTotal 应为 1(1 个批次写成功) if f.metrics.FlushSuccessTotal.Load() != 1 { t.Errorf("90%% company: FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) } // 无错误、无 Readd if f.metrics.FlushErrorTotal.Load() != 0 { t.Errorf("90%% company: FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) } if f.metrics.DirtyReaddTotal.Load() != 0 { t.Errorf("90%% company: DirtyReaddTotal = %d, want 0", f.metrics.DirtyReaddTotal.Load()) } if len(cache.readdCalled) != 0 { t.Errorf("90%% company: Readd should not be called, got %d calls", len(cache.readdCalled)) } }