rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Centralize Postgres test fixture in internal/test

Browse Source 
Each package that exercises the database against a real Postgres
carried its own copy of the connection bootstrap and schema setup.
Those copies had already drifted: some keyed off PROBO_TEST_PG_ADDR
with hardcoded defaults, others off PROBO_TEST_PG_URL, and the
agentrun/coredata suites hand-applied individual agent_runs
migrations to ensure the table existed.

Introduce a single test.PGClient helper that parses PROBO_TEST_PG_URL
(falling back to the local compose database), runs the full coredata
migration set once per process, and skips when no database is
reachable so make test stays a pure unit-test run. Migrate the
agentrun, coredata, cookiebanner, iam, and thirdparty suites onto it
and delete the duplicated helpers so the bootstrap can no longer
diverge.

Signed-off-by: Bryan Frimin <bryan@probo.com>
This commit is contained in:
Bryan Frimin
2026-06-07 12:14:03 +02:00
parent 0a1b47607b
commit a8d4da3916
13 changed files with 233 additions and 513 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
pkg/agentrun/helpers_test.go
View File

@@ -19,7 +19,6 @@ import (
"encoding/json"
"errors"
"fmt"
"os"
"sync"
"testing"
"time"
@@ -34,14 +33,6 @@ import (
"go.probo.inc/probo/pkg/llm"
)
var (
sharedPGClient *pg.Client
pgOnce sync.Once
pgInitErr error
ensureTableOnce sync.Once
ensureTableErr error
)
func testLogger() *log.Logger {
return log.NewLogger(log.WithFormat(log.FormatPretty))
}
@@ -153,114 +144,6 @@ func (r *simpleRegistry) Agent(name string) (*agent.Agent, error) {
return a, nil
}
func pgClient(t *testing.T) *pg.Client {
t.Helper()
pgOnce.Do(func() {
addr := os.Getenv("PROBO_TEST_PG_ADDR")
if addr == "" {
addr = "localhost:5432"
}
user := os.Getenv("PROBO_TEST_PG_USER")
if user == "" {
user = "probod"
}
password := os.Getenv("PROBO_TEST_PG_PASSWORD")
if password == "" {
password = "probod"
}
database := os.Getenv("PROBO_TEST_PG_DATABASE")
if database == "" {
database = "probod_test"
}
sharedPGClient, pgInitErr = pg.NewClient(
pg.WithAddr(addr),
pg.WithUser(user),
pg.WithPassword(password),
pg.WithDatabase(database),
pg.WithPoolSize(5),
)
if pgInitErr != nil {
return
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
pgInitErr = sharedPGClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
_, err := conn.Exec(ctx, "SELECT 1")
return err
})
})
if pgInitErr != nil {
t.Skipf("cannot connect to test database: %v", pgInitErr)
}
ensureAgentRunsTable(t, sharedPGClient)
return sharedPGClient
}
func ensureAgentRunsTable(t *testing.T, client *pg.Client) {
t.Helper()
ensureTableOnce.Do(func() {
ctx := context.Background()
ensureTableErr = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
var exists bool
if err := conn.QueryRow(
ctx,
`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'agent_runs')`,
).Scan(&exists); err != nil {
return fmt.Errorf("cannot check agent_runs existence: %w", err)
}
if !exists {
ddl, err := coredata.Migrations.ReadFile("migrations/20260424T173529Z.sql")
if err != nil {
return fmt.Errorf("cannot read agent_runs base migration: %w", err)
}
if _, err := conn.Exec(ctx, string(ddl)); err != nil {
return fmt.Errorf("cannot apply agent_runs base migration: %w", err)
}
}
var hasLeaseGeneration bool
if err := conn.QueryRow(
ctx,
`SELECT EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_name = 'agent_runs'
AND column_name = 'lease_generation'
)`,
).Scan(&hasLeaseGeneration); err != nil {
return fmt.Errorf("cannot check lease_generation column: %w", err)
}
if !hasLeaseGeneration {
ddl, err := coredata.Migrations.ReadFile("migrations/20260607T060000Z.sql")
if err != nil {
return fmt.Errorf("cannot read agent_runs lease generation migration: %w", err)
}
if _, err := conn.Exec(ctx, string(ddl)); err != nil {
return fmt.Errorf("cannot apply agent_runs lease generation migration: %w", err)
}
}
return nil
})
})
require.NoError(t, ensureTableErr, "cannot ensure agent_runs table")
}
func insertTestOrganization(t *testing.T, client *pg.Client) gid.GID {
t.Helper()

7
pkg/agentrun/service_test.go
View File

@@ -20,6 +20,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.probo.inc/probo/internal/test"
"go.probo.inc/probo/pkg/agentrun"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
@@ -27,7 +28,7 @@ import (
)
func TestService_Get(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
svc := agentrun.NewService(client)
run := insertPendingRun(
@@ -49,7 +50,7 @@ func TestService_Get(t *testing.T) {
}
func TestService_ListForOrganizationID(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
svc := agentrun.NewService(client)
orgID := insertTestOrganization(t, client)
@@ -81,7 +82,7 @@ func TestService_ListForOrganizationID(t *testing.T) {
}
func TestService_CountForOrganizationID(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
svc := agentrun.NewService(client)
orgID := insertTestOrganization(t, client)

21
pkg/agentrun/worker_test.go
View File

@@ -30,6 +30,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.probo.inc/probo/internal/test"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agentrun"
"go.probo.inc/probo/pkg/coredata"
@@ -37,7 +38,7 @@ import (
)
func TestWorker_PicksUpAndCompletes(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
ag := newDummyAgent(
"echo-agent",
[]*llm.ChatCompletionResponse{
@@ -79,7 +80,7 @@ func TestWorker_PicksUpAndCompletes(t *testing.T) {
}
func TestWorker_StopAndResume(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
store := coredata.NewPGCheckpointer(client)
toolReady := make(chan struct{})
@@ -187,7 +188,7 @@ func TestWorker_StopAndResume(t *testing.T) {
// child as active, and restore must resolve it from the registry so the
// resumed run continues in that branch and completes.
func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
store := coredata.NewPGCheckpointer(client)
toolReady := make(chan struct{})
@@ -314,7 +315,7 @@ func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) {
}
func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
store := coredata.NewPGCheckpointer(client)
toolReady := make(chan struct{})
@@ -444,7 +445,7 @@ func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) {
}
func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
store := coredata.NewPGCheckpointer(client)
toolReady := make(chan struct{})
@@ -593,7 +594,7 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) {
}
func TestWorker_HeartbeatLeaseLostLeavesRunForRecovery(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
toolReady := make(chan struct{})
toolRelease := make(chan struct{})
@@ -685,7 +686,7 @@ func TestWorker_HeartbeatLeaseLostLeavesRunForRecovery(t *testing.T) {
}
func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
toolReady := make(chan struct{})
toolRelease := make(chan struct{})
@@ -798,7 +799,7 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) {
}
func TestWorker_UnknownAgentFails(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
run := insertPendingRun(
t,
@@ -833,7 +834,7 @@ func TestWorker_UnknownAgentFails(t *testing.T) {
}
func TestWorker_InvalidInputMessagesFails(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
ag := newDummyAgent(
"worker-agent",
[]*llm.ChatCompletionResponse{
@@ -929,7 +930,7 @@ func TestWorker_SIGTERM(t *testing.T) {
}
func runSIGTERMSubprocess(t *testing.T) {
client := pgClient(t)
client := test.PGClient(t)
workStarted := make(chan struct{})