Centralize Postgres test fixture in internal/test
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:
@@ -18,8 +18,6 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -30,123 +28,6 @@ import (
|
||||
"go.probo.inc/probo/pkg/llm"
|
||||
)
|
||||
|
||||
var (
|
||||
sharedPGClientCoredata *pg.Client
|
||||
pgOnceCoredata sync.Once
|
||||
pgInitErrCoredata error
|
||||
ensureTableOnceCoredata sync.Once
|
||||
ensureTableErrCoredata error
|
||||
)
|
||||
|
||||
func pgClient(t *testing.T) *pg.Client {
|
||||
t.Helper()
|
||||
|
||||
pgOnceCoredata.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"
|
||||
}
|
||||
|
||||
sharedPGClientCoredata, pgInitErrCoredata = pg.NewClient(
|
||||
pg.WithAddr(addr),
|
||||
pg.WithUser(user),
|
||||
pg.WithPassword(password),
|
||||
pg.WithDatabase(database),
|
||||
pg.WithPoolSize(5),
|
||||
)
|
||||
if pgInitErrCoredata != nil {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
pgInitErrCoredata = sharedPGClientCoredata.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
|
||||
_, err := conn.Exec(ctx, "SELECT 1")
|
||||
return err
|
||||
})
|
||||
})
|
||||
|
||||
if pgInitErrCoredata != nil {
|
||||
t.Skipf("cannot connect to test database: %v", pgInitErrCoredata)
|
||||
}
|
||||
|
||||
ensureAgentRunsTable(t, sharedPGClientCoredata)
|
||||
|
||||
return sharedPGClientCoredata
|
||||
}
|
||||
|
||||
func ensureAgentRunsTable(t *testing.T, client *pg.Client) {
|
||||
t.Helper()
|
||||
|
||||
ensureTableOnceCoredata.Do(func() {
|
||||
ctx := context.Background()
|
||||
ensureTableErrCoredata = 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, ensureTableErrCoredata, "cannot ensure agent_runs table")
|
||||
}
|
||||
|
||||
func insertPendingRun(
|
||||
t *testing.T,
|
||||
client *pg.Client,
|
||||
|
||||
Reference in New Issue
Block a user