// store_test.go:migration + 队列原子抢单语义 + 配额预扣(A6.1 关键路径)。 package store import ( "path/filepath" "sync" "testing" "time" ) // fmtInt 简单 int→string(测试内避免多余导入分支)。 func fmtInt(n int) string { if n == 0 { return "0" } digits := "" for n > 0 { digits = string(rune('0'+n%10)) + digits n /= 10 } return digits } func openTestDB(t *testing.T) *DB { t.Helper() db, err := Open(filepath.Join(t.TempDir(), "test.db")) if err != nil { t.Fatalf("Open: %v", err) } t.Cleanup(func() { db.Close() }) if err := db.Migrate(); err != nil { t.Fatalf("Migrate: %v", err) } return db } // TestMigration 幂等 + 七表存在。 func TestMigration(t *testing.T) { db := openTestDB(t) // 重复 migration 不报错 if err := db.Migrate(); err != nil { t.Fatalf("二次 migration 应幂等: %v", err) } for _, tbl := range []string{"consumers", "api_keys", "quota_ledger", "jobs", "dead_letters", "audit", "rules", "robots_cache"} { var n int if err := db.Raw().QueryRow( `SELECT COUNT(name) FROM sqlite_master WHERE type='table' AND name=?`, tbl).Scan(&n); err != nil { t.Fatal(err) } if n != 1 { t.Errorf("表 %s 不存在", tbl) } } } // TestClaimAtomic 并发抢单:N 个 worker 抢 M 个任务,每任务恰好被抢一次。 func TestClaimAtomic(t *testing.T) { db := openTestDB(t) const N, M = 8, 30 payload := []byte(`{"id":"t1"}`) for i := 0; i < M; i++ { reqID := "req-" + fmtInt(i) if _, err := db.EnqueueJob(reqID, "search", payload, 100); err != nil { t.Fatalf("EnqueueJob: %v", err) } } var mu sync.Mutex claimed := make(map[int64]bool) var wg sync.WaitGroup for w := 0; w < N; w++ { wg.Add(1) go func(id int) { defer wg.Done() for { job, err := db.ClaimNext("worker-"+string(rune('A'+id)), 30*time.Second) if err == ErrNotFound { return } if err != nil { t.Errorf("ClaimNext: %v", err) return } mu.Lock() if claimed[job.ID] { t.Errorf("任务 %d 被重复抢", job.ID) } claimed[job.ID] = true mu.Unlock() } }(w) } wg.Wait() if len(claimed) != M { t.Fatalf("应抢满 %d 任务,实际 %d", M, len(claimed)) } } // TestJobLifecycle 入队 → 抢 → 完成 → 重试退避 → 死信。 func TestJobLifecycle(t *testing.T) { db := openTestDB(t) id, err := db.EnqueueJob("req-lc", "read", []byte(`{"id":"lc"}`), 100) if err != nil { t.Fatal(err) } job, err := db.ClaimNext("w1", time.Minute) if err != nil || job.ID != id { t.Fatalf("抢单失败: %v", err) } if job.Status != "running" || job.Attempts != 1 { t.Fatalf("抢单后状态不符: %+v", job) } // 瞬时错误重试(attempts=1 < 2) retry, err := db.RetryWithBackoff(id, "upstream 5xx", 2, 10*time.Millisecond) if err != nil || !retry { t.Fatalf("应进入重试: ok=%v err=%v", retry, err) } // 退避期内(退避最小 1s)立刻抢应拿不到 if _, err := db.ClaimNext("w2", time.Minute); err != ErrNotFound { t.Fatalf("退避期内不应可抢: %v", err) } time.Sleep(1100 * time.Millisecond) job2, err := db.ClaimNext("w2", time.Minute) if err != nil || job2.ID != id || job2.Attempts != 2 { t.Fatalf("退避后应可重抢: %v", err) } // 第二次失败 → 死信 retry, err = db.RetryWithBackoff(id, "still failing", 2, 10*time.Millisecond) if err != nil || retry { t.Fatalf("attempts 尽应入死信: retry=%v err=%v", retry, err) } running, _ := db.CountRunning() queued, _ := db.CountQueued() if running != 0 || queued != 0 { t.Fatalf("终态后 running/queued 应为 0: %d/%d", running, queued) } } // TestReaper 租约过期收割。 func TestReaper(t *testing.T) { db := openTestDB(t) id, _ := db.EnqueueJob("req-reap", "search", []byte(`{}`), 100) if _, err := db.ClaimNext("w-died", 30*time.Millisecond); err != nil { t.Fatal(err) } time.Sleep(1100 * time.Millisecond) // 退避最小 1s,等 avail 过期 n, err := db.ReapExpired(3) if err != nil || n != 1 { t.Fatalf("reaper 应收割 1 任务: n=%d err=%v", n, err) } var status string if err := db.Raw().QueryRow(`SELECT status FROM jobs WHERE id=?`, id).Scan(&status); err != nil { t.Fatal(err) } if status != "queued" { t.Fatalf("收割后应回队列: %s", status) } } // TestQuotaReserve 预扣防超卖:limit=3 时第 4 次预扣失败。 func TestQuotaReserve(t *testing.T) { db := openTestDB(t) cid, _ := db.CreateConsumer("c1", "") kid, err := db.CreateKey(&ApiKey{ConsumerID: cid, Prefix: "bs_test", Hash: "h1", Salt: "s1", Scopes: []string{"search"}, RPM: 10, DailyQuota: 3, MonthlyQuota: 100, ConcurrentSessions: 1, Status: "enabled"}) if err != nil { t.Fatal(err) } now := time.Now() for i := 0; i < 3; i++ { if _, err := db.QuotaReserve(kid, 3, now); err != nil { t.Fatalf("第 %d 次预扣不应失败: %v", i+1, err) } } if _, err := db.QuotaReserve(kid, 3, now); err != ErrNotFound { t.Fatalf("超限预扣应 ErrNotFound: %v", err) } if err := db.QuotaSettle(kid, now); err != nil { t.Fatal(err) } used, _ := db.QuotaUsed(kid, now) if used != 1 { t.Fatalf("结算后 used 应为 1: %d", used) } if err := db.QuotaRelease(kid, now); err != nil { t.Fatal(err) } } // TestRulesUpsertDelete 规则表 CRUD。 func TestRulesUpsertDelete(t *testing.T) { db := openTestDB(t) if err := db.RuleUpsert("suffix", ".example.com", "direct", 10); err != nil { t.Fatal(err) } if err := db.RuleUpsert("suffix", ".example.com", "deny", 5); err != nil { t.Fatal(err) } rules, err := db.RulesAll() if err != nil || len(rules) != 1 { t.Fatalf("upsert 应去重: n=%d err=%v", len(rules), err) } if rules[0].Action != "deny" || rules[0].Sort != 5 { t.Fatalf("upsert 未更新: %+v", rules[0]) } if err := db.RuleDelete("suffix", ".example.com"); err != nil { t.Fatal(err) } if err := db.RuleDelete("suffix", ".example.com"); err != ErrNotFound { t.Fatalf("删不存在规则应 ErrNotFound: %v", err) } }