chore: directory migration - gateway→server, web→client/electron

server/internal/repo/budget_test.go

@@ -0,0 +1,165 @@
+package repo
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"sync"
+	"testing"
+
+	"github.com/google/uuid"
+)
+
+// setupBudgetDB 复用 quota_test.go 的临时库，再建一个会话。
+// 用真实 SQLite 而非 mock：预算的正确性核心是「判断与自增在同一条 UPDATE 里」，
+// 那正是只有真实数据库才能验证的部分。
+func setupBudgetDB(t *testing.T) uuid.UUID {
+	t.Helper()
+	setupTestDB(t)
+	id, err := CreateSession(context.Background(), nil, "bot", "预算测试", "")
+	if err != nil {
+		t.Fatalf("create session: %v", err)
+	}
+	return id
+}
+
+func TestSessionBudgetZeroMeansUnlimited(t *testing.T) {
+	id := setupBudgetDB(t)
+	ctx := context.Background()
+
+	// 默认 0 = 不限：引入预算不该把已在进行的会话卡死
+	b, err := GetSessionBudget(ctx, id)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !b.Unlimited || b.Remaining != -1 {
+		t.Fatalf("默认应为不限：%+v", b)
+	}
+	// 不限时反复占用都成功
+	for i := 0; i < 5; i++ {
+		if _, err := ConsumeSessionBudget(ctx, id); err != nil {
+			t.Fatalf("不限额下第 %d 次占用失败: %v", i+1, err)
+		}
+	}
+}
+
+func TestSessionBudgetExhausts(t *testing.T) {
+	id := setupBudgetDB(t)
+	ctx := context.Background()
+
+	if _, err := SetSessionBudget(ctx, id, 2); err != nil {
+		t.Fatal(err)
+	}
+	for i := 1; i <= 2; i++ {
+		b, err := ConsumeSessionBudget(ctx, id)
+		if err != nil {
+			t.Fatalf("第 %d 次应成功: %v", i, err)
+		}
+		if b.Used != i {
+			t.Fatalf("第 %d 次后 used = %d", i, b.Used)
+		}
+	}
+	b, err := ConsumeSessionBudget(ctx, id)
+	if !errors.Is(err, ErrSessionBudgetExhausted) {
+		t.Fatalf("第 3 次应耗尽，得到 err=%v b=%+v", err, b)
+	}
+	if b.Remaining != 0 {
+		t.Fatalf("耗尽后剩余应为 0：%+v", b)
+	}
+}
+
+// 判断与自增必须在同一条 UPDATE 里，否则并发下会把预算刷穿
+func TestSessionBudgetConcurrentDoesNotOverdraw(t *testing.T) {
+	id := setupBudgetDB(t)
+	ctx := context.Background()
+	const limit = 10
+	if _, err := SetSessionBudget(ctx, id, limit); err != nil {
+		t.Fatal(err)
+	}
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	ok := 0
+	for i := 0; i < 40; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if _, err := ConsumeSessionBudget(ctx, id); err == nil {
+				mu.Lock()
+				ok++
+				mu.Unlock()
+			}
+		}()
+	}
+	wg.Wait()
+
+	if ok != limit {
+		t.Fatalf("40 并发下成功 %d 次，期望恰好 %d 次（预算被刷穿或误拒）", ok, limit)
+	}
+	b, _ := GetSessionBudget(ctx, id)
+	if b.Used != limit {
+		t.Fatalf("used = %d，期望 %d", b.Used, limit)
+	}
+}
+
+func TestSessionBudgetResetAndLowerBelowUsed(t *testing.T) {
+	id := setupBudgetDB(t)
+	ctx := context.Background()
+
+	SetSessionBudget(ctx, id, 5)
+	for i := 0; i < 3; i++ {
+		ConsumeSessionBudget(ctx, id)
+	}
+
+	// 调到低于已用次数 = 「就到这里为止」，是人的合法意图，不该报错
+	b, err := SetSessionBudget(ctx, id, 1)
+	if err != nil {
+		t.Fatalf("下调预算不该失败: %v", err)
+	}
+	if b.Remaining != 0 {
+		t.Fatalf("已用 3 上限 1 时剩余应为 0：%+v", b)
+	}
+	if _, err := ConsumeSessionBudget(ctx, id); !errors.Is(err, ErrSessionBudgetExhausted) {
+		t.Fatal("下调后应立即拦住")
+	}
+
+	// 重置只清已用次数，不动上限
+	b, err = ResetSessionBudget(ctx, id)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if b.Used != 0 || b.Max != 1 {
+		t.Fatalf("重置后应为 0/1：%+v", b)
+	}
+}
+
+// 会话预算先扣、全局配额后扣；全局拦下时必须把会话那次退回去，
+// 否则那格白掉了 —— 那次往返实际上没有发生
+func TestRefundSessionBudget(t *testing.T) {
+	id := setupBudgetDB(t)
+	ctx := context.Background()
+	SetSessionBudget(ctx, id, 3)
+	ConsumeSessionBudget(ctx, id)
+
+	RefundSessionBudget(ctx, id)
+	b, _ := GetSessionBudget(ctx, id)
+	if b.Used != 0 {
+		t.Fatalf("退还后 used 应为 0：%+v", b)
+	}
+
+	// 已经是 0 时再退不该变成负数
+	RefundSessionBudget(ctx, id)
+	b, _ = GetSessionBudget(ctx, id)
+	if b.Used != 0 {
+		t.Fatalf("重复退还把 used 变成了 %d", b.Used)
+	}
+}
+
+func TestGetSessionBudgetMissingSession(t *testing.T) {
+	setupBudgetDB(t)
+	if _, err := GetSessionBudget(context.Background(), uuid.New()); err == nil {
+		t.Fatal("不存在的会话应报错")
+	} else if errors.Is(err, sql.ErrNoRows) {
+		t.Fatal("应包装成可读错误而不是裸 sql.ErrNoRows")
+	}
+}