Repository: Azure/ARO-RP
Pull request #1790 (2021-10-13 16:29:58)

Add SKU availability and restriction checks to dynamic validation
### Which issue this PR addresses:

<!--
Please include a link to the ADO work item as well as any GitHub issues.

Usage: `Fixes #<GitHub issue number>`, or `Fixes (paste link of issue)`.
-->
Fixes https://msazure.visualstudio.com/AzureRedHatOpenShift/_workitems/edit/11028028

### What this PR does / why we need it:

If an install is attempted using a VM SKU that is not available or restricted in the target install region, the install fails late in the process because the worker nodes never come up. This PR adds an API call and SKU checking to the dynamic validation process, so that unavailable SKUs or SKUs with subscription restrictions fail early in the install process.

Prerequisite for https://github.com/Azure/ARO-RP/pull/1751

<!--
Include a brief summary of what the PR is intended to accomplish and how the PR
does it. (2-3 sentences)
-->

### Test plan for issue:

Unit testing, and manual testing steps (for workers) below.

1) Using a SKU that is unavailable in the region (`Standard_L8s` in `eastus`, for example), comment out line 44 `FeatureRequireD2sV3Workers` in pkg/env/dev.go, and then hardcode the unavailable SKU in line 417 of pkg/util/cluster/cluster.go (`oc.Properties.WorkerProfiles[0].VMSize = "Standard_L8s"`). This will change the SKU that is used for a development cluster to `Standard_L8s`.
2) Hardcode the SKU value in pkg/api/validate/vm.go, and create a fake quota in pkg/api/validate/dynamic/quota.go. For example:
```
	case "Standard_L8s":
		requiredResources["standardLFamily"] += (count * 1)
		requiredResources["cores"] += (count * 1)
```
This prevents the install from failing during static validation.

3) Attempt to deploy a cluster using the hack script https://github.com/Azure/ARO-RP/blob/master/docs/deploy-development-rp.md and observe either a restrictions failure (if available but restricted in subscription) or an unavailable failure (if not available in region at all). 

<!--
How did you test that this PR works?

- Are there unit tests?
- Are there integration/e2e tests?
- If it is not possible to write automated tests, explain why and document how
  to manually test and verify the feature.
-->

### Is there any documentation that needs to be updated for this PR?

No
<!--
- If yes and the docs are in GitHub, include doc updates in the PR.
- If yes and the docs are not in GitHub (i.e. ADO wiki), include a link to the
  docs.
- If no, explain why (e.g. "tech debt cleanup, N/A").
-->



Patch:
diff --git a/pkg/api/validate/dynamic/dynamic.go b/pkg/api/validate/dynamic/dynamic.go
index af2029a86..8af3d815f 100644
--- a/pkg/api/validate/dynamic/dynamic.go
+++ b/pkg/api/validate/dynamic/dynamic.go
@@ -53,6 +53,7 @@ type Dynamic interface {
 	ValidateQuota(ctx context.Context, oc *api.OpenShiftCluster) error
 	ValidateDiskEncryptionSets(ctx context.Context, oc *api.OpenShiftCluster) error
 	ValidateEncryptionAtHost(ctx context.Context, oc *api.OpenShiftCluster) error
+	ValidateVMSku(ctx context.Context, location string, subscriptionID string, oc *api.OpenShiftCluster) error
 }
 
 type dynamic struct {
@@ -65,14 +66,17 @@ type dynamic struct {
 	providers          features.ProvidersClient
 	virtualNetworks    virtualNetworksGetClient
 	diskEncryptionSets compute.DiskEncryptionSetsClient
+	resourceSkusClient compute.ResourceSkusClient
 	spComputeUsage     compute.UsageClient
 	spNetworkUsage     network.UsageClient
 }
 
 type AuthorizerType string
 
-const AuthorizerFirstParty AuthorizerType = "resource provider"
-const AuthorizerClusterServicePrincipal AuthorizerType = "cluster"
+const (
+	AuthorizerFirstParty              AuthorizerType = "resource provider"
+	AuthorizerClusterServicePrincipal AuthorizerType = "cluster"
+)
 
 func NewValidator(log *logrus.Entry, env env.Interface, azEnv *azureclient.AROEnvironment, subscriptionID string, authorizer refreshable.Authorizer, authorizerType AuthorizerType) (Dynamic, error) {
 	return &dynamic{
@@ -87,6 +91,7 @@ func NewValidator(log *logrus.Entry, env env.Interface, azEnv *azureclient.AROEn
 		permissions:        authorization.NewPermissionsClient(azEnv, subscriptionID, authorizer),
 		virtualNetworks:    newVirtualNetworksCache(network.NewVirtualNetworksClient(azEnv, subscriptionID, authorizer)),
 		diskEncryptionSets: compute.NewDiskEncryptionSetsClient(azEnv, subscriptionID, authorizer),
+		resourceSkusClient: compute.NewResourceSkusClient(azEnv, subscriptionID, authorizer),
 	}, nil
 }
 
diff --git a/pkg/api/validate/dynamic/sku.go b/pkg/api/validate/dynamic/sku.go
new file mode 100644
index 000000000..ae4885c83
--- /dev/null
+++ b/pkg/api/validate/dynamic/sku.go
@@ -0,0 +1,63 @@
+package dynamic
+
+// Copyright (c) Microsoft Corporation.
+// Licensed under the Apache License 2.0.
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	mgmtcompute "github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-01/compute"
+
+	"github.com/Azure/ARO-RP/pkg/api"
+	"github.com/Azure/ARO-RP/pkg/util/computeskus"
+)
+
+// ValidateWorkerSku uses resourceSkusClient to ensure that the VM sizes listed in the cluster document are available for use in the target region.
+func (dv *dynamic) ValidateVMSku(ctx context.Context, location string, subscriptionID string, oc *api.OpenShiftCluster) error {
+	// Get a list of available worker SKUs, filtering by location. We initialize a new resourceSkusClient here instead of using the one in dv.env,
+	// so that we can determine SKU availability within target cluster subscription instead of within RP subscription.
+	filter := fmt.Sprintf("location eq %s", location)
+	skus, err := dv.resourceSkusClient.List(ctx, filter)
+	if err != nil {
+		return err
+	}
+
+	filteredSkus := computeskus.FilterVMSizes(skus, location)
+	masterProfileSku := string(oc.Properties.MasterProfile.VMSize)
+
+	err = checkSKUAvailability(filteredSkus, location, "properties.masterProfile.VMSize", masterProfileSku)
+	if err != nil {
+		return err
+	}
+
+	// In case there are multiple WorkerProfiles listed in the cluster document (such as post-install),
+	// compare VMSize in each WorkerProfile to the resourceSkusClient call above to ensure that the sku is available in region.
+	for i, workerprofile := range oc.Properties.WorkerProfiles {
+		workerProfileSku := string(workerprofile.VMSize)
+
+		err = checkSKUAvailability(filteredSkus, location, fmt.Sprintf("properties.workerProfiles[%d].VMSize", i), workerProfileSku)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func checkSKUAvailability(skus map[string]*mgmtcompute.ResourceSku, location, path, vmsize string) error {
+	// Ensure desired sku exists in target region
+	if skus[vmsize] == nil {
+		return api.NewCloudError(http.StatusBadRequest, api.CloudErrorCodeInvalidParameter, path, "The selected SKU '%v' is unavailable in region '%v'", vmsize, location)
+	}
+
+	// Fail if sku is available, but restricted within the subscription. Restrictions are subscription-specific.
+	// https://docs.microsoft.com/en-us/azure/azure-resource-manager/templates/error-sku-not-available
+	isRestricted := computeskus.IsRestricted(skus, location, vmsize)
+	if isRestricted {
+		return api.NewCloudError(http.StatusBadRequest, api.CloudErrorCodeInvalidParameter, path, "The selected SKU '%v' is restricted in region '%v' for selected subscription", vmsize, location)
+	}
+
+	return nil
+}
diff --git a/pkg/api/validate/openshiftcluster_validatedynamic.go b/pkg/api/validate/openshiftcluster_validatedynamic.go
index 0f7cabc01..4ecab714e 100644
--- a/pkg/api/validate/openshiftcluster_validatedynamic.go
+++ b/pkg/api/validate/openshiftcluster_validatedynamic.go
@@ -121,5 +121,10 @@ func (dv *openShiftClusterDynamicValidator) Dynamic(ctx context.Context) error {
 		return err
 	}
 
+	err = spDynamic.ValidateVMSku(ctx, dv.oc.Location, dv.subscriptionDoc.ID, dv.oc)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
diff --git a/pkg/env/prod.go b/pkg/env/prod.go
index e2d231cb5..e30d54eb6 100644
--- a/pkg/env/prod.go
+++ b/pkg/env/prod.go
@@ -22,6 +22,7 @@ import (
 	"github.com/Azure/ARO-RP/pkg/proxy"
 	"github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/compute"
 	"github.com/Azure/ARO-RP/pkg/util/clientauthorizer"
+	"github.com/Azure/ARO-RP/pkg/util/computeskus"
 	"github.com/Azure/ARO-RP/pkg/util/keyvault"
 	"github.com/Azure/ARO-RP/pkg/util/refreshable"
 	"github.com/Azure/ARO-RP/pkg/util/version"
@@ -274,29 +275,7 @@ func (p *prod) populateVMSkus(ctx context.Context, resourceSkusClient compute.Re
 		return err
 	}
 
-	p.vmskus = map[string]*mgmtcompute.ResourceSku{}
-	for _, sku := range skus {
-		// TODO(mjudeikis): At some point some SKU's stopped returning zones and
-		// locations. IcM is open with MSFT but this might take a while.
-		// Revert once we find out right behaviour.
-		// https://github.com/Azure/ARO-RP/issues/1515
-		if len(*sku.Locations) == 0 || !strings.EqualFold((*sku.Locations)[0], p.Location()) ||
-			*sku.ResourceType != "virtualMachines" {
-			continue
-		}
-
-		if len(*sku.LocationInfo) == 0 { // happened in eastus2euap
-			continue
-		}
-
-		// We copy only part of the object so we don't have to keep
-		// a lot of data in memory.
-		p.vmskus[*sku.Name] = &mgmtcompute.ResourceSku{
-			Name:         sku.Name,
-			LocationInfo: sku.LocationInfo,
-			Capabilities: sku.Capabilities,
-		}
-	}
+	p.vmskus = computeskus.FilterVMSizes(skus, p.Location())
 
 	return nil
 }
diff --git a/pkg/util/computeskus/computeskus.go b/pkg/util/computeskus/computeskus.go
index 7ea6a69fb..5bd534c79 100644
--- a/pkg/util/computeskus/computeskus.go
+++ b/pkg/util/computeskus/computeskus.go
@@ -4,6 +4,8 @@ package computeskus
 // Licensed under the Apache License 2.0.
 
 import (
+	"strings"
+
 	mgmtcompute "github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-01/compute"
 )
 
@@ -18,7 +20,7 @@ func Zones(sku *mgmtcompute.ResourceSku) []string {
 	return *(*sku.LocationInfo)[0].Zones
 }
 
-// Capabilities checks whether given resource SKU has specific capability
+// HasCapability checks whether given resource SKU has specific capability
 func HasCapability(sku *mgmtcompute.ResourceSku, capabilityName string) bool {
 	if sku.Capabilities == nil {
 		return false
@@ -32,3 +34,46 @@ func HasCapability(sku *mgmtcompute.ResourceSku, capabilityName string) bool {
 
 	return false
 }
+
+// IsRestricted checks whether given resource SKU is restricted in a given location
+func IsRestricted(skus map[string]*mgmtcompute.ResourceSku, location, VMSize string) bool {
+	for _, restriction := range *skus[VMSize].Restrictions {
+		for _, restrictedLocation := range *restriction.RestrictionInfo.Locations {
+			if restrictedLocation == location {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+// FilterVMSizes filters resource SKU by location and returns only virtual machines, their names, restrictions, location info, and capabilities.
+func FilterVMSizes(skus []mgmtcompute.ResourceSku, location string) map[string]*mgmtcompute.ResourceSku {
+	vmskus := map[string]*mgmtcompute.ResourceSku{}
+	for _, sku := range skus {
+		// TODO(mjudeikis): At some point some SKU's stopped returning zones and
+		// locations. IcM is open with MSFT but this might take a while.
+		// Revert once we find out right behaviour.
+		// https://github.com/Azure/ARO-RP/issues/1515
+		if len(*sku.Locations) == 0 || !strings.EqualFold((*sku.Locations)[0], location) ||
+			*sku.ResourceType != "virtualMachines" {
+			continue
+		}
+
+		if len(*sku.LocationInfo) == 0 { // happened in eastus2euap
+			continue
+		}
+
+		// We copy only part of the object so we don't have to keep
+		// a lot of data in memory.
+		vmskus[*sku.Name] = &mgmtcompute.ResourceSku{
+			Name:         sku.Name,
+			Restrictions: sku.Restrictions,
+			LocationInfo: sku.LocationInfo,
+			Capabilities: sku.Capabilities,
+		}
+	}
+
+	return vmskus
+}


Test patch:
diff --git a/pkg/api/validate/dynamic/sku_test.go b/pkg/api/validate/dynamic/sku_test.go
new file mode 100644
index 000000000..1a31b013d
--- /dev/null
+++ b/pkg/api/validate/dynamic/sku_test.go
@@ -0,0 +1,172 @@
+package dynamic
+
+// Copyright (c) Microsoft Corporation.
+// Licensed under the Apache License 2.0.
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	mgmtcompute "github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-01/compute"
+	"github.com/Azure/go-autorest/autorest/to"
+	"github.com/golang/mock/gomock"
+	"github.com/sirupsen/logrus"
+
+	"github.com/Azure/ARO-RP/pkg/api"
+	mock_compute "github.com/Azure/ARO-RP/pkg/util/mocks/azureclient/mgmt/compute"
+)
+
+func TestValidateVMSku(t *testing.T) {
+	for _, tt := range []struct {
+		name                  string
+		restrictions          mgmtcompute.ResourceSkuRestrictionsReasonCode
+		restrictionLocation   *[]string
+		targetLocation        string
+		workerProfile1Sku     string
+		workerProfile2Sku     string
+		masterProfileSku      string
+		availableSku          string
+		restrictedSku         string
+		resourceSkusClientErr error
+		wantErr               string
+	}{
+		{
+			name:              "worker and master sku are valid",
+			workerProfile1Sku: "Standard_D4s_v2",
+			workerProfile2Sku: "Standard_D4s_v2",
+			masterProfileSku:  "Standard_D4s_v2",
+			availableSku:      "Standard_D4s_v2",
+		},
+		{
+			name:                  "unable to retrieve skus information",
+			workerProfile1Sku:     "Standard_D4s_v2",
+			workerProfile2Sku:     "Standard_D4s_v2",
+			resourceSkusClientErr: errors.New("unable to retrieve skus information"),
+			wantErr:               "unable to retrieve skus information",
+		},
+		{
+			name:              "desired worker sku doesn't exist in the target region",
+			workerProfile1Sku: "Standard_L80",
+			workerProfile2Sku: "Standard_L80",
+			masterProfileSku:  "Standard_D4s_v2",
+			availableSku:      "Standard_D4s_v2",
+			wantErr:           "400: InvalidParameter: properties.workerProfiles[0].VMSize: The selected SKU 'Standard_L80' is unavailable in region 'eastus'",
+		},
+		{
+			name:              "desired master sku doesn't exist in the target region",
+			workerProfile1Sku: "Standard_D4s_v2",
+			workerProfile2Sku: "Standard_D4s_v2",
+			masterProfileSku:  "Standard_L80",
+			availableSku:      "Standard_D4s_v2",
+			wantErr:           "400: InvalidParameter: properties.masterProfile.VMSize: The selected SKU 'Standard_L80' is unavailable in region 'eastus'",
+		},
+		{
+			name:              "one valid workerprofile and one invalid workerprofile",
+			workerProfile1Sku: "Standard_L80",
+			workerProfile2Sku: "Standard_D4s_v2",
+			masterProfileSku:  "Standard_D4s_v2",
+			availableSku:      "Standard_D4s_v2",
+			wantErr:           "400: InvalidParameter: properties.workerProfiles[0].VMSize: The selected SKU 'Standard_L80' is unavailable in region 'eastus'",
+		},
+		{
+			name:         "worker sku exists in region but is not available in subscription",
+			restrictions: mgmtcompute.NotAvailableForSubscription,
+			restrictionLocation: &[]string{
+				"eastus",
+			},
+			workerProfile1Sku: "Standard_L80",
+			workerProfile2Sku: "Standard_L80",
+			masterProfileSku:  "Standard_D4s_v2",
+			availableSku:      "Standard_D4s_v2",
+			restrictedSku:     "Standard_L80",
+			wantErr:           "400: InvalidParameter: properties.workerProfiles[0].VMSize: The selected SKU 'Standard_L80' is restricted in region 'eastus' for selected subscription",
+		},
+		{
+			name:         "master sku exists in region but is not available in subscription",
+			restrictions: mgmtcompute.NotAvailableForSubscription,
+			restrictionLocation: &[]string{
+				"eastus",
+			},
+			workerProfile1Sku: "Standard_D4s_v2",
+			workerProfile2Sku: "Standard_D4s_v2",
+			masterProfileSku:  "Standard_L80",
+			availableSku:      "Standard_D4s_v2",
+			restrictedSku:     "Standard_L80",
+			wantErr:           "400: InvalidParameter: properties.masterProfile.VMSize: The selected SKU 'Standard_L80' is restricted in region 'eastus' for selected subscription",
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.targetLocation == "" {
+				tt.targetLocation = "eastus"
+			}
+
+			controller := gomock.NewController(t)
+			defer controller.Finish()
+
+			oc := &api.OpenShiftCluster{
+				Properties: api.OpenShiftClusterProperties{
+					WorkerProfiles: []api.WorkerProfile{
+						{
+							VMSize: api.VMSize(tt.workerProfile1Sku),
+						},
+						{
+							VMSize: api.VMSize(tt.workerProfile2Sku),
+						},
+					},
+					MasterProfile: api.MasterProfile{
+						VMSize: api.VMSize(tt.masterProfileSku),
+					},
+				},
+			}
+
+			skus := []mgmtcompute.ResourceSku{
+				{
+					Name:      &tt.availableSku,
+					Locations: &[]string{"eastus"},
+					LocationInfo: &[]mgmtcompute.ResourceSkuLocationInfo{
+						{Zones: &[]string{"1, 2, 3"}},
+					},
+					Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{},
+					Capabilities: &[]mgmtcompute.ResourceSkuCapabilities{},
+					ResourceType: to.StringPtr("virtualMachines"),
+				},
+				{
+					Name:      &tt.restrictedSku,
+					Locations: &[]string{tt.targetLocation},
+					LocationInfo: &[]mgmtcompute.ResourceSkuLocationInfo{
+						{Zones: &[]string{"1, 2, 3"}},
+					},
+					Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{
+						{
+							ReasonCode: tt.restrictions,
+							RestrictionInfo: &mgmtcompute.ResourceSkuRestrictionInfo{
+								Locations: tt.restrictionLocation,
+							},
+						},
+					},
+					Capabilities: &[]mgmtcompute.ResourceSkuCapabilities{},
+					ResourceType: to.StringPtr("virtualMachines"),
+				},
+			}
+
+			resourceSkusClient := mock_compute.NewMockResourceSkusClient(controller)
+			resourceSkusClient.EXPECT().
+				List(gomock.Any(), fmt.Sprintf("location eq %v", tt.targetLocation)).
+				Return(skus, tt.resourceSkusClientErr)
+
+			dv := dynamic{
+				authorizerType:     AuthorizerClusterServicePrincipal,
+				log:                logrus.NewEntry(logrus.StandardLogger()),
+				resourceSkusClient: resourceSkusClient,
+			}
+
+			err := dv.ValidateVMSku(context.Background(), tt.targetLocation, subscriptionID, oc)
+			if err != nil && err.Error() != tt.wantErr ||
+				err == nil && tt.wantErr != "" {
+				t.Error(err)
+			}
+		})
+	}
+}
diff --git a/pkg/util/computeskus/computeskus_test.go b/pkg/util/computeskus/computeskus_test.go
index 014b11924..cfa7ef7da 100644
--- a/pkg/util/computeskus/computeskus_test.go
+++ b/pkg/util/computeskus/computeskus_test.go
@@ -102,3 +102,157 @@ func TestHasCapability(t *testing.T) {
 		})
 	}
 }
+
+func TestFilterVmSizes(t *testing.T) {
+	for _, tt := range []struct {
+		name             string
+		providedLocation string
+		resourceType     string
+		skuLocation      []string
+		skuRestrictions  mgmtcompute.ResourceSkuRestrictions
+		skuLocationInfo  []mgmtcompute.ResourceSkuLocationInfo
+		skuCapabilities  string
+		wantResult       map[string]*mgmtcompute.ResourceSku
+	}{
+		{
+			name:             "resource type is a virtual machine",
+			providedLocation: "eastus",
+			resourceType:     "virtualMachines",
+			skuRestrictions:  mgmtcompute.ResourceSkuRestrictions{ReasonCode: mgmtcompute.NotAvailableForSubscription},
+			skuLocation:      []string{"eastus"},
+			skuLocationInfo:  []mgmtcompute.ResourceSkuLocationInfo{{Zones: &[]string{"eastus-2"}}},
+			skuCapabilities:  "some-capability",
+
+			wantResult: map[string]*mgmtcompute.ResourceSku{
+				"Fake_Sku": {
+					Name: to.StringPtr("Fake_Sku"),
+					Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{{
+						ReasonCode: mgmtcompute.NotAvailableForSubscription}},
+					LocationInfo: &[]mgmtcompute.ResourceSkuLocationInfo{{
+						Zones: &[]string{"eastus-2"}},
+					},
+					Capabilities: &[]mgmtcompute.ResourceSkuCapabilities{{
+						Name: to.StringPtr("some-capability"),
+					}},
+				},
+			},
+		},
+		{
+			name:             "resource type not a virtual machine",
+			providedLocation: "eastus",
+			resourceType:     "disk",
+			skuLocation:      []string{"eastus"},
+			skuLocationInfo:  []mgmtcompute.ResourceSkuLocationInfo{{Zones: &[]string{"eastus-2"}}},
+			wantResult:       map[string]*mgmtcompute.ResourceSku{},
+		},
+		{
+			name:             "sku Location doesn't match provided location",
+			providedLocation: "mars",
+			resourceType:     "virtualMachines",
+			skuLocation:      []string{"eastus"},
+			skuLocationInfo:  []mgmtcompute.ResourceSkuLocationInfo{{Zones: &[]string{"eastus-2"}}},
+			wantResult:       map[string]*mgmtcompute.ResourceSku{},
+		},
+		{
+			name:             "sku Location has length of 0",
+			providedLocation: "eastus",
+			resourceType:     "virtualMachines",
+			skuLocation:      []string{},
+			skuLocationInfo:  []mgmtcompute.ResourceSkuLocationInfo{{Zones: &[]string{"eastus-2"}}},
+			wantResult:       map[string]*mgmtcompute.ResourceSku{},
+		},
+		{
+			name:             "sku LocationInfo has length of 0",
+			providedLocation: "eastus",
+			resourceType:     "virtualMachines",
+			skuLocation:      []string{"eastus"},
+			skuLocationInfo:  []mgmtcompute.ResourceSkuLocationInfo{},
+			wantResult:       map[string]*mgmtcompute.ResourceSku{},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+
+			sku := []mgmtcompute.ResourceSku{
+				{
+					Name: to.StringPtr("Fake_Sku"),
+					Capabilities: &[]mgmtcompute.ResourceSkuCapabilities{
+						{
+							Name: to.StringPtr(tt.skuCapabilities),
+						},
+					},
+					Locations:    &tt.skuLocation,
+					Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{tt.skuRestrictions},
+					LocationInfo: &tt.skuLocationInfo,
+					ResourceType: to.StringPtr(tt.resourceType),
+				},
+			}
+
+			result := FilterVMSizes(sku, tt.providedLocation)
+
+			if !reflect.DeepEqual(result, tt.wantResult) {
+				t.Error(cmp.Diff(result, tt.wantResult))
+			}
+		})
+	}
+}
+
+func TestIsRestricted(t *testing.T) {
+	for _, tt := range []struct {
+		name       string
+		location   string
+		vmsize     string
+		sku        map[string]*mgmtcompute.ResourceSku
+		wantResult bool
+	}{
+		{
+			name:     "sku is restricted in one location",
+			location: "eastus",
+			vmsize:   "Standard_Sku_1",
+			sku: map[string]*mgmtcompute.ResourceSku{
+				"Standard_Sku_1": {Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{
+					{
+						RestrictionInfo: &mgmtcompute.ResourceSkuRestrictionInfo{Locations: &[]string{"eastus"}},
+					},
+				}},
+			},
+			wantResult: true,
+		},
+		{
+			name:     "sku is restricted in multiple locations",
+			location: "eastus",
+			vmsize:   "Standard_Sku_1",
+			sku: map[string]*mgmtcompute.ResourceSku{
+				"Standard_Sku_1": {Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{
+					{
+						RestrictionInfo: &mgmtcompute.ResourceSkuRestrictionInfo{Locations: &[]string{
+							"eastus",
+							"eastus2",
+						}},
+					},
+				}},
+			},
+			wantResult: true,
+		},
+		{
+			name:     "sku is not restricted",
+			location: "eastus",
+			vmsize:   "Standard_Sku_2",
+			sku: map[string]*mgmtcompute.ResourceSku{
+				"Standard_Sku_2": {Restrictions: &[]mgmtcompute.ResourceSkuRestrictions{
+					{
+						RestrictionInfo: &mgmtcompute.ResourceSkuRestrictionInfo{Locations: &[]string{""}},
+					},
+				}},
+			},
+			wantResult: false,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			result := IsRestricted(tt.sku, tt.location, tt.vmsize)
+
+			if result != tt.wantResult {
+				t.Error(result)
+			}
+		})
+	}
+}