feat(azure): add WithClient variants and azfake unit tests for all modules

Add exported WithClient variants for Azure module functions that accept
a pre-built SDK client, enabling unit testing with fake servers. Each
existing ContextE function now delegates its SDK call to the
corresponding WithClient variant.

Pattern established across all modules:
- GetDiskContextE(ctx, name, rg, sub) -- convenience, creates its own client
- GetDiskWithClient(ctx, client, rg, name) -- testable, accepts injected client

Functions that extract fields from response objects (not clients) drop
the WithClient suffix since they don't take a client:
- GetAvailabilitySetFaultDomainCount(avs)
- GetLoadBalancerFrontendIPConfigNames(lb)
- GetNetworkInterfacePrivateIPs(nic)
- GetIPOfPublicIPAddressByName(pip)
- GetVirtualNetworkDNSServerIPs(vnet)

Nil-deref guards added on all newly-exported helpers that dereference
response properties (avs.Properties, lb.Properties, nic.Properties,
feConfig.Properties, vnet.Properties, pip.Name).

Modules updated: disk, availabilityset, loadbalancer, networkinterface,
publicaddress, virtualnetwork, cosmosdb, servicebus

All WithClient functions are tested using Azure SDK's azfake framework
(httptest for servicebus beta SDK), running in CI without credentials.

Author: james00012

modules/azure/availabilityset.go

@@ -2,6 +2,7 @@ package azure
 
 import (
 	"context"
+	"errors"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
@@ -94,12 +95,21 @@ func CheckAvailabilitySetContainsVMContextE(t testing.TestingT, ctx context.Cont
 		return false, err
 	}
 
+	return CheckAvailabilitySetContainsVMWithClient(ctx, client, resGroupName, avsName, vmName)
+}
+
+// CheckAvailabilitySetContainsVMWithClient checks if the Virtual Machine is contained in the Availability Set VMs
+// using the provided AvailabilitySetsClient.
+func CheckAvailabilitySetContainsVMWithClient(ctx context.Context, client *armcompute.AvailabilitySetsClient, resGroupName string, avsName string, vmName string) (bool, error) {
 	resp, err := client.Get(ctx, resGroupName, avsName, nil)
 	if err != nil {
 		return false, err
 	}
 
 	for _, vm := range resp.Properties.VirtualMachines {
+		if vm.ID == nil {
+			continue
+		}
 		// VM IDs are always ALL CAPS in this property so ignoring case
 		if strings.EqualFold(vmName, GetNameFromResourceID(*vm.ID)) {
 			return true, nil
@@ -148,6 +158,12 @@ func GetAvailabilitySetVMNamesInCapsContextE(t testing.TestingT, ctx context.Con
 		return nil, err
 	}
 
+	return GetAvailabilitySetVMNamesInCapsWithClient(ctx, client, resGroupName, avsName)
+}
+
+// GetAvailabilitySetVMNamesInCapsWithClient gets a list of VM names in the specified Azure Availability Set
+// using the provided AvailabilitySetsClient.
+func GetAvailabilitySetVMNamesInCapsWithClient(ctx context.Context, client *armcompute.AvailabilitySetsClient, resGroupName string, avsName string) ([]string, error) {
 	resp, err := client.Get(ctx, resGroupName, avsName, nil)
 	if err != nil {
 		return nil, err
@@ -156,6 +172,9 @@ func GetAvailabilitySetVMNamesInCapsContextE(t testing.TestingT, ctx context.Con
 	vms := []string{}
 
 	for _, vm := range resp.Properties.VirtualMachines {
+		if vm.ID == nil {
+			continue
+		}
 		// IDs are returned in ALL CAPS for this property
 		if vmName := GetNameFromResourceID(*vm.ID); len(vmName) > 0 {
 			vms = append(vms, vmName)
@@ -204,6 +223,15 @@ func GetAvailabilitySetFaultDomainCountContextE(t testing.TestingT, ctx context.
 		return -1, err
 	}
 
+	return ExtractAvailabilitySetFaultDomainCount(avs)
+}
+
+// ExtractAvailabilitySetFaultDomainCount gets the Fault Domain Count from the provided AvailabilitySet.
+func ExtractAvailabilitySetFaultDomainCount(avs *armcompute.AvailabilitySet) (int32, error) {
+	if avs.Properties == nil || avs.Properties.PlatformFaultDomainCount == nil {
+		return -1, errors.New("availability set has no fault domain count")
+	}
+
 	return *avs.Properties.PlatformFaultDomainCount, nil
 }
 
@@ -229,6 +257,11 @@ func GetAvailabilitySetContextE(t testing.TestingT, ctx context.Context, avsName
 		return nil, err
 	}
 
+	return GetAvailabilitySetWithClient(ctx, client, resGroupName, avsName)
+}
+
+// GetAvailabilitySetWithClient gets an Availability Set using the provided AvailabilitySetsClient.
+func GetAvailabilitySetWithClient(ctx context.Context, client *armcompute.AvailabilitySetsClient, resGroupName string, avsName string) (*armcompute.AvailabilitySet, error) {
 	resp, err := client.Get(ctx, resGroupName, avsName, nil)
 	if err != nil {
 		return nil, err

modules/azure/availabilityset_withclient_test.go

@@ -0,0 +1,149 @@
+package azure_test
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	azfake "github.com/Azure/azure-sdk-for-go/sdk/azcore/fake"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
+	computefake "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6/fake"
+	"github.com/gruntwork-io/terratest/modules/azure"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func newFakeAvailabilitySetsClient(t *testing.T, srv *computefake.AvailabilitySetsServer) *armcompute.AvailabilitySetsClient {
+	t.Helper()
+
+	transport := computefake.NewAvailabilitySetsServerTransport(srv)
+	client, err := armcompute.NewAvailabilitySetsClient("fake-sub", &azfake.TokenCredential{}, &arm.ClientOptions{
+		ClientOptions: policy.ClientOptions{Transport: transport},
+	})
+	require.NoError(t, err)
+
+	return client
+}
+
+func fakeAvsGetHandler(avsName string, vmIDs []string, faultDomainCount int32) func(context.Context, string, string, *armcompute.AvailabilitySetsClientGetOptions) (azfake.Responder[armcompute.AvailabilitySetsClientGetResponse], azfake.ErrorResponder) {
+	return func(_ context.Context, _ string, _ string, _ *armcompute.AvailabilitySetsClientGetOptions) (resp azfake.Responder[armcompute.AvailabilitySetsClientGetResponse], errResp azfake.ErrorResponder) {
+		vms := make([]*armcompute.SubResource, len(vmIDs))
+		for i, id := range vmIDs {
+			vms[i] = &armcompute.SubResource{ID: to.Ptr(id)}
+		}
+
+		resp.SetResponse(http.StatusOK, armcompute.AvailabilitySetsClientGetResponse{
+			AvailabilitySet: armcompute.AvailabilitySet{
+				Name: to.Ptr(avsName),
+				Properties: &armcompute.AvailabilitySetProperties{
+					VirtualMachines:          vms,
+					PlatformFaultDomainCount: to.Ptr(faultDomainCount),
+				},
+			},
+		}, nil)
+
+		return
+	}
+}
+
+func TestGetAvailabilitySetWithClient(t *testing.T) {
+	t.Parallel()
+
+	srv := &computefake.AvailabilitySetsServer{
+		Get: fakeAvsGetHandler("my-avs", nil, 2),
+	}
+	client := newFakeAvailabilitySetsClient(t, srv)
+
+	avs, err := azure.GetAvailabilitySetWithClient(t.Context(), client, "rg", "my-avs")
+	require.NoError(t, err)
+	assert.Equal(t, "my-avs", *avs.Name)
+}
+
+func TestCheckAvailabilitySetContainsVMWithClient(t *testing.T) {
+	t.Parallel()
+
+	vmIDs := []string{
+		"/subscriptions/sub/resourceGroups/RG/providers/Microsoft.Compute/virtualMachines/VM-ONE",
+		"/subscriptions/sub/resourceGroups/RG/providers/Microsoft.Compute/virtualMachines/VM-TWO",
+	}
+
+	tests := []struct {
+		name    string
+		vmName  string
+		found   bool
+		wantErr bool
+	}{
+		{name: "exact case match", vmName: "VM-ONE", found: true},
+		{name: "case insensitive match", vmName: "vm-one", found: true},
+		{name: "not found", vmName: "vm-three", found: false, wantErr: true},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			srv := &computefake.AvailabilitySetsServer{
+				Get: fakeAvsGetHandler("avs", vmIDs, 2),
+			}
+			client := newFakeAvailabilitySetsClient(t, srv)
+
+			found, err := azure.CheckAvailabilitySetContainsVMWithClient(t.Context(), client, "rg", "avs", tc.vmName)
+
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			assert.Equal(t, tc.found, found)
+		})
+	}
+}
+
+func TestGetAvailabilitySetVMNamesInCapsWithClient(t *testing.T) {
+	t.Parallel()
+
+	vmIDs := []string{
+		"/subscriptions/sub/resourceGroups/RG/providers/Microsoft.Compute/virtualMachines/VM-ALPHA",
+		"/subscriptions/sub/resourceGroups/RG/providers/Microsoft.Compute/virtualMachines/VM-BETA",
+	}
+
+	srv := &computefake.AvailabilitySetsServer{
+		Get: fakeAvsGetHandler("avs", vmIDs, 3),
+	}
+	client := newFakeAvailabilitySetsClient(t, srv)
+
+	names, err := azure.GetAvailabilitySetVMNamesInCapsWithClient(t.Context(), client, "rg", "avs")
+	require.NoError(t, err)
+	assert.Equal(t, []string{"VM-ALPHA", "VM-BETA"}, names)
+}
+
+func TestExtractAvailabilitySetFaultDomainCount(t *testing.T) {
+	t.Parallel()
+
+	t.Run("valid", func(t *testing.T) {
+		t.Parallel()
+
+		avs := &armcompute.AvailabilitySet{
+			Properties: &armcompute.AvailabilitySetProperties{
+				PlatformFaultDomainCount: to.Ptr[int32](3),
+			},
+		}
+
+		count, err := azure.ExtractAvailabilitySetFaultDomainCount(avs)
+		require.NoError(t, err)
+		assert.Equal(t, int32(3), count)
+	})
+
+	t.Run("nil properties", func(t *testing.T) {
+		t.Parallel()
+
+		avs := &armcompute.AvailabilitySet{}
+
+		_, err := azure.ExtractAvailabilitySetFaultDomainCount(avs)
+		require.Error(t, err)
+	})
+}

modules/azure/cosmosdb.go

@@ -73,13 +73,16 @@ func GetCosmosDBAccountContextE(ctx context.Context, subscriptionID string, reso
 		return nil, err
 	}
 
-	// Get the corresponding database account
-	resp, err := cosmosClient.Get(ctx, resourceGroupName, accountName, nil)
+	return GetCosmosDBAccountWithClient(ctx, cosmosClient, resourceGroupName, accountName)
+}
+
+// GetCosmosDBAccountWithClient gets a database account using the provided DatabaseAccountsClient.
+func GetCosmosDBAccountWithClient(ctx context.Context, client *armcosmos.DatabaseAccountsClient, resourceGroupName string, accountName string) (*armcosmos.DatabaseAccountGetResults, error) {
+	resp, err := client.Get(ctx, resourceGroupName, accountName, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	// Return DB
 	return &resp.DatabaseAccountGetResults, nil
 }
 
@@ -148,13 +151,16 @@ func GetCosmosDBSQLDatabaseContextE(ctx context.Context, subscriptionID string,
 		return nil, err
 	}
 
-	// Get the corresponding database
-	resp, err := cosmosClient.GetSQLDatabase(ctx, resourceGroupName, accountName, databaseName, nil)
+	return GetCosmosDBSQLDatabaseWithClient(ctx, cosmosClient, resourceGroupName, accountName, databaseName)
+}
+
+// GetCosmosDBSQLDatabaseWithClient gets a SQL database using the provided SQLResourcesClient.
+func GetCosmosDBSQLDatabaseWithClient(ctx context.Context, client *armcosmos.SQLResourcesClient, resourceGroupName string, accountName string, databaseName string) (*armcosmos.SQLDatabaseGetResults, error) {
+	resp, err := client.GetSQLDatabase(ctx, resourceGroupName, accountName, databaseName, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	// Return DB
 	return &resp.SQLDatabaseGetResults, nil
 }
 
@@ -192,13 +198,16 @@ func GetCosmosDBSQLContainerContextE(ctx context.Context, subscriptionID string,
 		return nil, err
 	}
 
-	// Get the corresponding SQL container
-	resp, err := cosmosClient.GetSQLContainer(ctx, resourceGroupName, accountName, databaseName, containerName, nil)
+	return GetCosmosDBSQLContainerWithClient(ctx, cosmosClient, resourceGroupName, accountName, databaseName, containerName)
+}
+
+// GetCosmosDBSQLContainerWithClient gets a SQL container using the provided SQLResourcesClient.
+func GetCosmosDBSQLContainerWithClient(ctx context.Context, client *armcosmos.SQLResourcesClient, resourceGroupName string, accountName string, databaseName string, containerName string) (*armcosmos.SQLContainerGetResults, error) {
+	resp, err := client.GetSQLContainer(ctx, resourceGroupName, accountName, databaseName, containerName, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	// Return container
 	return &resp.SQLContainerGetResults, nil
 }
 
@@ -236,13 +245,17 @@ func GetCosmosDBSQLDatabaseThroughputContextE(ctx context.Context, subscriptionI
 		return nil, err
 	}
 
-	// Get the corresponding database throughput config
-	resp, err := cosmosClient.GetSQLDatabaseThroughput(ctx, resourceGroupName, accountName, databaseName, nil)
+	return GetCosmosDBSQLDatabaseThroughputWithClient(ctx, cosmosClient, resourceGroupName, accountName, databaseName)
+}
+
+// GetCosmosDBSQLDatabaseThroughputWithClient gets a SQL database throughput configuration
+// using the provided SQLResourcesClient.
+func GetCosmosDBSQLDatabaseThroughputWithClient(ctx context.Context, client *armcosmos.SQLResourcesClient, resourceGroupName string, accountName string, databaseName string) (*armcosmos.ThroughputSettingsGetResults, error) {
+	resp, err := client.GetSQLDatabaseThroughput(ctx, resourceGroupName, accountName, databaseName, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	// Return throughput config
 	return &resp.ThroughputSettingsGetResults, nil
 }
 
@@ -280,12 +293,16 @@ func GetCosmosDBSQLContainerThroughputContextE(ctx context.Context, subscription
 		return nil, err
 	}
 
-	// Get the corresponding container throughput config
-	resp, err := cosmosClient.GetSQLContainerThroughput(ctx, resourceGroupName, accountName, databaseName, containerName, nil)
+	return GetCosmosDBSQLContainerThroughputWithClient(ctx, cosmosClient, resourceGroupName, accountName, databaseName, containerName)
+}
+
+// GetCosmosDBSQLContainerThroughputWithClient gets a SQL container throughput configuration
+// using the provided SQLResourcesClient.
+func GetCosmosDBSQLContainerThroughputWithClient(ctx context.Context, client *armcosmos.SQLResourcesClient, resourceGroupName string, accountName string, databaseName string, containerName string) (*armcosmos.ThroughputSettingsGetResults, error) {
+	resp, err := client.GetSQLContainerThroughput(ctx, resourceGroupName, accountName, databaseName, containerName, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	// Return throughput config
 	return &resp.ThroughputSettingsGetResults, nil
 }