chore: integration testing portability enhancement and schema cleanup (…

…#1024) * chore: introduce harness initializer that uses the GraphTestHarness interface for harness setup * chore: fix typo and some cleanup * NewADContainer * add logic to add tierzero property for test ADUsers * add nil check * add tenantID property when creating new azure tenant * chore: remove unimplemented edges from schema for better testing alignment --------- Co-authored-by: Brandon Shearin <[email protected]>
SpecterOps · Dec 18, 2024 · 51b29a4 · 51b29a4
1 parent 057d892
commit 51b29a4
Show file tree

Hide file tree

Showing 7 changed files with 263 additions and 261 deletions.
diff --git a/cmd/api/src/analysis/azure/queries_test.go b/cmd/api/src/analysis/azure/queries_test.go
@@ -33,6 +33,7 @@ import (
 
 func TestAnalysisAzure_GraphStats(t *testing.T) {
 	testCtx := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
+	testCtx.SetupAzure()
 	testCtx.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 
 		_, agg, err := azure2.GraphStats(context.TODO(), testCtx.Graph.Database)

diff --git a/cmd/api/src/queries/graph_integration_test.go b/cmd/api/src/queries/graph_integration_test.go
@@ -165,6 +165,7 @@ func TestGetEntityResults(t *testing.T) {
 	queryCache, err := cache.NewCache(cache.Config{MaxSize: 1})
 	require.Nil(t, err)
 
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		objectID, err := harness.InboundControl.ControlledUser.Properties.Get(common.ObjectID.String()).String()
 		require.Nil(t, err)
@@ -197,6 +198,7 @@ func TestGetEntityResults_QueryShorterThanSlowQueryThreshold(t *testing.T) {
 	queryCache, err := cache.NewCache(cache.Config{MaxSize: 1})
 	require.Nil(t, err)
 
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		objectID, err := harness.InboundControl.ControlledUser.Properties.Get(common.ObjectID.String()).String()
 		require.Nil(t, err)
@@ -230,6 +232,7 @@ func TestGetEntityResults_Cache(t *testing.T) {
 	queryCache, err := cache.NewCache(cache.Config{MaxSize: 2})
 	require.Nil(t, err)
 
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		objectID, err := harness.InboundControl.ControlledUser.Properties.Get(common.ObjectID.String()).String()
 		require.Nil(t, err)
@@ -270,6 +273,7 @@ func TestGetEntityResults_Cache(t *testing.T) {
 
 func TestGetAssetGroupComboNode(t *testing.T) {
 	testContext := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		graphQuery := queries.NewGraphQuery(db, cache.Cache{}, config.Configuration{})
 		comboNode, err := graphQuery.GetAssetGroupComboNode(context.Background(), "", ad.AdminTierZero)

diff --git a/cmd/api/src/test/integration/graph.go b/cmd/api/src/test/integration/graph.go
@@ -82,10 +82,14 @@ func (s *GraphTestContext) UpdateNode(node *graph.Node) {
 	})
 }
 
-func (s *GraphTestContext) DatabaseTest(dbDelegate func(harness HarnessDetails, db graph.Database)) {
-	s.setupActiveDirectory()
-	s.setupAzure()
+func (s *GraphTestContext) InitializeHarness(harness GraphTestHarness) {
+	s.Graph.WriteTransaction(s.testCtx, func(tx graph.Transaction) error {
+		harness.Setup(s)
+		return nil
+	})
+}
 
+func (s *GraphTestContext) DatabaseTest(dbDelegate func(harness HarnessDetails, db graph.Database)) {
 	dbDelegate(s.Harness, s.Graph.Database)
 }
 
@@ -109,8 +113,7 @@ func (s *GraphTestContext) DatabaseTestWithSetup(setup func(harness *HarnessDeta
 }
 
 func (s *GraphTestContext) BatchTest(batchDelegate func(harness HarnessDetails, batch graph.Batch), assertionDelegate func(details HarnessDetails, tx graph.Transaction)) {
-	s.setupActiveDirectory()
-	s.setupAzure()
+	s.SetupAzureAndActiveDirectory()
 
 	s.Graph.BatchOperation(s.testCtx, func(batch graph.Batch) error {
 		batchDelegate(s.Harness, batch)
@@ -124,8 +127,7 @@ func (s *GraphTestContext) BatchTest(batchDelegate func(harness HarnessDetails,
 }
 
 func (s *GraphTestContext) TransactionalTest(txDelegate func(harness HarnessDetails, tx graph.Transaction)) {
-	s.setupActiveDirectory()
-	s.setupAzure()
+	s.SetupAzureAndActiveDirectory()
 
 	s.Graph.WriteTransaction(s.testCtx, func(tx graph.Transaction) error {
 		txDelegate(s.Harness, tx)
@@ -319,6 +321,7 @@ func (s *GraphTestContext) NewAzureTenant(tenantID string) *graph.Node {
 	return s.NewNode(graph.AsProperties(graph.PropertyMap{
 		common.Name:     "New Tenant",
 		common.ObjectID: tenantID,
+		azure.TenantID:  tenantID,
 		azure.License:   "license",
 	}), azure.Entity, azure.Tenant)
 }
@@ -345,12 +348,27 @@ func (s *GraphTestContext) NewActiveDirectoryComputer(name, domainSID string) *g
 	}), ad.Entity, ad.Computer)
 }
 
-func (s *GraphTestContext) NewActiveDirectoryUser(name, domainSID string, isTierZero ...bool) *graph.Node {
+func (s *GraphTestContext) NewActiveDirectoryContainer(name, domainSID string) *graph.Node {
 	return s.NewNode(graph.AsProperties(graph.PropertyMap{
+		common.Name:     name,
+		common.ObjectID: must.NewUUIDv4().String(),
+		ad.DomainSID:    domainSID,
+	}), ad.Entity, ad.Container)
+}
+
+func (s *GraphTestContext) NewActiveDirectoryUser(name, domainSID string, isTierZero ...bool) *graph.Node {
+
+	propertyMap := graph.PropertyMap{
 		common.Name:     name,
 		common.ObjectID: strings.ToUpper(must.NewUUIDv4().String()),
 		ad.DomainSID:    domainSID,
-	}), ad.Entity, ad.User)
+	}
+
+	if isTierZero != nil && isTierZero[0] {
+		propertyMap[common.SystemTags] = ad.AdminTierZero
+	}
+
+	return s.NewNode(graph.AsProperties(propertyMap), ad.Entity, ad.User)
 }
 
 func (s *GraphTestContext) NewCustomActiveDirectoryUser(properties *graph.Properties) *graph.Node {
@@ -515,7 +533,12 @@ type CertTemplateData struct {
 	CertificatePolicy             []string
 }
 
-func (s *GraphTestContext) setupAzure() {
+func (s *GraphTestContext) SetupAzureAndActiveDirectory() {
+	s.SetupAzure()
+	s.SetupActiveDirectory()
+}
+
+func (s *GraphTestContext) SetupAzure() {
 	s.Harness.AZBaseHarness.Setup(s)
 	s.Harness.AZGroupMembership.Setup(s)
 	s.Harness.AZEntityPanelHarness.Setup(s)
@@ -530,7 +553,7 @@ func (s *GraphTestContext) setupAzure() {
 	s.Harness.AZManagementGroup.Setup(s)
 }
 
-func (s *GraphTestContext) setupActiveDirectory() {
+func (s *GraphTestContext) SetupActiveDirectory() {
 	// startServer a host of Tier Zero tagged assets
 	s.Harness.RootADHarness.Setup(s)