diff --git a/integration/harness.go b/integration/harness.go index fa568613..1b23136f 100644 --- a/integration/harness.go +++ b/integration/harness.go @@ -71,8 +71,14 @@ type Session struct { DB graph.Database PGPool *pgxpool.Pool Ctx context.Context + cleanupMode CleanupMode } +var ( + errConnectionStringMissing = errors.New("connection string is not set") + errDriverMismatch = errors.New("connection string driver mismatch") +) + // DriverFromConnStr returns the dawgs driver name based on the connection string scheme. func DriverFromConnectionString(connStr string) (string, error) { u, err := url.Parse(connStr) @@ -90,9 +96,7 @@ func DriverFromConnectionString(connStr string) (string, error) { } } -func Open(t *testing.T, opts Options) *Session { - t.Helper() - +func OpenSession(opts Options) (*Session, error) { ctx := context.Background() connEnv := opts.ConnectionStringEnvVar if connEnv == "" { @@ -101,22 +105,16 @@ func Open(t *testing.T, opts Options) *Session { connStr := os.Getenv(connEnv) if connStr == "" { - if opts.SkipIfNoConnection { - t.Skipf("%s env var is not set", connEnv) - } - t.Fatalf("%s env var is not set", connEnv) + return nil, fmt.Errorf("%w: %s env var is not set", errConnectionStringMissing, connEnv) } driver, err := DriverFromConnectionString(connStr) if err != nil { - t.Fatalf("failed to detect driver: %v", err) + return nil, fmt.Errorf("failed to detect driver: %w", err) } if opts.RequireDriver != "" && driver != opts.RequireDriver { - if opts.SkipIfDriverMismatch { - t.Skipf("%s is not a %s connection string", connEnv, opts.RequireDriver) - } - t.Fatalf("%s is not a %s connection string", connEnv, opts.RequireDriver) + return nil, fmt.Errorf("%w: %s is not a %s connection string", errDriverMismatch, connEnv, opts.RequireDriver) } cfg := dawgs.Config{ @@ -128,16 +126,17 @@ func Open(t *testing.T, opts Options) *Session { ConnectionString: connStr, Driver: driver, Ctx: ctx, + cleanupMode: opts.CleanupMode, } if driver == pg.DriverName { poolCfg, err := pgxpool.ParseConfig(connStr) if err != nil { - t.Fatalf("failed to parse pool configuration: %v", err) + return nil, fmt.Errorf("failed to parse pool configuration: %w", err) } pool, err := pg.NewPool(poolCfg) if err != nil { - t.Fatalf("failed to create PG pool: %v", err) + return nil, fmt.Errorf("failed to create PG pool: %w", err) } cfg.Pool = pool session.PGPool = pool @@ -145,38 +144,88 @@ func Open(t *testing.T, opts Options) *Session { db, err := dawgs.Open(ctx, driver, cfg) if err != nil { - t.Fatalf("failed to open database: %v", err) + if session.PGPool != nil { + session.PGPool.Close() + } + return nil, fmt.Errorf("failed to open database: %w", err) } session.DB = db schema := opts.Schema if schema == nil { - schema = buildSchema(t, opts) + schema, err = buildSchema(opts) + if err != nil { + _ = db.Close(ctx) + return nil, fmt.Errorf("failed to build schema: %w", err) + } } if schema != nil { if err := db.AssertSchema(ctx, *schema); err != nil { - t.Fatalf("failed to assert schema: %v", err) + _ = db.Close(ctx) + return nil, fmt.Errorf("failed to assert schema: %w", err) } } + return session, nil +} + +func Open(t *testing.T, opts Options) *Session { + t.Helper() + + session, err := OpenSession(opts) + if err != nil { + if errors.Is(err, errConnectionStringMissing) && opts.SkipIfNoConnection { + t.Skipf("%v", err) + } + if errors.Is(err, errDriverMismatch) && opts.SkipIfDriverMismatch { + t.Skipf("%v", err) + } + t.Fatalf("failed to open integration database: %v", err) + } + t.Cleanup(func() { - if opts.CleanupMode != CloseOnly { - _ = db.WriteTransaction(ctx, func(tx graph.Transaction) error { - return tx.Nodes().Delete() - }) + if err := session.Close(); err != nil { + t.Errorf("failed to clean up integration database: %v", err) } - db.Close(ctx) }) return session } +// ResetGraph removes all nodes and cascading edges from the session database. +func (s *Session) ResetGraph() error { + return s.DB.WriteTransaction(s.Ctx, func(tx graph.Transaction) error { + return tx.Nodes().Delete() + }) +} + +// RegisterGraphCleanup resets the graph at the end of the supplied test scope. +// It is intended for subtests sharing a session opened by their parent test. +func (s *Session) RegisterGraphCleanup(t testing.TB) { + t.Helper() + t.Cleanup(func() { + if err := s.ResetGraph(); err != nil { + t.Errorf("failed to clear graph: %v", err) + } + }) +} + +// Close applies the configured cleanup mode and closes the database connection. +func (s *Session) Close() error { + var cleanupErr error + if s.cleanupMode != CloseOnly { + cleanupErr = s.ResetGraph() + } + if err := s.DB.Close(s.Ctx); err != nil && cleanupErr == nil { + cleanupErr = err + } + return cleanupErr +} + func (s *Session) ClearGraph(t *testing.T) { t.Helper() - if err := s.DB.WriteTransaction(s.Ctx, func(tx graph.Transaction) error { - return tx.Nodes().Delete() - }); err != nil { + if err := s.ResetGraph(); err != nil { t.Fatalf("failed to clear graph: %v", err) } } @@ -243,15 +292,16 @@ func (s *Session) withRollback(t *testing.T, delegate func(tx graph.Transaction) return err } -func buildSchema(t *testing.T, opts Options) *graph.Schema { - t.Helper() - - nodeKinds, edgeKinds := collectKinds(t, opts.Datasets, opts.datasetPath()) +func buildSchema(opts Options) (*graph.Schema, error) { + nodeKinds, edgeKinds, err := collectKinds(opts.Datasets, opts.datasetPath()) + if err != nil { + return nil, err + } nodeKinds = nodeKinds.Add(opts.ExtraNodeKinds...) edgeKinds = edgeKinds.Add(opts.ExtraEdgeKinds...) if len(nodeKinds) == 0 && len(edgeKinds) == 0 { - return nil + return nil, nil } graphName := opts.GraphName @@ -266,25 +316,23 @@ func buildSchema(t *testing.T, opts Options) *graph.Schema { Edges: edgeKinds, }}, DefaultGraph: graph.Graph{Name: graphName}, - } + }, nil } // collectKinds parses the given datasets and returns the union of all node and edge kinds. -func collectKinds(t *testing.T, datasets []string, datasetPath func(name string) string) (graph.Kinds, graph.Kinds) { - t.Helper() - +func collectKinds(datasets []string, datasetPath func(name string) string) (graph.Kinds, graph.Kinds, error) { var nodeKinds, edgeKinds graph.Kinds for _, name := range datasets { f, err := os.Open(datasetPath(name)) if err != nil { - t.Fatalf("failed to open dataset %q for kind scanning: %v", name, err) + return nil, nil, fmt.Errorf("failed to open dataset %q for kind scanning: %w", name, err) } doc, err := opengraph.ParseDocument(f) f.Close() if err != nil { - t.Fatalf("failed to parse dataset %q: %v", name, err) + return nil, nil, fmt.Errorf("failed to parse dataset %q: %w", name, err) } nk, ek := doc.Graph.Kinds() @@ -292,7 +340,7 @@ func collectKinds(t *testing.T, datasets []string, datasetPath func(name string) edgeKinds = edgeKinds.Add(ek...) } - return nodeKinds, edgeKinds + return nodeKinds, edgeKinds, nil } func (s *Options) datasetPath() func(name string) string { diff --git a/integration/harness_test.go b/integration/harness_test.go new file mode 100644 index 00000000..ca34f40b --- /dev/null +++ b/integration/harness_test.go @@ -0,0 +1,40 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package integration + +import ( + "errors" + "testing" +) + +func TestOpenSessionReturnsConnectionStringError(t *testing.T) { + t.Setenv(ConnectionStringEnv, "") + + _, err := OpenSession(Options{}) + if !errors.Is(err, errConnectionStringMissing) { + t.Fatalf("OpenSession() error = %v, want connection string missing error", err) + } +} + +func TestOpenSessionReturnsDriverMismatchError(t *testing.T) { + t.Setenv(ConnectionStringEnv, "postgresql://user:password@localhost:5432/database") + + _, err := OpenSession(Options{RequireDriver: "neo4j"}) + if !errors.Is(err, errDriverMismatch) { + t.Fatalf("OpenSession() error = %v, want driver mismatch error", err) + } +} diff --git a/integration/pgsql_delete_by_kind_test.go b/integration/pgsql_delete_by_kind_test.go index 917db51b..330c0ebc 100644 --- a/integration/pgsql_delete_by_kind_test.go +++ b/integration/pgsql_delete_by_kind_test.go @@ -89,6 +89,11 @@ func TestPostgreSQLDeleteNodesByKinds(t *testing.T) { } t.Run("includeAny deletes matching kinds and cascades edges", func(t *testing.T) { + t.Cleanup(func() { + if err := deleter.DeleteNodesByKinds(ctx, nil, nil); err != nil { + t.Fatalf("failed to clean up nodes between subtests: %v", err) + } + }) createFixture() if err := deleter.DeleteNodesByKinds(ctx, graph.Kinds{kindA}, nil); err != nil { @@ -105,10 +110,14 @@ func TestPostgreSQLDeleteNodesByKinds(t *testing.T) { t.Fatalf("edge count after cascade: got %d, want 0", count) } - cleanupAll(t, ctx, deleter) }) t.Run("excludeAny protects matching kinds", func(t *testing.T) { + t.Cleanup(func() { + if err := deleter.DeleteNodesByKinds(ctx, nil, nil); err != nil { + t.Fatalf("failed to clean up nodes between subtests: %v", err) + } + }) createFixture() // Delete every node except those carrying kindB. @@ -123,10 +132,14 @@ func TestPostgreSQLDeleteNodesByKinds(t *testing.T) { t.Fatalf("kindB node count: got %d, want 2", count) } - cleanupAll(t, ctx, deleter) }) t.Run("undefined include kinds are a safe no-op", func(t *testing.T) { + t.Cleanup(func() { + if err := deleter.DeleteNodesByKinds(ctx, nil, nil); err != nil { + t.Fatalf("failed to clean up nodes between subtests: %v", err) + } + }) createFixture() if err := deleter.DeleteNodesByKinds(ctx, graph.Kinds{missing}, nil); err != nil { @@ -137,10 +150,14 @@ func TestPostgreSQLDeleteNodesByKinds(t *testing.T) { t.Fatalf("node count after no-op delete: got %d, want 4", count) } - cleanupAll(t, ctx, deleter) }) t.Run("undefined exclude kinds fail closed and delete nothing", func(t *testing.T) { + t.Cleanup(func() { + if err := deleter.DeleteNodesByKinds(ctx, nil, nil); err != nil { + t.Fatalf("failed to clean up nodes between subtests: %v", err) + } + }) createFixture() // An unresolved exclusion would otherwise collapse to an unguarded delete; the driver must refuse instead. @@ -152,15 +169,5 @@ func TestPostgreSQLDeleteNodesByKinds(t *testing.T) { t.Fatalf("node count after failed delete: got %d, want 4", count) } - cleanupAll(t, ctx, deleter) }) } - -// cleanupAll removes every node between subtests so each starts from an empty graph. -func cleanupAll(t *testing.T, ctx context.Context, deleter nodesByKindDeleter) { - t.Helper() - - if err := deleter.DeleteNodesByKinds(ctx, nil, nil); err != nil { - t.Fatalf("failed to clean up nodes between subtests: %v", err) - } -} diff --git a/integration/pgsql_delete_relationships_by_kind_test.go b/integration/pgsql_delete_relationships_by_kind_test.go index 5600bb70..068c3d21 100644 --- a/integration/pgsql_delete_relationships_by_kind_test.go +++ b/integration/pgsql_delete_relationships_by_kind_test.go @@ -92,6 +92,7 @@ func TestPostgreSQLDeleteRelationshipsByKinds(t *testing.T) { } t.Run("deletes relationships of the given kind and leaves nodes", func(t *testing.T) { + t.Cleanup(func() { ClearGraph(t, db, ctx) }) createFixture() if err := deleter.DeleteRelationshipsByKinds(ctx, graph.Kinds{edgeKindA}); err != nil { @@ -108,10 +109,10 @@ func TestPostgreSQLDeleteRelationshipsByKinds(t *testing.T) { t.Fatalf("node count: got %d, want 3", count) } - ClearGraph(t, db, ctx) }) t.Run("multiple kinds delete every matching relationship", func(t *testing.T) { + t.Cleanup(func() { ClearGraph(t, db, ctx) }) createFixture() if err := deleter.DeleteRelationshipsByKinds(ctx, graph.Kinds{edgeKindA, edgeKindB}); err != nil { @@ -125,10 +126,10 @@ func TestPostgreSQLDeleteRelationshipsByKinds(t *testing.T) { t.Fatalf("node count: got %d, want 3", count) } - ClearGraph(t, db, ctx) }) t.Run("undefined and empty kinds are a safe no-op", func(t *testing.T) { + t.Cleanup(func() { ClearGraph(t, db, ctx) }) createFixture() if err := deleter.DeleteRelationshipsByKinds(ctx, graph.Kinds{missing}); err != nil { @@ -145,6 +146,5 @@ func TestPostgreSQLDeleteRelationshipsByKinds(t *testing.T) { t.Fatalf("edge count after no-op delete: got %d, want 3", count) } - ClearGraph(t, db, ctx) }) } diff --git a/integration/wipe_graph_test.go b/integration/wipe_graph_test.go index 0e1390f2..71a7454b 100644 --- a/integration/wipe_graph_test.go +++ b/integration/wipe_graph_test.go @@ -70,6 +70,7 @@ func TestWipeGraph(t *testing.T) { } t.Run("truncates all partitions and atomically retains the recreated node", func(t *testing.T) { + session.RegisterGraphCleanup(t) session.ClearGraph(t) seed(t) @@ -101,6 +102,7 @@ func TestWipeGraph(t *testing.T) { }) t.Run("rolls back the truncate when the retain delegate fails", func(t *testing.T) { + session.RegisterGraphCleanup(t) session.ClearGraph(t) seed(t) @@ -118,6 +120,7 @@ func TestWipeGraph(t *testing.T) { }) t.Run("truncates without a retain delegate", func(t *testing.T) { + session.RegisterGraphCleanup(t) session.ClearGraph(t) seed(t)