MailUI4Agents/gateway/internal/repo/budget_test.go

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")
	}
}