Agent-System/internal/dataset/pg_test.go

package dataset_test

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/url"
	"os"
	"sync"
	"testing"
	"time"

	"chat-agent-system/internal/artifactvault"
	"chat-agent-system/internal/dataset"
	"chat-agent-system/internal/db"

	"github.com/jackc/pgx/v5/pgxpool"
)

func init() {
	if _, err := os.Stat("deploy/postgres/goose"); err == nil {
		return
	}
	for _, cand := range []string{"../..", "../../.."} {
		if fi, err := os.Stat(cand + "/deploy/postgres/goose"); err == nil && fi.IsDir() {
			_ = os.Chdir(cand)
			return
		}
	}
}

func rewriteDatabaseName(dsn, name string) string {
	if u, err := url.Parse(dsn); err == nil && u.Scheme != "" {
		u.Path = "/" + name
		return u.String()
	}
	return dsn
}

var (
	isoOnce            sync.Once
	isoPoolA, isoPoolB *pgxpool.Pool
	isoErr             error
	isoCleanup         func()
)

func TestMain(m *testing.M) {
	code := m.Run()
	if isoCleanup != nil {
		isoCleanup()
	}
	os.Exit(code)
}

func openIsolatedU9(t *testing.T) (*pgxpool.Pool, *pgxpool.Pool) {
	t.Helper()
	if testing.Short() {
		t.Skip("U9 PG races are not part of -short")
	}
	dsn := db.SkipIfNoTestDatabaseURL(t)
	db.RefuseSharedEnvDSN(t, dsn)

	isoOnce.Do(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
		defer cancel()
		adminDSN := rewriteDatabaseName(dsn, "postgres")
		admin, err := pgxpool.New(ctx, adminDSN)
		if err != nil {
			isoErr = err
			return
		}
		const isolated = "agent_rel_m3_u9"
		_, _ = admin.Exec(ctx, `SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname=$1 AND pid <> pg_backend_pid()`, isolated)
		_, _ = admin.Exec(ctx, `DROP DATABASE IF EXISTS `+isolated)
		if _, err := admin.Exec(ctx, `CREATE DATABASE `+isolated); err != nil {
			admin.Close()
			isoErr = err
			return
		}
		isoDSN := rewriteDatabaseName(dsn, isolated)
		poolA, err := pgxpool.New(ctx, isoDSN)
		if err != nil {
			admin.Close()
			isoErr = err
			return
		}
		if err := db.RunMigrations(ctx, poolA, db.ResolveMigrationsDir()); err != nil {
			poolA.Close()
			admin.Close()
			isoErr = err
			return
		}
		poolB, err := pgxpool.New(ctx, isoDSN)
		if err != nil {
			poolA.Close()
			admin.Close()
			isoErr = err
			return
		}
		isoPoolA, isoPoolB = poolA, poolB
		isoCleanup = func() {
			poolA.Close()
			poolB.Close()
			_, _ = admin.Exec(context.Background(), `SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname=$1 AND pid <> pg_backend_pid()`, isolated)
			_, _ = admin.Exec(context.Background(), `DROP DATABASE IF EXISTS `+isolated)
			admin.Close()
		}
	})
	if isoErr != nil {
		t.Fatalf("isolated db: %v", isoErr)
	}
	return isoPoolA, isoPoolB
}

func shaBody(s string) (string, string) {
	sum := sha256.Sum256([]byte(s))
	return hex.EncodeToString(sum[:]), s
}

func seedCommit(t *testing.T, pool *pgxpool.Pool, docID, body string) {
	t.Helper()
	ctx := context.Background()
	sha, md := shaBody(body)
	if _, err := pool.Exec(ctx, `
		INSERT INTO document_content_blobs (content_sha256, content_markdown, byte_size)
		VALUES ($1,$2,$3) ON CONFLICT DO NOTHING
	`, sha, md, len(md)); err != nil {
		t.Fatal(err)
	}
	var snapID, verID int64
	if err := pool.QueryRow(ctx, `
		INSERT INTO document_snapshots (doc_id, content_sha256, source_kind)
		VALUES ($1,$2,'agent') RETURNING id
	`, docID, sha).Scan(&snapID); err != nil {
		t.Fatal(err)
	}
	if err := pool.QueryRow(ctx, `
		INSERT INTO document_versions (doc_id, version_no, status, origin_kind, baseline_snapshot_id, head_snapshot_id)
		VALUES ($1,1,'sealed','agent',$2,$2) RETURNING id
	`, docID, snapID).Scan(&verID); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `
		INSERT INTO document_change_events (
			doc_id, version_id, event_no, to_snapshot_id, change_kind, actor_kind, idempotency_key
		) VALUES ($1,$2,1,$3,'agent_commit','agent',$4)
	`, docID, verID, snapID, "u9:"+docID); err != nil {
		t.Fatal(err)
	}
}

func TestPG_DuplicateEnqueueAndSkipLocked(t *testing.T) {
	poolA, poolB := openIsolatedU9(t)
	ctx := context.Background()
	seedCommit(t, poolA, "doc_u9_commit", "# u9 agent commit body")

	job := dataset.NewJob("u9-race", "dev", "scratch-u9", 38)
	pgA := &dataset.PG{Pool: poolA, Env: "dev"}
	pgB := &dataset.PG{Pool: poolB, Env: "dev"}
	if err := pgA.EnsureJob(ctx, job); err != nil {
		t.Fatal(err)
	}
	n1, err := pgA.EnqueueKind(ctx, job, artifactvault.KindAgentCommit, 0)
	if err != nil {
		t.Fatal(err)
	}
	n2, err := pgA.EnqueueKind(ctx, job, artifactvault.KindAgentCommit, 0)
	if err != nil {
		t.Fatal(err)
	}
	if n1 < 1 {
		t.Fatalf("expected seeded commit, enqueued %d", n1)
	}
	pending, err := pgA.CountQueue(ctx, job.ID, "pending")
	if err != nil {
		t.Fatal(err)
	}
	if pending != n1 {
		t.Fatalf("duplicate enqueue grew queue pending=%d first=%d second=%d", pending, n1, n2)
	}

	var wg sync.WaitGroup
	errCh := make(chan error, 2)
	for _, pg := range []*dataset.PG{pgA, pgB} {
		wg.Add(1)
		go func(p *dataset.PG) {
			defer wg.Done()
			_, err := p.ClaimAndProcess(ctx, job, 50)
			errCh <- err
		}(pg)
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatal(err)
		}
	}
	arts, err := pgA.CountArtifacts(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if arts != 1 {
		t.Fatalf("want 1 artifact after two workers, got %d", arts)
	}
	st, err := pgA.RecalcJob(ctx, job.ID)
	if err != nil {
		t.Fatal(err)
	}
	if st.Succeeded != 1 {
		t.Fatalf("recalc succeeded=%d scanned=%d", st.Succeeded, st.Scanned)
	}
}

func TestPG_PauseStatusAndRegistryImport(t *testing.T) {
	poolA, _ := openIsolatedU9(t)
	ctx := context.Background()
	pg := &dataset.PG{Pool: poolA, Env: "dev"}
	job := dataset.NewJob("u9-pause", "dev", "scratch-u9", 38)
	if err := pg.EnsureJob(ctx, job); err != nil {
		t.Fatal(err)
	}
	if err := pg.SetJobStatus(ctx, job.ID, dataset.JobPaused); err != nil {
		t.Fatal(err)
	}
	var status string
	if err := poolA.QueryRow(ctx, `SELECT status FROM dataset_backfill_jobs WHERE job_id=$1`, job.ID).Scan(&status); err != nil {
		t.Fatal(err)
	}
	if status != dataset.JobPaused {
		t.Fatalf("status=%s", status)
	}
	for i, d := range []dataset.Dataset{
		{ID: "pinned-clean", Class: dataset.ClassPinned, DisplayName: "a", StorageIdentity: "scratch-u9", Status: "approved"},
		{ID: "demo-minimal", Class: dataset.ClassDemo, DisplayName: "b", StorageIdentity: "scratch-u9", Status: "approved"},
	} {
		d.SourceSchema = 38
		d.SourceCommit = fmt.Sprintf("c%d", i)
		if err := pg.UpsertDataset(ctx, d); err != nil {
			t.Fatal(err)
		}
	}
	var n int
	if err := poolA.QueryRow(ctx, `SELECT COUNT(*) FROM dataset_registry`).Scan(&n); err != nil || n < 2 {
		t.Fatalf("registry count=%d err=%v", n, err)
	}
}