test: modifying fake to only initialize the query once

Azure · Jan 28, 2025 · 38666ed · 38666ed
1 parent 9d49d25
commit 38666ed
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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)