[Azure] Add sovereign cloud support for availabilityset module

modules/azure/availabilityset.go

@@ -39,7 +39,7 @@ func CheckAvailabilitySetContainsVM(t testing.TestingT, vmName string, avsName s
 
 // CheckAvailabilitySetContainsVME checks if the Virtual Machine is contained in the Availability Set VMs
 func CheckAvailabilitySetContainsVME(t testing.TestingT, vmName string, avsName string, resGroupName string, subscriptionID string) (bool, error) {
-	client, err := GetAvailabilitySetClientE(subscriptionID)
+	client, err := CreateAvailabilitySetClientE(subscriptionID)
 	if err != nil {
 		return false, err
 	}
@@ -71,7 +71,7 @@ func GetAvailabilitySetVMNamesInCaps(t testing.TestingT, avsName string, resGrou
 
 // GetAvailabilitySetVMNamesInCapsE gets a list of VM names in the specified Azure Availability Set
 func GetAvailabilitySetVMNamesInCapsE(t testing.TestingT, avsName string, resGroupName string, subscriptionID string) ([]string, error) {
-	client, err := GetAvailabilitySetClientE(subscriptionID)
+	client, err := CreateAvailabilitySetClientE(subscriptionID)
 	if err != nil {
 		return nil, err
 	}
@@ -120,7 +120,7 @@ func GetAvailabilitySetE(t testing.TestingT, avsName string, resGroupName string
 	}
 
 	// Get the client reference
-	client, err := GetAvailabilitySetClientE(subscriptionID)
+	client, err := CreateAvailabilitySetClientE(subscriptionID)
 	if err != nil {
 		return nil, err
 	}
@@ -135,6 +135,7 @@ func GetAvailabilitySetE(t testing.TestingT, avsName string, resGroupName string
 }
 
 // GetAvailabilitySetClientE gets a new Availability Set client in the specified Azure Subscription
+// TODO: remove in next version
 func GetAvailabilitySetClientE(subscriptionID string) (*compute.AvailabilitySetsClient, error) {
 	// Validate Azure subscription ID
 	subscriptionID, err := getTargetAzureSubscription(subscriptionID)

modules/azure/availabilityset_test.go

@@ -17,12 +17,12 @@ The below tests are currently stubbed out, with the expectation that they will t
 If/when methods to create and delete network resources are added, these tests can be extended.
 */
 
-func TestGetAvailabilitySetClientE(t *testing.T) {
+func TestCreateAvailabilitySetClientE(t *testing.T) {
 	t.Parallel()
 
 	subscriptionID := ""
 
-	client, err := GetAvailabilitySetClientE(subscriptionID)
+	client, err := CreateAvailabilitySetClientE(subscriptionID)
 
 	require.NoError(t, err)
 	assert.NotEmpty(t, *client)

modules/azure/client_factory.go

@@ -206,6 +206,33 @@ func CreateStorageBlobContainerClientE(subscriptionID string) (*storage.BlobCont
 	return &blobContainerClient, nil
 }
 
+// CreateAvailabilitySetClientE creates a new Availability Set client in the specified Azure Subscription
+func CreateAvailabilitySetClientE(subscriptionID string) (*compute.AvailabilitySetsClient, error) {
+	// Validate Azure subscription ID
+	subscriptionID, err := getTargetAzureSubscription(subscriptionID)
+	if err != nil {
+		return nil, err
+	}
+
+	// Lookup environment URI
+	baseURI, err := getBaseURI()
+	if err != nil {
+		return nil, err
+	}
+
+	// Get the Availability Set client
+	client := compute.NewAvailabilitySetsClientWithBaseURI(baseURI, subscriptionID)
+
+	// Create an authorizer
+	authorizer, err := NewAuthorizer()
+	if err != nil {
+		return nil, err
+	}
+	client.Authorizer = *authorizer
+
+	return &client, nil
+}
+
 // GetKeyVaultURISuffixE returns the proper KeyVault URI suffix for the configured Azure environment.
 // This function would fail the test if there is an error.
 func GetKeyVaultURISuffixE() (string, error) {