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:
151
internal/test/pg.go
Normal file
151
internal/test/pg.go
Normal file
@@ -0,0 +1,151 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
//
|
||||
// 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")
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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{})
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
52
pkg/thirdparty/resolver_test.go
vendored
52
pkg/thirdparty/resolver_test.go
vendored
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user