onesvm-browser-server/server/internal/store/store_test.go

// 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",
		"fp_templates", "fp_state", "session_cookies"} {
		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)
	}
}