vojo/apps/ai-bot/store_test.go

package main

import (
	"sync"
	"sync/atomic"
	"testing"
)

// These tests exercise the Postgres-backed store directly. They run only when
// AI_BOT_TEST_DATABASE_URL points at a throwaway database (openTestStore skips
// otherwise) and start from a clean slate (openTestStore truncates).

func TestStoreTxnDedup(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()

	if got, err := st.HasTxn("txn-1"); err != nil || got {
		t.Fatalf("fresh txn: got (%v,%v), want (false,nil)", got, err)
	}
	if err := st.MarkTxn("txn-1"); err != nil {
		t.Fatalf("mark: %v", err)
	}
	if got, err := st.HasTxn("txn-1"); err != nil || !got {
		t.Fatalf("marked txn: got (%v,%v), want (true,nil)", got, err)
	}
	// Re-marking is idempotent (a retried transaction).
	if err := st.MarkTxn("txn-1"); err != nil {
		t.Fatalf("re-mark: %v", err)
	}
	if got, _ := st.HasTxn("txn-2"); got {
		t.Fatalf("unrelated txn must be unseen")
	}
}

func TestStoreSeenEvent(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()

	first, err := st.SeenEvent("$ev1")
	if err != nil || !first {
		t.Fatalf("first SeenEvent: got (%v,%v), want (true,nil)", first, err)
	}
	again, err := st.SeenEvent("$ev1")
	if err != nil || again {
		t.Fatalf("repeat SeenEvent: got (%v,%v), want (false,nil)", again, err)
	}
	other, err := st.SeenEvent("$ev2")
	if err != nil || !other {
		t.Fatalf("new SeenEvent: got (%v,%v), want (true,nil)", other, err)
	}
}

// Dedup state must survive a process restart — the whole point of the durable store.
func TestStoreDedupSurvivesRestart(t *testing.T) {
	st := openTestStore(t)
	if _, err := st.SeenEvent("$ev-restart"); err != nil {
		t.Fatalf("seen: %v", err)
	}
	if err := st.MarkTxn("txn-restart"); err != nil {
		t.Fatalf("mark: %v", err)
	}
	st.Close()

	// Reopen the same database WITHOUT truncating: simulates a container restart.
	st2, err := OpenStore(testDSN())
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	defer st2.Close()

	if isNew, err := st2.SeenEvent("$ev-restart"); err != nil || isNew {
		t.Fatalf("event after restart must be already-seen: got (%v,%v)", isNew, err)
	}
	if seen, err := st2.HasTxn("txn-restart"); err != nil || !seen {
		t.Fatalf("txn after restart must be seen: got (%v,%v)", seen, err)
	}
}

func TestStoreLimiterPerUserCap(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()

	const user = "@u:vojo.chat"
	const cap, ceiling = 2, 100.0

	for i := 0; i < cap; i++ {
		if res, err := st.Reserve(user, cap, ceiling); err != nil || res != reserveOK {
			t.Fatalf("reserve %d: got (%v,%v), want reserveOK", i, res, err)
		}
	}
	// The (cap+1)th request is denied per-user.
	if res, err := st.Reserve(user, cap, ceiling); err != nil || res != reserveDeniedUser {
		t.Fatalf("over-cap reserve: got (%v,%v), want reserveDeniedUser", res, err)
	}
	// A different user is unaffected.
	if res, err := st.Reserve("@v:vojo.chat", cap, ceiling); err != nil || res != reserveOK {
		t.Fatalf("other user reserve: got (%v,%v), want reserveOK", res, err)
	}
	// Refund returns a slot, so the first user can reserve once more.
	if err := st.RefundRequest(user); err != nil {
		t.Fatalf("refund: %v", err)
	}
	if res, err := st.Reserve(user, cap, ceiling); err != nil || res != reserveOK {
		t.Fatalf("post-refund reserve: got (%v,%v), want reserveOK", res, err)
	}
}

// A zero per-user cap denies even the first request — the SQLite store's
// requests(0) >= cap(0) behaviour, preserved.
func TestStoreLimiterZeroCap(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()
	if res, err := st.Reserve("@u:vojo.chat", 0, 100.0); err != nil || res != reserveDeniedUser {
		t.Fatalf("zero-cap reserve: got (%v,%v), want reserveDeniedUser", res, err)
	}
}

// A zero ceiling denies the very first request of the day even before any spend row
// exists — the SQLite store treated SUM(NULL) as 0.0 (0 >= 0), and the PG store must
// match (SUM over zero rows is NULL).
func TestStoreLimiterZeroCeiling(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()
	if res, err := st.Reserve("@u:vojo.chat", 1_000_000, 0); err != nil || res != reserveDeniedGlobal {
		t.Fatalf("zero-ceiling reserve on empty store: got (%v,%v), want reserveDeniedGlobal", res, err)
	}
}

func TestStoreLimiterGlobalCeiling(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()

	const ceiling = 1.0
	// Book spend up to the ceiling (Reconcile is what feeds the global gate).
	if err := st.Reconcile("@a:vojo.chat", 0.6); err != nil {
		t.Fatalf("reconcile a: %v", err)
	}
	if err := st.Reconcile("@b:vojo.chat", 0.5); err != nil {
		t.Fatalf("reconcile b: %v", err)
	}
	if spent, err := st.SpentTodayUSD(); err != nil || spent < 1.1 {
		t.Fatalf("spent today: got (%v,%v), want >= 1.1", spent, err)
	}
	// Now any reservation is denied globally, regardless of the per-user cap.
	if res, err := st.Reserve("@c:vojo.chat", 1_000_000, ceiling); err != nil || res != reserveDeniedGlobal {
		t.Fatalf("over-ceiling reserve: got (%v,%v), want reserveDeniedGlobal", res, err)
	}
}

// The pgx pool is concurrent (the SQLite store serialized on one connection). The
// advisory lock in Reserve must still admit EXACTLY perUserCap requests when many
// arrive at once for the same user — the same user messaging from several rooms
// simultaneously must not slip past the cap.
func TestStoreReserveConcurrentRespectsCap(t *testing.T) {
	st := openTestStore(t)
	defer st.Close()

	const user = "@race:vojo.chat"
	const cap = 10
	const goroutines = 50

	var ok int64
	var wg sync.WaitGroup
	for i := 0; i < goroutines; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			res, err := st.Reserve(user, cap, 1e9)
			if err != nil {
				t.Errorf("reserve: %v", err)
				return
			}
			if res == reserveOK {
				atomic.AddInt64(&ok, 1)
			}
		}()
	}
	wg.Wait()
	if ok != cap {
		t.Fatalf("concurrent reserves admitted %d, want exactly %d (the per-user cap)", ok, cap)
	}
}

func TestStoreWarnedEncrypted(t *testing.T) {
	st := openTestStore(t)
	const room = "!enc:vojo.chat"
	if warned, err := st.HasWarnedEncrypted(room); err != nil || warned {
		t.Fatalf("fresh room: got (%v,%v), want (false,nil)", warned, err)
	}
	if err := st.SetWarnedEncrypted(room); err != nil {
		t.Fatalf("set: %v", err)
	}
	// Setting twice is idempotent.
	if err := st.SetWarnedEncrypted(room); err != nil {
		t.Fatalf("re-set: %v", err)
	}
	if warned, err := st.HasWarnedEncrypted(room); err != nil || !warned {
		t.Fatalf("warned room: got (%v,%v), want (true,nil)", warned, err)
	}
	st.Close()

	// The one-shot flag must outlive a restart (F5: no re-react after restart).
	st2, err := OpenStore(testDSN())
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	defer st2.Close()
	if warned, err := st2.HasWarnedEncrypted(room); err != nil || !warned {
		t.Fatalf("warned after restart: got (%v,%v), want (true,nil)", warned, err)
	}
}