From 38666ed11efe95b9c73decfbee94b8204411a79e Mon Sep 17 00:00:00 2001
From: Bryce Soghigian <49734722+Bryce-Soghigian@users.noreply.github.com>
Date: Tue, 28 Jan 2025 19:21:49 +0000
Subject: [PATCH] test: modifying fake to only initialize the query once

---
 pkg/fake/azureresourcegraphapi.go      | 18 ++++++++++++++++--
 pkg/fake/azureresourcegraphapi_test.go |  2 +-
 pkg/test/environment.go                |  7 +------
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/pkg/fake/azureresourcegraphapi.go b/pkg/fake/azureresourcegraphapi.go
index 39e7c29bc..9eea36287 100644
--- a/pkg/fake/azureresourcegraphapi.go
+++ b/pkg/fake/azureresourcegraphapi.go
@@ -44,9 +44,23 @@ type AzureResourceGraphBehavior struct {
 var _ instance.AzureResourceGraphAPI = &AzureResourceGraphAPI{}
 
 type AzureResourceGraphAPI struct {
+	vmListQuery  string
+	nicListQuery string
 	AzureResourceGraphBehavior
 }
 
+func NewAzureResourceGraphAPI(resourceGroup string, virtualMachinesAPI *VirtualMachinesAPI, networkInterfacesAPI *NetworkInterfacesAPI) *AzureResourceGraphAPI {
+	return &AzureResourceGraphAPI{
+		vmListQuery:  instance.GetVMListQueryBuilder(resourceGroup).String(),
+		nicListQuery: instance.GetNICListQueryBuilder(resourceGroup).String(),
+		AzureResourceGraphBehavior: AzureResourceGraphBehavior{
+			VirtualMachinesAPI:   virtualMachinesAPI,
+			NetworkInterfacesAPI: networkInterfacesAPI,
+			ResourceGroup:        resourceGroup,
+		},
+	}
+}
+
 // Reset must be called between tests otherwise tests will pollute each other.
 func (c *AzureResourceGraphAPI) Reset() {}
 
@@ -68,7 +82,7 @@ func (c *AzureResourceGraphAPI) Resources(_ context.Context, query armresourcegr
 
 func (c *AzureResourceGraphAPI) getResourceList(query string) []interface{} {
 	switch query {
-	case instance.GetVMListQueryBuilder(c.ResourceGroup).String():
+	case c.vmListQuery:
 		vmList := lo.Filter(c.loadVMObjects(), func(vm armcompute.VirtualMachine, _ int) bool {
 			return vm.Tags != nil && vm.Tags[instance.NodePoolTagKey] != nil
 		})
@@ -77,7 +91,7 @@ func (c *AzureResourceGraphAPI) getResourceList(query string) []interface{} {
 			return convertBytesToInterface(b)
 		})
 		return resourceList
-	case instance.GetNICListQueryBuilder(c.ResourceGroup).String():
+	case c.nicListQuery:
 		nicList := lo.Filter(c.loadNicObjects(), func(nic armnetwork.Interface, _ int) bool {
 			return nic.Tags != nil && nic.Tags[instance.NodePoolTagKey] != nil
 		})
diff --git a/pkg/fake/azureresourcegraphapi_test.go b/pkg/fake/azureresourcegraphapi_test.go
index ebd7f1856..f5bc7b1ef 100644
--- a/pkg/fake/azureresourcegraphapi_test.go
+++ b/pkg/fake/azureresourcegraphapi_test.go
@@ -32,7 +32,7 @@ func TestAzureResourceGraphAPI_Resources_VM(t *testing.T) {
 	resourceGroup := "test_managed_cluster_rg"
 	subscriptionID := "test_sub"
 	virtualMachinesAPI := &VirtualMachinesAPI{}
-	azureResourceGraphAPI := &AzureResourceGraphAPI{AzureResourceGraphBehavior{VirtualMachinesAPI: virtualMachinesAPI, ResourceGroup: resourceGroup}}
+	azureResourceGraphAPI := NewAzureResourceGraphAPI(resourceGroup, virtualMachinesAPI, nil)
 	cases := []struct {
 		testName      string
 		vmNames       []string
diff --git a/pkg/test/environment.go b/pkg/test/environment.go
index b08eeb74f..bcf8ab48f 100644
--- a/pkg/test/environment.go
+++ b/pkg/test/environment.go
@@ -101,12 +101,7 @@ func NewRegionalEnvironment(ctx context.Context, env *coretest.Environment, regi
 	loadBalancersAPI := &fake.LoadBalancersAPI{}
 	nodeImageVersionsAPI := &fake.NodeImageVersionsAPI{}
 
-	azureResourceGraphAPI := &fake.AzureResourceGraphAPI{
-		AzureResourceGraphBehavior: fake.AzureResourceGraphBehavior{
-			VirtualMachinesAPI:   virtualMachinesAPI,
-			NetworkInterfacesAPI: networkInterfacesAPI,
-			ResourceGroup:        resourceGroup,
-		}}
+	azureResourceGraphAPI := fake.NewAzureResourceGraphAPI(resourceGroup, virtualMachinesAPI, networkInterfacesAPI)
 	// Cache
 	kubernetesVersionCache := cache.New(azurecache.KubernetesVersionTTL, azurecache.DefaultCleanupInterval)
 	instanceTypeCache := cache.New(instancetype.InstanceTypesCacheTTL, azurecache.DefaultCleanupInterval)