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

64
pkg/coredata/access_entry_upsert_test.go
View File

@@ -16,73 +16,17 @@ package coredata_test
import (
"context"
"net"
"net/url"
"os"
"testing"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/internal/test"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
// testPgDSNEnvVar is the environment variable that points the integration
// tests at a migrated test database. When unset, the tests are skipped so
// `make test` stays a pure unit-test run.
const testPgDSNEnvVar = "PROBO_TEST_PG_URL"
// newTestPgClient returns a pg.Client connected to the test database, or
// skips the test if no DSN is configured.
func newTestPgClient(t *testing.T) *pg.Client {
t.Helper()
dsn := os.Getenv(testPgDSNEnvVar)
if dsn == "" {
t.Skipf("skipping: %s not set (requires a migrated test database)", testPgDSNEnvVar)
}
u, err := url.Parse(dsn)
require.NoError(t, err, "invalid %s value", testPgDSNEnvVar)
// Each test builds its own pg.Client, so we provide a fresh Prometheus
// registry every time to avoid "duplicate collector" panics when tests
// run in parallel.
opts := []pg.Option{pg.WithRegisterer(prometheus.NewRegistry())}
if u.Host != "" {
host := u.Host
if u.Port() == "" {
host = net.JoinHostPort(u.Hostname(), "5432")
}
opts = append(opts, pg.WithAddr(host))
}
if u.User != nil {
opts = append(opts, pg.WithUser(u.User.Username()))
if password, ok := u.User.Password(); ok {
opts = append(opts, pg.WithPassword(password))
}
}
if len(u.Path) > 1 {
opts = append(opts, pg.WithDatabase(u.Path[1:]))
}
client, err := pg.NewClient(opts...)
require.NoError(t, err)
t.Cleanup(func() {
client.Close()
})
return client
}
// accessEntryFixture bootstraps the parent rows (organization, campaign,
// source) that the access_entries FKs require.
type accessEntryFixture struct {
@@ -179,7 +123,7 @@ func seedAccessEntryFixture(t *testing.T, ctx context.Context, client *pg.Client
func TestAccessEntry_Upsert_FreezesDecidedFields(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
fx := seedAccessEntryFixture(t, ctx, client)
@@ -324,7 +268,7 @@ func TestAccessEntry_Upsert_FreezesDecidedFields(t *testing.T) {
func TestAccessEntry_Upsert_RefreshesSourceTrackingFields(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
fx := seedAccessEntryFixture(t, ctx, client)
@@ -414,7 +358,7 @@ func TestAccessEntry_Upsert_RefreshesSourceTrackingFields(t *testing.T) {
func TestAccessEntry_Upsert_InsertsActiveAccount(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
fx := seedAccessEntryFixture(t, ctx, client)

9
pkg/coredata/common_tracker_pattern_test.go
View File

@@ -22,6 +22,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/internal/test"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
@@ -107,7 +108,7 @@ func loadCommonTrackerPattern(
func TestCommonTrackerPattern_SetEnriched_AllowsEmptyDescription(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
now := time.Now().UTC().Truncate(time.Microsecond)
@@ -151,7 +152,7 @@ func TestCommonTrackerPattern_SetEnriched_AllowsEmptyDescription(t *testing.T) {
func TestCommonTrackerPattern_SetEnriched_LinksThirdPartyWithoutOverride(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
party := seedCommonThirdParty(t, ctx, client)
@@ -211,7 +212,7 @@ func TestCommonTrackerPattern_SetEnriched_LinksThirdPartyWithoutOverride(t *test
func TestCommonTrackerPattern_Upsert_RequeuesBlankRowOnThirdPartyLink(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
party := seedCommonThirdParty(t, ctx, client)
@@ -272,7 +273,7 @@ func TestCommonTrackerPattern_Upsert_RequeuesBlankRowOnThirdPartyLink(t *testing
func TestCommonTrackerPattern_Upsert_KeepsDescribedRowTerminal(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
party := seedCommonThirdParty(t, ctx, client)

3
pkg/coredata/pg_checkpointer_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/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
@@ -29,7 +30,7 @@ import (
func TestPGCheckpointer(t *testing.T) {
t.Parallel()
client := pgClient(t)
client := test.PGClient(t)
store := coredata.NewPGCheckpointer(client)
t.Run(

119
pkg/coredata/pg_test_helpers_test.go
View File

@@ -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,

7
pkg/coredata/tracker_pattern_test.go
View File

@@ -22,6 +22,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/internal/test"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
@@ -173,7 +174,7 @@ func seedTrackerPattern(
func TestTrackerPattern_Update_WritesSource(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
fx := seedTrackerPatternFixture(t, ctx, client)
@@ -221,7 +222,7 @@ func TestTrackerPattern_Update_WritesSource(t *testing.T) {
func TestTrackerPattern_Update_NotFoundForMissingRow(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
fx := seedTrackerPatternFixture(t, ctx, client)
@@ -259,7 +260,7 @@ func TestTrackerPattern_Update_NotFoundForMissingRow(t *testing.T) {
func TestResetStaleMappings(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
client := test.PGClient(t)
ctx := context.Background()
fx := seedTrackerPatternFixture(t, ctx, client)