272 lines
8.3 KiB
Go
272 lines
8.3 KiB
Go
package store
|
||
|
||
import (
|
||
"path/filepath"
|
||
"sync"
|
||
"testing"
|
||
)
|
||
|
||
func openTemp(t *testing.T) (*Store, string) {
|
||
t.Helper()
|
||
path := filepath.Join(t.TempDir(), "test.db")
|
||
s, err := Open(path)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() { s.Close() })
|
||
return s, path
|
||
}
|
||
|
||
func mustSnapshotAndJob(t *testing.T, s *Store) *Job {
|
||
t.Helper()
|
||
if err := s.UpsertSnapshot("snap1", "brief1", `{}`); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
job, err := s.EnsureJob("book", 1, "draft", "snap1")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return job
|
||
}
|
||
|
||
func TestReserveSettleLifecycle(t *testing.T) {
|
||
s, _ := openTemp(t)
|
||
job := mustSnapshotAndJob(t, s)
|
||
caps := Ceilings{BookUSD: 1.0, DayUSD: 2.0}
|
||
|
||
res, verdict, err := s.Reserve("book", 0.10, caps)
|
||
if err != nil || verdict != ReserveOK {
|
||
t.Fatalf("reserve: %v %v", verdict, err)
|
||
}
|
||
committed, reserved, _ := s.SpentUSD("book")
|
||
if committed != 0 || reserved != 0.10 {
|
||
t.Fatalf("after reserve: committed=%v reserved=%v", committed, reserved)
|
||
}
|
||
|
||
cp := Checkpoint{RequestHash: "h1", JobID: job.ID, Stage: "draft", Role: "translator",
|
||
ModelRequested: "m", ModelActual: "m", ResponseText: "текст", UsageJSON: "{}", CostUSD: 0.04}
|
||
if err := s.SettleWithCheckpoint(res, 0.04, cp, nil); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
committed, reserved, _ = s.SpentUSD("book")
|
||
if committed != 0.04 || reserved != 0 {
|
||
t.Fatalf("after settle: committed=%v reserved=%v", committed, reserved)
|
||
}
|
||
|
||
got, err := s.GetCheckpoint("h1")
|
||
if err != nil || got == nil || got.ResponseText != "текст" {
|
||
t.Fatalf("checkpoint: %+v err=%v", got, err)
|
||
}
|
||
|
||
// A repeat settle of the same hash (a caller retry) is idempotent:
|
||
// the reservation is released, money is not booked a second time.
|
||
res2, verdict, err := s.Reserve("book", 0.10, caps)
|
||
if err != nil || verdict != ReserveOK {
|
||
t.Fatalf("reserve2: %v %v", verdict, err)
|
||
}
|
||
if err := s.SettleWithCheckpoint(res2, 0.04, cp, nil); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
committed, reserved, _ = s.SpentUSD("book")
|
||
if committed != 0.04 || reserved != 0 {
|
||
t.Fatalf("duplicate settle must book nothing: committed=%v reserved=%v", committed, reserved)
|
||
}
|
||
}
|
||
|
||
// TestPutDerivedCheckpoint pins the $0 sanitized-export checkpoint (D38 infra-pack): it persists a
|
||
// retrievable response, books NO spend, and is idempotent (re-derive on resume is free).
|
||
func TestPutDerivedCheckpoint(t *testing.T) {
|
||
s, _ := openTemp(t)
|
||
job := mustSnapshotAndJob(t, s)
|
||
|
||
cp := Checkpoint{RequestHash: "tm-sanitized-v1:abc", JobID: job.ID, ChunkIdx: 0, Stage: "edit",
|
||
Role: "editor", ModelRequested: "m", ModelActual: "m", ResponseText: "очищенный текст",
|
||
UsageJSON: "{}", FinishReason: "sanitized_export"}
|
||
if err := s.PutDerivedCheckpoint(cp); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// Retrievable via the standard checkpoint read (the final_hash→checkpoint export contract).
|
||
got, err := s.GetCheckpoint("tm-sanitized-v1:abc")
|
||
if err != nil || got == nil || got.ResponseText != "очищенный текст" {
|
||
t.Fatalf("derived checkpoint: %+v err=%v", got, err)
|
||
}
|
||
if got.CostUSD != 0 {
|
||
t.Fatalf("derived checkpoint must be $0, got %v", got.CostUSD)
|
||
}
|
||
// No spend booked (it is a post-processing artifact, not a provider call).
|
||
committed, reserved, _ := s.SpentUSD("book")
|
||
if committed != 0 || reserved != 0 {
|
||
t.Fatalf("derived checkpoint booked spend: committed=%v reserved=%v", committed, reserved)
|
||
}
|
||
// Idempotent: a resume re-derives the identical row for free (ON CONFLICT DO NOTHING).
|
||
if err := s.PutDerivedCheckpoint(cp); err != nil {
|
||
t.Fatalf("re-put must be idempotent: %v", err)
|
||
}
|
||
committed, _, _ = s.SpentUSD("book")
|
||
if committed != 0 {
|
||
t.Fatalf("idempotent re-put booked spend: committed=%v", committed)
|
||
}
|
||
}
|
||
|
||
func TestCeilingsDeny(t *testing.T) {
|
||
s, _ := openTemp(t)
|
||
mustSnapshotAndJob(t, s)
|
||
|
||
// Book ceiling: committed+reserved+estimate > cap ⇒ deny.
|
||
res, verdict, err := s.Reserve("book", 0.9, Ceilings{BookUSD: 1.0})
|
||
if err != nil || verdict != ReserveOK {
|
||
t.Fatalf("first reserve: %v %v", verdict, err)
|
||
}
|
||
_, verdict, err = s.Reserve("book", 0.2, Ceilings{BookUSD: 1.0})
|
||
if err != nil || verdict != ReserveDeniedBook {
|
||
t.Fatalf("second reserve must hit the book ceiling, got %v %v", verdict, err)
|
||
}
|
||
// Release gives headroom back.
|
||
if err := s.Release(res); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
_, verdict, err = s.Reserve("book", 0.2, Ceilings{BookUSD: 1.0})
|
||
if err != nil || verdict != ReserveOK {
|
||
t.Fatalf("after release: %v %v", verdict, err)
|
||
}
|
||
|
||
// The day ceiling is counted across all books over the day.
|
||
_, verdict, err = s.Reserve("other-book", 0.1, Ceilings{DayUSD: 0.25})
|
||
if err != nil || verdict != ReserveDeniedDay {
|
||
t.Fatalf("day ceiling must deny across books, got %v %v", verdict, err)
|
||
}
|
||
}
|
||
|
||
func TestConcurrentReservesRespectCeiling(t *testing.T) {
|
||
s, _ := openTemp(t)
|
||
mustSnapshotAndJob(t, s)
|
||
caps := Ceilings{BookUSD: 0.5}
|
||
|
||
// 20 concurrent reservations of 0.1 at a ceiling of 0.5: at most 5 can pass.
|
||
var wg sync.WaitGroup
|
||
var mu sync.Mutex
|
||
admitted := 0
|
||
for i := 0; i < 20; i++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
_, verdict, err := s.Reserve("book", 0.1, caps)
|
||
if err != nil {
|
||
t.Error(err)
|
||
return
|
||
}
|
||
if verdict == ReserveOK {
|
||
mu.Lock()
|
||
admitted++
|
||
mu.Unlock()
|
||
}
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
if admitted != 5 {
|
||
t.Fatalf("ceiling 0.5 with 0.1 reserves must admit exactly 5, admitted %d", admitted)
|
||
}
|
||
}
|
||
|
||
func TestRecoveryClearsStaleReservations(t *testing.T) {
|
||
path := filepath.Join(t.TempDir(), "test.db")
|
||
s, err := Open(path)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, verdict, err := s.Reserve("book", 0.3, Ceilings{BookUSD: 1}); err != nil || verdict != ReserveOK {
|
||
t.Fatalf("%v %v", verdict, err)
|
||
}
|
||
s.Close() // "crashed" without releasing the reservation
|
||
|
||
s2, err := Open(path)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer s2.Close()
|
||
committed, reserved, _ := s2.SpentUSD("book")
|
||
if committed != 0 || reserved != 0 {
|
||
t.Fatalf("reopen must clear stale reservations: committed=%v reserved=%v", committed, reserved)
|
||
}
|
||
}
|
||
|
||
func TestJobKeepsOriginalSnapshot(t *testing.T) {
|
||
s, _ := openTemp(t)
|
||
if err := s.UpsertSnapshot("snapA", "brief", `{}`); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if err := s.UpsertSnapshot("snapB", "brief", `{"v":2}`); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
j1, err := s.EnsureJob("book", 1, "draft", "snapA")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// A repeat EnsureJob with a different snapshot does NOT rebind the job (Р6:
|
||
// the snapshot is pinned to the job until it completes).
|
||
j2, err := s.EnsureJob("book", 1, "draft", "snapB")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if j1.ID != j2.ID || j2.SnapshotID != "snapA" {
|
||
t.Fatalf("job must keep its original snapshot: %+v", j2)
|
||
}
|
||
}
|
||
|
||
// RequestLogRows must return ALL rows readable by the caller AFTER the method
|
||
// returns: previously it handed back *sql.Rows with `defer cancel()`, and the context
|
||
// was canceled mid lazy iteration in tmctl report («context canceled», table cut off —
|
||
// a finding from the real Phase 0 acceptance).
|
||
func TestRequestLogRowsMaterializes(t *testing.T) {
|
||
s, _ := openTemp(t)
|
||
for i := 1; i <= 3; i++ {
|
||
if err := s.InsertRequestLog(RequestLog{BookID: "book", Stage: "draft", PromptTokens: 10 * i, OK: true}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
}
|
||
// A row from another book must not leak into the result set.
|
||
if err := s.InsertRequestLog(RequestLog{BookID: "other", Stage: "draft", OK: true}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
rows, err := s.RequestLogRows("book")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if len(rows) != 3 {
|
||
t.Fatalf("want 3 materialized rows, got %d (context canceled mid-iteration?)", len(rows))
|
||
}
|
||
if rows[0].PromptTokens != 10 || rows[2].PromptTokens != 30 {
|
||
t.Fatalf("rows not in id order or wrong data: %+v", rows)
|
||
}
|
||
}
|
||
|
||
func TestMigrateIsIdempotent(t *testing.T) {
|
||
path := filepath.Join(t.TempDir(), "test.db")
|
||
for i := 0; i < 3; i++ {
|
||
s, err := Open(path)
|
||
if err != nil {
|
||
t.Fatalf("open #%d: %v", i, err)
|
||
}
|
||
s.Close()
|
||
}
|
||
}
|
||
|
||
func TestSingleProcessOwnership(t *testing.T) {
|
||
// A second Open of the same project must be denied: a parallel process
|
||
// (report during translate) would zero out live reservations with its recovery pass.
|
||
path := filepath.Join(t.TempDir(), "test.db")
|
||
s1, err := Open(path)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer s1.Close()
|
||
if _, err := Open(path); err == nil {
|
||
t.Fatal("second Open must be denied while the project is locked")
|
||
}
|
||
s1.Close()
|
||
s2, err := Open(path)
|
||
if err != nil {
|
||
t.Fatalf("after Close the lock must be free: %v", err)
|
||
}
|
||
s2.Close()
|
||
}
|