diff --git a/internal/test/pg.go b/internal/test/pg.go new file mode 100644 index 000000000..ecfecba22 --- /dev/null +++ b/internal/test/pg.go @@ -0,0 +1,151 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +// Package test provides shared Postgres test fixtures for packages that +// exercise the database layer against a real Postgres instance. Centralizing +// the connection bootstrap and agent_runs schema setup here keeps a single +// copy of that logic so it cannot drift between the coredata and agentrun +// test suites. +package test + +import ( + "context" + "fmt" + "io" + "net" + "net/url" + "os" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.gearno.de/kit/log" + "go.gearno.de/kit/migrator" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/coredata" +) + +const ( + // pgURLEnvVar points the integration tests at a migrated test database. + pgURLEnvVar = "PROBO_TEST_PG_URL" + + // defaultPGURL targets the local compose Postgres so tests run with zero + // configuration against a developer's stack. When the database is not + // reachable (e.g. in CI, where `make test` runs without Postgres) the + // connection check fails and the test is skipped. + defaultPGURL = "postgres://probod:probod@localhost:5432/probod_test" +) + +var ( + sharedPGClient *pg.Client + pgOnce sync.Once + pgInitErr error + migrateOnce sync.Once + migrateErr error +) + +// PGClient returns a process-wide shared pg.Client connected to the test +// database described by the PROBO_TEST_PG_URL environment variable (falling +// back to a local compose Postgres), applying the agent_runs migrations on +// first use. The test is skipped when no database is reachable so `make test` +// stays a pure unit-test run. +func PGClient(t *testing.T) *pg.Client { + t.Helper() + + pgOnce.Do( + func() { + dsn := os.Getenv(pgURLEnvVar) + if dsn == "" { + dsn = defaultPGURL + } + + u, err := url.Parse(dsn) + if err != nil { + pgInitErr = fmt.Errorf("cannot parse %s: %w", pgURLEnvVar, err) + return + } + + opts := []pg.Option{pg.WithPoolSize(25)} + + 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:])) + } + + sharedPGClient, pgInitErr = pg.NewClient(opts...) + 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) + } + + migrateSchema(t, sharedPGClient) + + return sharedPGClient +} + +// migrateSchema applies the full coredata migration set to the shared test +// database. The migrator is idempotent (it records applied versions in +// schema_versions and serializes through an advisory lock), so running it +// once per process brings any reachable database up to date regardless of +// its starting state. +func migrateSchema(t *testing.T, client *pg.Client) { + t.Helper() + + migrateOnce.Do( + func() { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + logger := log.NewLogger(log.WithOutput(io.Discard)) + + migrateErr = migrator. + NewMigrator(client, coredata.Migrations, logger). + Run(ctx, "migrations") + }, + ) + + require.NoError(t, migrateErr, "cannot migrate test database schema") +} diff --git a/pkg/agentrun/helpers_test.go b/pkg/agentrun/helpers_test.go index c9168ba05..8fb34bc8e 100644 --- a/pkg/agentrun/helpers_test.go +++ b/pkg/agentrun/helpers_test.go @@ -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() diff --git a/pkg/agentrun/service_test.go b/pkg/agentrun/service_test.go index 81d7a56cb..9e0ccacd5 100644 --- a/pkg/agentrun/service_test.go +++ b/pkg/agentrun/service_test.go @@ -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) diff --git a/pkg/agentrun/worker_test.go b/pkg/agentrun/worker_test.go index 8a2256a45..9848d64b6 100644 --- a/pkg/agentrun/worker_test.go +++ b/pkg/agentrun/worker_test.go @@ -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{}) diff --git a/pkg/cookiebanner/pattern_analysis_worker_process_test.go b/pkg/cookiebanner/pattern_analysis_worker_process_test.go index 52f466bfc..6118eb5ef 100644 --- a/pkg/cookiebanner/pattern_analysis_worker_process_test.go +++ b/pkg/cookiebanner/pattern_analysis_worker_process_test.go @@ -18,66 +18,18 @@ import ( "context" "errors" "io" - "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/log" "go.gearno.de/kit/pg" + "go.probo.inc/probo/internal/test" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" ) -const testPgDSNEnvVar = "PROBO_TEST_PG_URL" - -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) - - 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 -} - // workerFixture bootstraps the parent rows the worker's transaction // needs: an organization, a cookie banner, an uncategorised category, // and a normal category. Patterns/detected trackers are seeded @@ -314,7 +266,7 @@ func seedThirdParty(t *testing.T, ctx context.Context, client *pg.Client, fx wor func TestPatternAnalysisWorker_PromotesSourceOnExistingGlob(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -394,7 +346,7 @@ func TestPatternAnalysisWorker_PromotesSourceOnExistingGlob(t *testing.T) { func TestPatternAnalysisWorker_AdoptionTriggersDraftVersion(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -480,7 +432,7 @@ func TestPatternAnalysisWorker_AdoptionTriggersDraftVersion(t *testing.T) { func TestPatternAnalysisWorker_AdoptionPromotesSourceCrossCategory(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -563,7 +515,7 @@ func TestPatternAnalysisWorker_AdoptionPromotesSourceCrossCategory(t *testing.T) func TestReportDetectedTrackers_PromotesSourceOnExistingGlob(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -628,7 +580,7 @@ func TestReportDetectedTrackers_PromotesSourceOnExistingGlob(t *testing.T) { func TestPatternAnalysisWorker_MergeWithoutAdoptionSkipsDraftVersion(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -690,7 +642,7 @@ func TestPatternAnalysisWorker_MergeWithoutAdoptionSkipsDraftVersion(t *testing. func TestPatternAnalysisWorker_GlobInheritsUnanimousMapping(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -749,7 +701,7 @@ func TestPatternAnalysisWorker_GlobInheritsUnanimousMapping(t *testing.T) { func TestPatternAnalysisWorker_GlobSkipsConflictingMapping(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) diff --git a/pkg/cookiebanner/tracker_mapping_worker_test.go b/pkg/cookiebanner/tracker_mapping_worker_test.go index 5bb459c37..c8c4571ca 100644 --- a/pkg/cookiebanner/tracker_mapping_worker_test.go +++ b/pkg/cookiebanner/tracker_mapping_worker_test.go @@ -24,6 +24,7 @@ import ( "github.com/stretchr/testify/require" "go.gearno.de/kit/log" "go.gearno.de/kit/pg" + "go.probo.inc/probo/internal/test" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" ) @@ -164,7 +165,7 @@ func promote( func TestPromoteThirdParty_ExactCommonLink(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -194,7 +195,7 @@ func TestPromoteThirdParty_ExactCommonLink(t *testing.T) { func TestPromoteThirdParty_HeuristicMatch(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -235,7 +236,7 @@ func TestPromoteThirdParty_HeuristicMatch(t *testing.T) { func TestPromoteThirdParty_FallbackCreate(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -265,7 +266,7 @@ func TestPromoteThirdParty_FallbackCreate(t *testing.T) { func TestResolveOrgThirdParty_CreationGated(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -286,7 +287,7 @@ func TestResolveOrgThirdParty_CreationGated(t *testing.T) { func TestProcess_PreservesCatalogMappingOnReTrigger(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -314,7 +315,7 @@ func TestProcess_PreservesCatalogMappingOnReTrigger(t *testing.T) { func TestProcess_UncategorisedPatternIsNotPromoted(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -359,7 +360,7 @@ func TestProcess_UncategorisedPatternIsNotPromoted(t *testing.T) { func TestProcess_ExtensionPatternIsNotPromoted(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -405,7 +406,7 @@ func TestProcess_ExtensionPatternIsNotPromoted(t *testing.T) { func TestProcess_NoOpWhenAlreadyPromoted(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -465,7 +466,7 @@ func TestProcess_NoOpWhenAlreadyPromoted(t *testing.T) { func TestMatchBySiblingOrigin_SiblingWithThirdPartyID(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -595,7 +596,7 @@ func TestMatchBySiblingOrigin_SiblingWithThirdPartyID(t *testing.T) { func TestMatchBySiblingOrigin_AmbiguousThirdParties(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -779,7 +780,7 @@ func TestMatchBySiblingOrigin_AmbiguousThirdParties(t *testing.T) { func TestMatchBySiblingOrigin_NoSiblings(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -820,7 +821,7 @@ func TestMatchBySiblingOrigin_NoSiblings(t *testing.T) { func TestMatchBySiblingOrigin_EmptyDomains(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedWorkerFixture(t, ctx, client) @@ -861,7 +862,7 @@ func TestMatchBySiblingOrigin_EmptyDomains(t *testing.T) { func TestMatchBySiblingOrigin_ConvergentSiblings(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1015,7 +1016,7 @@ func TestMatchBySiblingOrigin_ConvergentSiblings(t *testing.T) { func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1065,7 +1066,7 @@ func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) { func TestProcess_BackfillsCommonThirdPartyFromSibling(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1219,7 +1220,7 @@ func TestProcess_BackfillsCommonThirdPartyFromSibling(t *testing.T) { func TestProcess_UncategorisedLinksExistingThirdParty(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1282,7 +1283,7 @@ func TestProcess_UncategorisedLinksExistingThirdParty(t *testing.T) { func TestProcess_SiblingPromotionOnFirstPartyOrigin(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1435,7 +1436,7 @@ func TestProcess_SiblingPromotionOnFirstPartyOrigin(t *testing.T) { func TestProcess_ReenqueuesUnmappedSiblingOnResolve(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1567,7 +1568,7 @@ func TestProcess_ReenqueuesUnmappedSiblingOnResolve(t *testing.T) { func TestProcess_DoesNotReenqueuePromotedOrExtensionSiblings(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) @@ -1709,7 +1710,7 @@ func TestProcess_DoesNotReenqueuePromotedOrExtensionSiblings(t *testing.T) { func TestProcess_NoReenqueueWhenCommonThirdPartyPreexisted(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) diff --git a/pkg/coredata/access_entry_upsert_test.go b/pkg/coredata/access_entry_upsert_test.go index d6639d5fb..f544e3161 100644 --- a/pkg/coredata/access_entry_upsert_test.go +++ b/pkg/coredata/access_entry_upsert_test.go @@ -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) diff --git a/pkg/coredata/common_tracker_pattern_test.go b/pkg/coredata/common_tracker_pattern_test.go index 5ecfc2ba9..300832533 100644 --- a/pkg/coredata/common_tracker_pattern_test.go +++ b/pkg/coredata/common_tracker_pattern_test.go @@ -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) diff --git a/pkg/coredata/pg_checkpointer_test.go b/pkg/coredata/pg_checkpointer_test.go index fcdf3c169..343b54a7c 100644 --- a/pkg/coredata/pg_checkpointer_test.go +++ b/pkg/coredata/pg_checkpointer_test.go @@ -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( diff --git a/pkg/coredata/pg_test_helpers_test.go b/pkg/coredata/pg_test_helpers_test.go index 405cd2ffd..e340ced56 100644 --- a/pkg/coredata/pg_test_helpers_test.go +++ b/pkg/coredata/pg_test_helpers_test.go @@ -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, diff --git a/pkg/coredata/tracker_pattern_test.go b/pkg/coredata/tracker_pattern_test.go index d1daa5cff..74c31c612 100644 --- a/pkg/coredata/tracker_pattern_test.go +++ b/pkg/coredata/tracker_pattern_test.go @@ -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) diff --git a/pkg/iam/authorizer_batch_test.go b/pkg/iam/authorizer_batch_test.go index 633b3b269..34a1b09f1 100644 --- a/pkg/iam/authorizer_batch_test.go +++ b/pkg/iam/authorizer_batch_test.go @@ -20,16 +20,14 @@ import ( "fmt" "io" "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/log" "go.gearno.de/kit/pg" + "go.probo.inc/probo/internal/test" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/iam" @@ -37,8 +35,6 @@ import ( "go.probo.inc/probo/pkg/mail" ) -const testPgDSNEnvVar = "PROBO_TEST_PG_URL" - type batchAuthorizeFixture struct { tenantID gid.TenantID identityID gid.GID @@ -56,7 +52,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("happy path", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -81,7 +77,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("mixed organization batch", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -114,7 +110,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("mixed entity type batch", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -144,7 +140,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("unsupported resource type for batch attributes", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -169,7 +165,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("single deny rolls back entire batch", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, &fixture.frameworkID1) @@ -195,7 +191,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("duplicate resources in batch", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -238,7 +234,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("dry-run does not write audit logs", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -264,7 +260,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("missing principal identity returns wrapped principal attributes error", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -288,7 +284,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("bulk insert failure aborts transaction", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithIdentityScopedStatements( @@ -320,7 +316,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("assumption required", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -348,7 +344,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("assumption succeeds with active child session", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -384,7 +380,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("assumption fails when child session is expired", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -421,7 +417,7 @@ func TestAuthorizer_AuthorizeBatch(t *testing.T) { t.Run("no membership ignores assumption check", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizer(client, action, nil) @@ -462,7 +458,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("returns nil decisions when every item is allowed", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -498,7 +494,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("returns per-item decisions on partial denial without aborting", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -541,7 +537,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("writes no audit logs when every item is denied", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -578,7 +574,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("skips audit log entries for dry-run allowed items", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -613,7 +609,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("rejects mixed organization batch", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -692,7 +688,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("assumption error is recorded only on items that require the check", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -734,7 +730,7 @@ func TestAuthorizer_AuthorizeMulti(t *testing.T) { t.Run("per-item resource attributes are honoured by the policy", func(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) fixture := seedBatchAuthorizeFixture(t, context.Background(), client) action := newBatchTestAction() authorizer := newTestAuthorizerWithStatements( @@ -1059,48 +1055,3 @@ func countAuditLogsForAction(t *testing.T, ctx context.Context, client *pg.Clien func newBatchTestAction() string { return fmt.Sprintf("test:framework-%d:get", time.Now().UnixNano()) } - -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) - - 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 -} diff --git a/pkg/thirdparty/resolver_test.go b/pkg/thirdparty/resolver_test.go index f676b18df..7f42a6032 100644 --- a/pkg/thirdparty/resolver_test.go +++ b/pkg/thirdparty/resolver_test.go @@ -17,67 +17,19 @@ package thirdparty import ( "context" "io" - "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/log" "go.gearno.de/kit/pg" + "go.probo.inc/probo/internal/test" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/slug" ) -const testPgDSNEnvVar = "PROBO_TEST_PG_URL" - -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) - - 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 -} - func discardLogger() *log.Logger { return log.NewLogger(log.WithOutput(io.Discard)) } @@ -126,7 +78,7 @@ func seedCatalogThirdParty( func TestResolveOrCreateCommonThirdParty(t *testing.T) { t.Parallel() - client := newTestPgClient(t) + client := test.PGClient(t) ctx := context.Background() logger := discardLogger()