cleaning

elastic · Feb 3, 2025 · 9cf9a36 · 9cf9a36
1 parent 59c786b
commit 9cf9a36
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 53 deletions.
diff --git a/internal/inventory/gcpfetcher/fetcher_assets.go b/internal/inventory/gcpfetcher/fetcher_assets.go
@@ -85,45 +85,50 @@ func (f *assetsInventory) fetch(ctx context.Context, assetChan chan<- inventory.
 	}
 
 	for _, item := range gcpAssets {
-		asset := inventory.NewAssetEvent(
-			classification,
-			item.Name,
-			item.Name,
-			inventory.WithRawAsset(item),
-			inventory.WithRelatedAssetIds(
-				f.findRelatedAssetIds(classification.Type, item),
-			),
-			inventory.WithCloud(inventory.Cloud{
-				Provider:    inventory.GcpCloudProvider,
-				AccountID:   item.CloudAccount.AccountId,
-				AccountName: item.CloudAccount.AccountName,
-				ProjectID:   item.CloudAccount.OrganisationId,
-				ProjectName: item.CloudAccount.OrganizationName,
-				ServiceName: assetType,
-			}),
-			inventory.WithLabels(getAssetLabels(item)),
-			inventory.WithTags(getAssetTags(item)),
-		)
-		enrichAsset(&asset, item)
+		asset := getAssetEvent(assetType, classification, item)
 		assetChan <- asset
 	}
 }
 
-func (f *assetsInventory) findRelatedAssetIds(t inventory.AssetType, item *gcpinventory.ExtendedGcpAsset) []string {
+func getAssetEvent(assetType string, classification inventory.AssetClassification, item *gcpinventory.ExtendedGcpAsset) inventory.AssetEvent {
+	asset := inventory.NewAssetEvent(
+		classification,
+		item.Name,
+		item.Name,
+		inventory.WithRawAsset(item),
+		inventory.WithRelatedAssetIds(
+			findRelatedAssetIds(classification.Type, item),
+		),
+		inventory.WithCloud(inventory.Cloud{
+			Provider:    inventory.GcpCloudProvider,
+			AccountID:   item.CloudAccount.AccountId,
+			AccountName: item.CloudAccount.AccountName,
+			ProjectID:   item.CloudAccount.OrganisationId,
+			ProjectName: item.CloudAccount.OrganizationName,
+			ServiceName: assetType,
+		}),
+		inventory.WithLabels(getAssetLabels(item)),
+		inventory.WithTags(getAssetTags(item)),
+	)
+	enrichAsset(&asset, item)
+	return asset
+}
+
+func findRelatedAssetIds(t inventory.AssetType, item *gcpinventory.ExtendedGcpAsset) []string {
 	ids := []string{}
 	ids = append(ids, item.Ancestors...)
 	if item.Resource != nil {
 		ids = append(ids, item.Resource.Parent)
 	}
 
-	ids = append(ids, f.findRelatedAssetIdsForType(t, item)...)
+	ids = append(ids, findRelatedAssetIdsForType(t, item)...)
 
 	ids = lo.Compact(ids)
 	ids = lo.Uniq(ids)
 	return ids
 }
 
-func (f *assetsInventory) findRelatedAssetIdsForType(t inventory.AssetType, item *gcpinventory.ExtendedGcpAsset) []string {
+func findRelatedAssetIdsForType(t inventory.AssetType, item *gcpinventory.ExtendedGcpAsset) []string {
 	ids := []string{}
 
 	var fields map[string]*structpb.Value

diff --git a/internal/inventory/gcpfetcher/fetcher_assets_test.go b/internal/inventory/gcpfetcher/fetcher_assets_test.go
@@ -220,42 +220,28 @@ func TestAccountFetcher_EnrichAsset(t *testing.T) {
 				OrganizationName: "<org name>",
 			},
 		}
-
-		actual := inventory.NewAssetEvent(
-			r.classification,
-			gcpAsset.Name,
-			gcpAsset.Name,
-			inventory.WithRawAsset(gcpAsset),
-			inventory.WithRelatedAssetIds([]string{}),
-			inventory.WithCloud(inventory.Cloud{
-				Provider:    inventory.GcpCloudProvider,
-				AccountID:   gcpAsset.CloudAccount.AccountId,
-				AccountName: gcpAsset.CloudAccount.AccountName,
-				ProjectID:   gcpAsset.CloudAccount.OrganisationId,
-				ProjectName: gcpAsset.CloudAccount.OrganizationName,
-				ServiceName: r.assetType,
-			}))
-
+		actual := getAssetEvent(r.assetType, r.classification, gcpAsset)
 		expected := item.enrichments
-		expected.Event = actual.Event                 // Event is not set in the enrichments
-		expected.Entity = actual.Entity               // Entity is not set in the enrichments
-		expected.RawAttributes = actual.RawAttributes // RawAttributes is not set in the enrichments
 
-		// When there are no cloud fields enrichments, use the actual cloud fields
+		// Set the common fields that are not set in the enrichments
+		expected.Event = actual.Event
+		expected.Entity = actual.Entity
+		expected.RawAttributes = actual.RawAttributes
+
+		// Cloud is the only field where we have both common and enriched fields
 		if expected.Cloud == nil {
+			// Use the actual cloud fields when there are no cloud enrichments
 			expected.Cloud = actual.Cloud
+		} else {
+			// Use common cloud fields when there are cloud enrichments
+			expected.Cloud.Provider = actual.Cloud.Provider
+			expected.Cloud.AccountID = actual.Cloud.AccountID
+			expected.Cloud.AccountName = actual.Cloud.AccountName
+			expected.Cloud.ProjectID = actual.Cloud.ProjectID
+			expected.Cloud.ProjectName = actual.Cloud.ProjectName
+			expected.Cloud.ServiceName = actual.Cloud.ServiceName
 		}
 
-		enrichAsset(&actual, gcpAsset)
-
-		// Add or safely override common cloud fields not set in the enrichments
-		expected.Cloud.Provider = actual.Cloud.Provider
-		expected.Cloud.AccountID = actual.Cloud.AccountID
-		expected.Cloud.AccountName = actual.Cloud.AccountName
-		expected.Cloud.ProjectID = actual.Cloud.ProjectID
-		expected.Cloud.ProjectName = actual.Cloud.ProjectName
-		expected.Cloud.ServiceName = actual.Cloud.ServiceName
-
 		assert.Equalf(t, expected, actual, "%v failed", "EnrichAsset")
 	}
 }