From 5a0c291ec4208f4b43a623eb9ecf6604776622cc Mon Sep 17 00:00:00 2001 From: sirisjo Date: Tue, 24 Feb 2026 12:05:36 -0800 Subject: [PATCH 1/4] insert kinds if not found --- packages/go/schemagen/generator/sql.go | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/packages/go/schemagen/generator/sql.go b/packages/go/schemagen/generator/sql.go index 2a6fd9e5294..e61ecd6a3df 100644 --- a/packages/go/schemagen/generator/sql.go +++ b/packages/go/schemagen/generator/sql.go @@ -177,26 +177,25 @@ func GenerateExtensionSQL(name string, displayName string, version string, names sb.WriteString(fmt.Sprintf("-- Code generated by Cuelang code gen. DO NOT EDIT!\n-- Cuelang source: %s/", SchemaSourceName)) sb.WriteString(` -CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS void AS $$ +CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$ BEGIN IF NOT EXISTS (SELECT id FROM kind WHERE kind.name = node_kind_name) THEN - INSERT INTO kind (name) VALUES (node_kind_name); + INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id; END IF; END $$ LANGUAGE plpgsql; `) sb.WriteString(` -CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS void AS $$ +CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS SMALLINT AS $$ DECLARE retrieved_kind_id SMALLINT; BEGIN SELECT k.id INTO retrieved_kind_id FROM kind k WHERE k.name = kind_name; IF retrieved_kind_id IS NULL THEN - INSERT INTO kind (name) VALUES (kind_name) - RETURNING id INTO retrieved_kind_id; + SELECT genscript_upsert_kind(kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id) THEN - INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id); + INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id; END IF; END $$ LANGUAGE plpgsql; `) @@ -208,7 +207,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_node_kinds nk WHERE nk.kind_id = retrieved_kind_id) THEN @@ -227,7 +226,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_relationship_kinds ek WHERE ek.kind_id = retrieved_kind_id) THEN @@ -248,12 +247,12 @@ DECLARE BEGIN SELECT id INTO retrieved_environment_kind_id FROM kind WHERE name = v_environment_kind_name; IF retrieved_environment_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_environment_kind_name) INTO retrieved_environment_kind_id; END IF; SELECT sk.id INTO retrieved_source_kind_id FROM source_kinds sk JOIN kind k ON sk.kind_id = k.id WHERE k.name = v_source_kind_name; IF retrieved_source_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_source_kind(v_source_kind_name) INTO retrieved_source_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN @@ -274,7 +273,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_principal_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_principal_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT 1 FROM schema_environments_principal_kinds pk WHERE pk.principal_kind = retrieved_kind_id) THEN From eadd6fc3a366aadd443653db7dd3ad7b319e1d21 Mon Sep 17 00:00:00 2001 From: sirisjo Date: Tue, 24 Feb 2026 12:06:01 -0800 Subject: [PATCH 2/4] add extension integration tests --- .../migration/extensions/ad_graph_schema.sql | 32 +-- .../migration/extensions/az_graph_schema.sql | 32 +-- .../migration/extensions_integration_test.go | 185 ++++++++++++++++++ packages/go/schemagen/generator/sql.go | 15 +- 4 files changed, 240 insertions(+), 24 deletions(-) create mode 100644 cmd/api/src/database/migration/extensions_integration_test.go diff --git a/cmd/api/src/database/migration/extensions/ad_graph_schema.sql b/cmd/api/src/database/migration/extensions/ad_graph_schema.sql index ff1e93ba2ad..b6c5cff22b1 100644 --- a/cmd/api/src/database/migration/extensions/ad_graph_schema.sql +++ b/cmd/api/src/database/migration/extensions/ad_graph_schema.sql @@ -15,25 +15,35 @@ -- SPDX-License-Identifier: Apache-2.0 -- Code generated by Cuelang code gen. DO NOT EDIT! -- Cuelang source: github.com/specterops/bloodhound/-/tree/main/packages/cue/schemas/ -CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS void AS $$ +CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$ +DECLARE + kind_id SMALLINT; BEGIN IF NOT EXISTS (SELECT id FROM kind WHERE kind.name = node_kind_name) THEN - INSERT INTO kind (name) VALUES (node_kind_name); + INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id INTO kind_id; + ELSE + SELECT id FROM kind WHERE kind.name = node_kind_name INTO kind_id; END IF; + + RETURN kind_id; END $$ LANGUAGE plpgsql; -CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS void AS $$ +CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS SMALLINT AS $$ DECLARE retrieved_kind_id SMALLINT; + source_kind_id SMALLINT; BEGIN SELECT k.id INTO retrieved_kind_id FROM kind k WHERE k.name = kind_name; IF retrieved_kind_id IS NULL THEN - INSERT INTO kind (name) VALUES (kind_name) - RETURNING id INTO retrieved_kind_id; + SELECT genscript_upsert_kind(kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id) THEN - INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id); + INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id INTO source_kind_id; + ELSE + SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id INTO source_kind_id; END IF; + + RETURN source_kind_id; END $$ LANGUAGE plpgsql; CREATE OR REPLACE FUNCTION genscript_upsert_schema_node_kind(v_extension_id INT, v_kind_name VARCHAR(256), v_display_name TEXT, v_description TEXT, v_is_display_kind BOOLEAN, v_icon TEXT, v_icon_color TEXT) RETURNS void AS $$ @@ -42,7 +52,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_node_kinds nk WHERE nk.kind_id = retrieved_kind_id) THEN @@ -59,7 +69,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_relationship_kinds ek WHERE ek.kind_id = retrieved_kind_id) THEN @@ -78,12 +88,12 @@ DECLARE BEGIN SELECT id INTO retrieved_environment_kind_id FROM kind WHERE name = v_environment_kind_name; IF retrieved_environment_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_environment_kind_name) INTO retrieved_environment_kind_id; END IF; SELECT sk.id INTO retrieved_source_kind_id FROM source_kinds sk JOIN kind k ON sk.kind_id = k.id WHERE k.name = v_source_kind_name; IF retrieved_source_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_source_kind(v_source_kind_name) INTO retrieved_source_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN @@ -102,7 +112,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_principal_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_principal_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT 1 FROM schema_environments_principal_kinds pk WHERE pk.principal_kind = retrieved_kind_id) THEN diff --git a/cmd/api/src/database/migration/extensions/az_graph_schema.sql b/cmd/api/src/database/migration/extensions/az_graph_schema.sql index 8d623cdd2d5..538e4dda78f 100644 --- a/cmd/api/src/database/migration/extensions/az_graph_schema.sql +++ b/cmd/api/src/database/migration/extensions/az_graph_schema.sql @@ -15,25 +15,35 @@ -- SPDX-License-Identifier: Apache-2.0 -- Code generated by Cuelang code gen. DO NOT EDIT! -- Cuelang source: github.com/specterops/bloodhound/-/tree/main/packages/cue/schemas/ -CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS void AS $$ +CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$ +DECLARE + kind_id SMALLINT; BEGIN IF NOT EXISTS (SELECT id FROM kind WHERE kind.name = node_kind_name) THEN - INSERT INTO kind (name) VALUES (node_kind_name); + INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id INTO kind_id; + ELSE + SELECT id FROM kind WHERE kind.name = node_kind_name INTO kind_id; END IF; + + RETURN kind_id; END $$ LANGUAGE plpgsql; -CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS void AS $$ +CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS SMALLINT AS $$ DECLARE retrieved_kind_id SMALLINT; + source_kind_id SMALLINT; BEGIN SELECT k.id INTO retrieved_kind_id FROM kind k WHERE k.name = kind_name; IF retrieved_kind_id IS NULL THEN - INSERT INTO kind (name) VALUES (kind_name) - RETURNING id INTO retrieved_kind_id; + SELECT genscript_upsert_kind(kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id) THEN - INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id); + INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id INTO source_kind_id; + ELSE + SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id INTO source_kind_id; END IF; + + RETURN source_kind_id; END $$ LANGUAGE plpgsql; CREATE OR REPLACE FUNCTION genscript_upsert_schema_node_kind(v_extension_id INT, v_kind_name VARCHAR(256), v_display_name TEXT, v_description TEXT, v_is_display_kind BOOLEAN, v_icon TEXT, v_icon_color TEXT) RETURNS void AS $$ @@ -42,7 +52,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_node_kinds nk WHERE nk.kind_id = retrieved_kind_id) THEN @@ -59,7 +69,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_relationship_kinds ek WHERE ek.kind_id = retrieved_kind_id) THEN @@ -78,12 +88,12 @@ DECLARE BEGIN SELECT id INTO retrieved_environment_kind_id FROM kind WHERE name = v_environment_kind_name; IF retrieved_environment_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_environment_kind_name) INTO retrieved_environment_kind_id; END IF; SELECT sk.id INTO retrieved_source_kind_id FROM source_kinds sk JOIN kind k ON sk.kind_id = k.id WHERE k.name = v_source_kind_name; IF retrieved_source_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_source_kind(v_source_kind_name) INTO retrieved_source_kind_id; END IF; IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN @@ -102,7 +112,7 @@ DECLARE BEGIN SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_principal_kind_name; IF retrieved_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + SELECT genscript_upsert_kind(v_principal_kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT 1 FROM schema_environments_principal_kinds pk WHERE pk.principal_kind = retrieved_kind_id) THEN diff --git a/cmd/api/src/database/migration/extensions_integration_test.go b/cmd/api/src/database/migration/extensions_integration_test.go new file mode 100644 index 00000000000..ed2be15e945 --- /dev/null +++ b/cmd/api/src/database/migration/extensions_integration_test.go @@ -0,0 +1,185 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// 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 + +//go:build integration + +package migration_test + +import ( + "context" + "fmt" + "net/url" + "strings" + "testing" + + "github.com/peterldowns/pgtestdb" + "gorm.io/gorm" + + "github.com/specterops/bloodhound/cmd/api/src/auth" + "github.com/specterops/bloodhound/cmd/api/src/database" + "github.com/specterops/bloodhound/cmd/api/src/model" + "github.com/specterops/bloodhound/cmd/api/src/test/integration/utils" + "github.com/stretchr/testify/require" +) + +type IntegrationTestSuite struct { + Context context.Context + BHDatabase *database.BloodhoundDB + DB *gorm.DB +} + +func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite { + t.Helper() + + var ( + ctx = context.Background() + connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{}) + ) + + // #region Setup for dbs + + gormDB, err := database.OpenDatabase(connConf.URL()) + require.NoError(t, err) + + db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver()) + + err = db.Migrate(ctx) + require.NoError(t, err) + + err = db.PopulateExtensionData(ctx) + require.NoError(t, err) + + // #endregion + + return IntegrationTestSuite{ + Context: ctx, + BHDatabase: db, + DB: gormDB, + } +} + +// getPostgresConfig reads key/value pairs from the default integration +// config file and creates a pgtestdb configuration object. +func getPostgresConfig(t *testing.T) pgtestdb.Config { + t.Helper() + + config, err := utils.LoadIntegrationTestConfig() + require.NoError(t, err) + + environmentMap := make(map[string]string) + for _, entry := range strings.Fields(config.Database.Connection) { + if parts := strings.SplitN(entry, "=", 2); len(parts) == 2 { + environmentMap[parts[0]] = parts[1] + } + } + + if strings.HasPrefix(environmentMap["host"], "/") { + return pgtestdb.Config{ + DriverName: "pgx", + User: environmentMap["user"], + Password: environmentMap["password"], + Database: environmentMap["dbname"], + Options: fmt.Sprintf("host=%s", url.PathEscape(environmentMap["host"])), + TestRole: &pgtestdb.Role{ + Username: environmentMap["user"], + Password: environmentMap["password"], + Capabilities: "NOSUPERUSER NOCREATEROLE", + }, + } + } + + return pgtestdb.Config{ + DriverName: "pgx", + Host: environmentMap["host"], + Port: environmentMap["port"], + User: environmentMap["user"], + Password: environmentMap["password"], + Database: environmentMap["dbname"], + Options: "sslmode=disable", + ForceTerminateConnections: true, + } +} + +func (suite *IntegrationTestSuite) teardownIntegrationTestSuite(t *testing.T) { + t.Helper() + + if suite.BHDatabase != nil { + suite.BHDatabase.Close(suite.Context) + } +} + +func TestExtensions_GetOnStartExtensionData(t *testing.T) { + var ( + testSuite = setupIntegrationTestSuite(t) + ) + + defer testSuite.teardownIntegrationTestSuite(t) + + err := testSuite.BHDatabase.PopulateExtensionData(testSuite.Context) + require.NoError(t, err) + + // Validate Both Schema Extensions Exist + extensions, totalRecords, err := testSuite.BHDatabase.GetGraphSchemaExtensions(testSuite.Context, model.Filters{}, model.Sort{}, 0, 0) + + require.NoError(t, err) + + require.Equal(t, 2, totalRecords) + + for _, extension := range extensions { + require.True(t, extension.IsBuiltin, "All extensions should be marked as built-in") + // Validate Schema Environments Exist + schemaEnvironments, err := testSuite.BHDatabase.GetEnvironmentsByExtensionId(testSuite.Context, extension.ID) + require.NoError(t, err) + + // There should only be one schema environment per built-in extension + require.Len(t, schemaEnvironments, 1) + schemaEnvironment := schemaEnvironments[0] + + // Validate Source Kinds Exist + sourceKind, err := testSuite.BHDatabase.GetSourceKindByID(testSuite.Context, int(schemaEnvironment.SourceKindId)) + require.NoError(t, err) + require.NotNil(t, sourceKind) + validateSourceKind(t, extension.Name, sourceKind.Name.String()) + + // Validate Environment Kinds Exist + environmentKind, err := testSuite.BHDatabase.GetKindById(testSuite.Context, schemaEnvironment.EnvironmentKindId) + require.NoError(t, err) + validateEnvironmentKind(t, extension.Name, environmentKind.Name) + } + +} + +func validateSourceKind(t *testing.T, extensionName, sourceKindName string) { + switch extensionName { + case "AD": + require.Equal(t, "Base", sourceKindName) + case "AZ": + require.Equal(t, "AZBase", sourceKindName) + default: + t.Error("Invalid extension name ", extensionName) + } +} + +func validateEnvironmentKind(t *testing.T, extensionName, environmentKindName string) { + switch extensionName { + case "AD": + require.Equal(t, "Domain", environmentKindName) + case "AZ": + require.Equal(t, "AZTenant", environmentKindName) + default: + t.Error("Invalid extension name ", extensionName) + } +} diff --git a/packages/go/schemagen/generator/sql.go b/packages/go/schemagen/generator/sql.go index e61ecd6a3df..c394caf7cfb 100644 --- a/packages/go/schemagen/generator/sql.go +++ b/packages/go/schemagen/generator/sql.go @@ -178,10 +178,16 @@ func GenerateExtensionSQL(name string, displayName string, version string, names sb.WriteString(` CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$ +DECLARE + kind_id SMALLINT; BEGIN IF NOT EXISTS (SELECT id FROM kind WHERE kind.name = node_kind_name) THEN - INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id; + INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id INTO kind_id; + ELSE + SELECT id FROM kind WHERE kind.name = node_kind_name INTO kind_id; END IF; + + RETURN kind_id; END $$ LANGUAGE plpgsql; `) @@ -189,14 +195,19 @@ END $$ LANGUAGE plpgsql; CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS SMALLINT AS $$ DECLARE retrieved_kind_id SMALLINT; + source_kind_id SMALLINT; BEGIN SELECT k.id INTO retrieved_kind_id FROM kind k WHERE k.name = kind_name; IF retrieved_kind_id IS NULL THEN SELECT genscript_upsert_kind(kind_name) INTO retrieved_kind_id; END IF; IF NOT EXISTS (SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id) THEN - INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id; + INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id INTO source_kind_id; + ELSE + SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id INTO source_kind_id; END IF; + + RETURN source_kind_id; END $$ LANGUAGE plpgsql; `) From 6d49c6ca65007ca8b5f48c703c9b20fa31325acc Mon Sep 17 00:00:00 2001 From: sirisjo Date: Tue, 24 Feb 2026 13:07:36 -0800 Subject: [PATCH 3/4] rabbit suggestions --- .../migration/extensions_integration_test.go | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/cmd/api/src/database/migration/extensions_integration_test.go b/cmd/api/src/database/migration/extensions_integration_test.go index ed2be15e945..e661da7cd1b 100644 --- a/cmd/api/src/database/migration/extensions_integration_test.go +++ b/cmd/api/src/database/migration/extensions_integration_test.go @@ -47,14 +47,17 @@ func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite { var ( ctx = context.Background() connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{}) + gormDB *gorm.DB + db *database.BloodhoundDB + err error ) // #region Setup for dbs - gormDB, err := database.OpenDatabase(connConf.URL()) + gormDB, err = database.OpenDatabase(connConf.URL()) require.NoError(t, err) - db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver()) + db = database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver()) err = db.Migrate(ctx) require.NoError(t, err) @@ -113,11 +116,11 @@ func getPostgresConfig(t *testing.T) pgtestdb.Config { } } -func (suite *IntegrationTestSuite) teardownIntegrationTestSuite(t *testing.T) { +func (s *IntegrationTestSuite) teardownIntegrationTestSuite(t *testing.T) { t.Helper() - if suite.BHDatabase != nil { - suite.BHDatabase.Close(suite.Context) + if s.BHDatabase != nil { + s.BHDatabase.Close(s.Context) } } From cb417eff9d1a3105ed972c0dff23c853568b6a7b Mon Sep 17 00:00:00 2001 From: sirisjo Date: Wed, 25 Feb 2026 09:54:19 -0800 Subject: [PATCH 4/4] change t.Error to t.Fatal --- .../src/database/migration/extensions_integration_test.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cmd/api/src/database/migration/extensions_integration_test.go b/cmd/api/src/database/migration/extensions_integration_test.go index e661da7cd1b..63fe99306d0 100644 --- a/cmd/api/src/database/migration/extensions_integration_test.go +++ b/cmd/api/src/database/migration/extensions_integration_test.go @@ -166,23 +166,25 @@ func TestExtensions_GetOnStartExtensionData(t *testing.T) { } func validateSourceKind(t *testing.T, extensionName, sourceKindName string) { + t.Helper() switch extensionName { case "AD": require.Equal(t, "Base", sourceKindName) case "AZ": require.Equal(t, "AZBase", sourceKindName) default: - t.Error("Invalid extension name ", extensionName) + t.Fatalf("Invalid extension name %s", extensionName) } } func validateEnvironmentKind(t *testing.T, extensionName, environmentKindName string) { + t.Helper() switch extensionName { case "AD": require.Equal(t, "Domain", environmentKindName) case "AZ": require.Equal(t, "AZTenant", environmentKindName) default: - t.Error("Invalid extension name ", extensionName) + t.Fatalf("Invalid extension name %s", extensionName) } }