Agent-System-Backend/internal/store/admission_pg_test.go

package store

import (
	"context"
	"errors"
	"sync"
	"testing"
)

// 协作存储 U3 切片 2（R8）：执行许可账本。预算判定原子性、让出/复活、
// 并发恰 N 个赢家。双档（PG 由 CI 真库 job 覆盖咨询锁与部分索引）。

func TestExecutionBudgetEnforced(t *testing.T) {
	st := openTestStore(t)
	budget := ExecutionBudget{Global: 2}
	if err := st.TryAcquireExecutionSlot("run_1", "o", budget); err != nil {
		t.Fatal(err)
	}
	if err := st.TryAcquireExecutionSlot("run_2", "o", budget); err != nil {
		t.Fatal(err)
	}
	err := st.TryAcquireExecutionSlot("run_3", "o", budget)
	if !errors.Is(err, ErrBudgetExceeded) {
		t.Fatalf("global=2 时第三个应拒绝：%v", err)
	}
	// 释放一个 → 第三个可进。
	if err := st.ReleaseExecutionSlot("run_1"); err != nil {
		t.Fatal(err)
	}
	if err := st.TryAcquireExecutionSlot("run_3", "o", budget); err != nil {
		t.Fatal(err)
	}
	// 复活（父让出后重取）：释放 run_2 再复活 run_2。
	if err := st.ReleaseExecutionSlot("run_2"); err != nil {
		t.Fatal(err)
	}
	if err := st.TryAcquireExecutionSlot("run_2", "o", budget); err != nil {
		t.Fatal(err)
	}
	global, byOwner, err := st.ActiveExecutionSlots()
	if err != nil || global != 2 || byOwner["o"] != 2 {
		t.Fatalf("活跃账本：global=%d byOwner=%v err=%v", global, byOwner, err)
	}
	// 预算关闭：零开销直通、不留租约。
	if err := st.TryAcquireExecutionSlot("run_9", "o", ExecutionBudget{}); err != nil {
		t.Fatal(err)
	}
}

func TestExecutionBudgetPerUser(t *testing.T) {
	st := openTestStore(t)
	budget := ExecutionBudget{PerUser: 1}
	if err := st.TryAcquireExecutionSlot("run_a", "alice", budget); err != nil {
		t.Fatal(err)
	}
	if err := st.TryAcquireExecutionSlot("run_b", "bob", budget); err != nil {
		t.Fatal(err)
	}
	err := st.TryAcquireExecutionSlot("run_c", "alice", budget)
	if !errors.Is(err, ErrBudgetExceeded) {
		t.Fatalf("perUser=1 时 alice 第二个应拒绝：%v", err)
	}
	// 同 run 重取不占新额度（复活语义）。
	if err := st.TryAcquireExecutionSlot("run_a", "alice", budget); err != nil {
		t.Fatal(err)
	}
}

// 并发获取恰 N 个赢家：预算 3、12 个竞争者 → 3 成功 9 拒绝（原子清点，
// 不超卖；PG 路径验证咨询锁串行化）。
func TestExecutionBudgetConcurrentNoOversell(t *testing.T) {
	st := openTestStore(t)
	budget := ExecutionBudget{Global: 3}
	const racers = 12
	var mu sync.Mutex
	ok, denied := 0, 0
	var wg sync.WaitGroup
	for i := 0; i < racers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			err := st.TryAcquireExecutionSlot(newID("run"), "o"+string(rune('a'+i%3)), budget)
			mu.Lock()
			defer mu.Unlock()
			switch {
			case err == nil:
				ok++
			case errors.Is(err, ErrBudgetExceeded):
				denied++
			default:
				t.Errorf("意外错误：%v", err)
			}
		}(i)
	}
	wg.Wait()
	if ok != 3 || denied != racers-3 {
		t.Fatalf("应恰 3 成功：ok=%d denied=%d", ok, denied)
	}
}

// 迁移到位验收（v21 表存在）由 TestMigrationUpgradesExistingDatabase 的
// 全量账本断言覆盖；此处补空库→全量后账本可清点的最小回归。
func TestAdmissionLeasesTablePresent(t *testing.T) {
	st := openTestStore(t)
	global, _, err := st.ActiveExecutionSlots()
	if err != nil || global != 0 {
		t.Fatalf("空账本：global=%d err=%v", global, err)
	}
	_ = context.Background()
}