Repository: kubernetes-sigs/cloud-provider-azure
Pull request #665 (2021-06-17 02:22:48)

cherry-pick #646 and #661 to release-1.0
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/release.md#issue-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. Follow the instructions for writing a release note: https://git.k8s.io/community/contributors/guide/release-notes.md
6. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

**What type of PR is this?**
> Uncomment only one ` /kind <>` line, hit enter to put that in a new line, and remove leading whitespaces from that line:
>
> /kind api-change
> /kind bug
> /kind cleanup
> /kind design
> /kind documentation
> /kind failing-test
> /kind feature
> /kind flake
/kind bug

**What this PR does / why we need it**:

cherry-pick #646 and #661 to release-1.0.

**Which issue(s) this PR fixes**:
<!-- 
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #

**Special notes for your reviewer**:


**Release note**:
```
wait for the success of the initial run of syncRegionZonesMap and ensure outdated routes are cleanup when dualstack is enabled
```



Patch:
diff --git a/cmd/cloud-controller-manager/app/controllermanager.go b/cmd/cloud-controller-manager/app/controllermanager.go
index 002943748..28cdaa4fa 100644
--- a/cmd/cloud-controller-manager/app/controllermanager.go
+++ b/cmd/cloud-controller-manager/app/controllermanager.go
@@ -263,12 +263,12 @@ func Run(ctx context.Context, c *cloudcontrollerconfig.CompletedConfig) error {
 	if c.DynamicReloadingConfig.EnableDynamicReloading {
 		cloud, err = provider.NewCloudFromSecret(c.ClientBuilder, c.DynamicReloadingConfig.CloudConfigSecretName, c.DynamicReloadingConfig.CloudConfigSecretNamespace, c.DynamicReloadingConfig.CloudConfigKey)
 		if err != nil {
-			klog.Fatalf("Run: Cloud provider could not be initialized dynamically from secret %s/%s: %v", c.DynamicReloadingConfig.CloudConfigSecretNamespace, c.DynamicReloadingConfig.CloudConfigSecretName, err)
+			klog.Fatalf("Run: Cloud provider azure could not be initialized dynamically from secret %s/%s: %v", c.DynamicReloadingConfig.CloudConfigSecretNamespace, c.DynamicReloadingConfig.CloudConfigSecretName, err)
 		}
 	} else {
-		cloud, err = cloudprovider.InitCloudProvider(c.ComponentConfig.KubeCloudShared.CloudProvider.Name, c.ComponentConfig.KubeCloudShared.CloudProvider.CloudConfigFile)
+		cloud, err = provider.NewCloudFromConfigFile(c.ComponentConfig.KubeCloudShared.CloudProvider.CloudConfigFile, true)
 		if err != nil {
-			klog.Fatalf("Cloud provider could not be initialized: %v", err)
+			klog.Fatalf("Cloud provider azure could not be initialized: %v", err)
 		}
 	}
 
diff --git a/pkg/node/node.go b/pkg/node/node.go
index ee3ee0c15..a36127a1e 100644
--- a/pkg/node/node.go
+++ b/pkg/node/node.go
@@ -38,7 +38,7 @@ func NewIMDSNodeProvider() *IMDSNodeProvider {
 	az, err := azureprovider.NewCloud(bytes.NewBuffer([]byte(`{
 			"useInstanceMetadata": true,
 			"vmType": "vmss"
-		}`)))
+		}`)), false)
 	if err != nil {
 		klog.Fatalf("Failed to initialize Azure cloud provider: %v", err)
 	}
diff --git a/pkg/node/nodearm.go b/pkg/node/nodearm.go
index 58904599d..110a5a462 100644
--- a/pkg/node/nodearm.go
+++ b/pkg/node/nodearm.go
@@ -47,7 +47,7 @@ func NewARMNodeProvider(cloudConfigFilePath string) *ARMNodeProvider {
 		}
 		defer configFile.Close()
 
-		az, err = azureprovider.NewCloud(configFile)
+		az, err = azureprovider.NewCloud(configFile, false)
 
 		if err != nil {
 			klog.Fatalf("Failed to initialize Azure cloud provider: %v", err)
diff --git a/pkg/provider/azure.go b/pkg/provider/azure.go
index 87d0378f9..dbe954a28 100644
--- a/pkg/provider/azure.go
+++ b/pkg/provider/azure.go
@@ -22,6 +22,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"strings"
 	"sync"
 	"time"
@@ -326,13 +327,11 @@ func init() {
 		}
 	}
 	autorest.StatusCodesForRetry = statusCodesForRetry
-
-	cloudprovider.RegisterCloudProvider(consts.CloudProviderName, NewCloud)
 }
 
 // NewCloud returns a Cloud with initialized clients
-func NewCloud(configReader io.Reader) (cloudprovider.Interface, error) {
-	az, err := NewCloudWithoutFeatureGates(configReader)
+func NewCloud(configReader io.Reader, syncZones bool) (cloudprovider.Interface, error) {
+	az, err := NewCloudWithoutFeatureGates(configReader, syncZones)
 	if err != nil {
 		return nil, err
 	}
@@ -341,6 +340,38 @@ func NewCloud(configReader io.Reader) (cloudprovider.Interface, error) {
 	return az, nil
 }
 
+func NewCloudFromConfigFile(configFilePath string, syncZones bool) (cloudprovider.Interface, error) {
+	var (
+		cloud cloudprovider.Interface
+		err   error
+	)
+
+	if configFilePath != "" {
+		var config *os.File
+		config, err = os.Open(configFilePath)
+		if err != nil {
+			klog.Fatalf("Couldn't open cloud provider configuration %s: %#v",
+				configFilePath, err)
+		}
+
+		defer config.Close()
+		cloud, err = NewCloud(config, syncZones)
+	} else {
+		// Pass explicit nil so plugins can actually check for nil. See
+		// "Why is my nil error value not equal to nil?" in golang.org/doc/faq.
+		cloud, err = NewCloud(nil, false)
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("could not init cloud provider azure: %v", err)
+	}
+	if cloud == nil {
+		return nil, fmt.Errorf("nil cloud")
+	}
+
+	return cloud, nil
+}
+
 func NewCloudFromSecret(clientBuilder cloudprovider.ControllerClientBuilder, secretName, secretNamespace, cloudConfigKey string) (cloudprovider.Interface, error) {
 	az := &Cloud{
 		nodeNames:          sets.NewString(),
@@ -369,7 +400,7 @@ func NewCloudFromSecret(clientBuilder cloudprovider.ControllerClientBuilder, sec
 
 // NewCloudWithoutFeatureGates returns a Cloud without trying to wire the feature gates.  This is used by the unit tests
 // that don't load the actual features being used in the cluster.
-func NewCloudWithoutFeatureGates(configReader io.Reader) (*Cloud, error) {
+func NewCloudWithoutFeatureGates(configReader io.Reader, syncZones bool) (*Cloud, error) {
 	config, err := parseConfig(configReader)
 	if err != nil {
 		return nil, err
@@ -383,7 +414,7 @@ func NewCloudWithoutFeatureGates(configReader io.Reader) (*Cloud, error) {
 		routeCIDRs:         map[string]string{},
 	}
 
-	err = az.InitializeCloudFromConfig(config, false)
+	err = az.InitializeCloudFromConfig(config, false, syncZones)
 	if err != nil {
 		return nil, err
 	}
@@ -392,7 +423,7 @@ func NewCloudWithoutFeatureGates(configReader io.Reader) (*Cloud, error) {
 }
 
 // InitializeCloudFromConfig initializes the Cloud from config.
-func (az *Cloud) InitializeCloudFromConfig(config *Config, fromSecret bool) error {
+func (az *Cloud) InitializeCloudFromConfig(config *Config, fromSecret, syncZones bool) error {
 	if config == nil {
 		// should not reach here
 		return fmt.Errorf("InitializeCloudFromConfig: cannot initialize from nil config")
@@ -504,6 +535,34 @@ func (az *Cloud) InitializeCloudFromConfig(config *Config, fromSecret bool) erro
 		}
 	}
 
+	err = az.initCaches()
+	if err != nil {
+		return err
+	}
+
+	if err := initDiskControllers(az); err != nil {
+		return err
+	}
+
+	// start delayed route updater.
+	az.routeUpdater = newDelayedRouteUpdater(az, routeUpdateInterval)
+	go az.routeUpdater.run()
+
+	if syncZones {
+		// wait for the success first time of syncing zones
+		err = az.syncRegionZonesMap()
+		if err != nil {
+			klog.Errorf("InitializeCloudFromConfig: failed to sync regional zones map for the first time: %s", err.Error())
+			return err
+		}
+
+		go az.refreshZones(az.syncRegionZonesMap)
+	}
+
+	return nil
+}
+
+func (az *Cloud) initCaches() (err error) {
 	az.vmCache, err = az.newVMCache()
 	if err != nil {
 		return err
@@ -524,16 +583,6 @@ func (az *Cloud) InitializeCloudFromConfig(config *Config, fromSecret bool) erro
 		return err
 	}
 
-	if err := initDiskControllers(az); err != nil {
-		return err
-	}
-
-	// start delayed route updater.
-	az.routeUpdater = newDelayedRouteUpdater(az, routeUpdateInterval)
-	go az.routeUpdater.run()
-
-	go az.refreshZones(az.syncRegionZonesMap)
-
 	return nil
 }
 
@@ -688,7 +737,10 @@ func (az *Cloud) configAzureClients(
 	az.PublicIPAddressesClient = publicipclient.New(publicIPClientConfig)
 	az.FileClient = fileclient.New(fileClientConfig)
 	az.AvailabilitySetsClient = vmasclient.New(vmasClientConfig)
-	az.ZoneClient = zoneclient.New(zoneClientConfig)
+
+	if az.ZoneClient == nil {
+		az.ZoneClient = zoneclient.New(zoneClientConfig)
+	}
 }
 
 func (az *Cloud) getAzureClientConfig(servicePrincipalToken *adal.ServicePrincipalToken) *azclients.ClientConfig {
diff --git a/pkg/provider/azure_config.go b/pkg/provider/azure_config.go
index 1fca4bdbe..b4478e94b 100644
--- a/pkg/provider/azure_config.go
+++ b/pkg/provider/azure_config.go
@@ -51,7 +51,7 @@ func (az *Cloud) InitializeCloudFromSecret() error {
 		return nil
 	}
 
-	if err := az.InitializeCloudFromConfig(config, true); err != nil {
+	if err := az.InitializeCloudFromConfig(config, true, true); err != nil {
 		klog.Errorf("Failed to initialize Azure cloud provider: %v", err)
 		return fmt.Errorf("InitializeCloudFromSecret: failed to initialize Azure cloud provider: %w", err)
 	}
diff --git a/pkg/provider/azure_routes.go b/pkg/provider/azure_routes.go
index b73229834..c8ed1fabf 100644
--- a/pkg/provider/azure_routes.go
+++ b/pkg/provider/azure_routes.go
@@ -146,12 +146,17 @@ func (d *delayedRouteUpdater) updateRoutes() {
 	}
 
 	// reconcile routes.
-	dirty := false
+	dirty, onlyUpdateTags := false, true
 	routes := []network.Route{}
 	if routeTable.RouteTablePropertiesFormat != nil && routeTable.RouteTablePropertiesFormat.Routes != nil {
 		routes = *routeTable.Routes
 	}
-	onlyUpdateTags := true
+
+	routes, dirty = d.cleanupOutdatedRoutes(routes)
+	if dirty {
+		onlyUpdateTags = false
+	}
+
 	for _, rt := range d.routesToUpdate {
 		if rt.operation == routeTableOperationUpdateTags {
 			routeTable.Tags = rt.routeTableTags
@@ -204,6 +209,34 @@ func (d *delayedRouteUpdater) updateRoutes() {
 	}
 }
 
+// cleanupOutdatedRoutes deletes all non-dualstack routes when dualstack is enabled,
+// and deletes all dualstack routes when dualstack is not enabled.
+func (d *delayedRouteUpdater) cleanupOutdatedRoutes(existingRoutes []network.Route) (routes []network.Route, changed bool) {
+	for i := len(existingRoutes) - 1; i >= 0; i-- {
+		existingRouteName := to.String(existingRoutes[i].Name)
+		split := strings.Split(existingRouteName, consts.RouteNameSeparator)
+
+		// filter out unmanaged routes
+		deleteRoute := false
+		if d.az.nodeNames.Has(split[0]) {
+			if d.az.ipv6DualStackEnabled && len(split) == 1 {
+				klog.V(2).Infof("cleanupOutdatedRoutes: deleting outdated non-dualstack route %s", existingRouteName)
+				deleteRoute = true
+			} else if !d.az.ipv6DualStackEnabled && len(split) == 2 {
+				klog.V(2).Infof("cleanupOutdatedRoutes: deleting outdated dualstack route %s", existingRouteName)
+				deleteRoute = true
+			}
+
+			if deleteRoute {
+				existingRoutes = append(existingRoutes[:i], existingRoutes[i+1:]...)
+				changed = true
+			}
+		}
+	}
+
+	return existingRoutes, changed
+}
+
 // addRouteOperation adds the routeOperation to delayedRouteUpdater and returns a delayedRouteOperation.
 func (d *delayedRouteUpdater) addRouteOperation(operation routeOperation, route network.Route) (*delayedRouteOperation, error) {
 	d.lock.Lock()
diff --git a/pkg/provider/azure_zones.go b/pkg/provider/azure_zones.go
index 10d07e5b1..e2180df60 100644
--- a/pkg/provider/azure_zones.go
+++ b/pkg/provider/azure_zones.go
@@ -24,9 +24,8 @@ import (
 	"strings"
 	"time"
 
-	"k8s.io/apimachinery/pkg/util/wait"
-
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/wait"
 	cloudprovider "k8s.io/cloud-provider"
 	"k8s.io/klog/v2"
 
@@ -34,32 +33,40 @@ import (
 	"sigs.k8s.io/cloud-provider-azure/pkg/consts"
 )
 
-func (az *Cloud) refreshZones(refreshFunc func()) {
+func (az *Cloud) refreshZones(refreshFunc func() error) {
 	ticker := time.NewTicker(consts.ZoneFetchingInterval)
 	defer ticker.Stop()
 
-	az.regionZonesMap = make(map[string][]string)
-	refreshFunc()
-	for ; true; <-ticker.C {
-		refreshFunc()
+	for range ticker.C {
+		_ = refreshFunc()
 	}
 }
 
-func (az *Cloud) syncRegionZonesMap() {
-	klog.V(4).Infof("refreshZones: starting to fetch all available zones for the subscription %s", az.SubscriptionID)
+func (az *Cloud) syncRegionZonesMap() error {
+	klog.V(2).Infof("refreshZones: starting to fetch all available zones for the subscription %s", az.SubscriptionID)
 	zones, rerr := az.ZoneClient.GetZones(context.Background(), az.SubscriptionID)
-	if len(zones) == 0 || rerr != nil {
-		klog.Warningf("refreshZones: error when get zones, will retry after %s", consts.ZoneFetchingInterval.String())
-		return
+	if rerr != nil {
+		klog.Warningf("refreshZones: error when get zones: %s, will retry after %s", rerr.Error().Error(), consts.ZoneFetchingInterval.String())
+		return rerr.Error()
+	}
+	if len(zones) == 0 {
+		klog.Warningf("refreshZones: empty zone list, will retry after %s", consts.ZoneFetchingInterval.String())
+		return fmt.Errorf("empty zone list")
 	}
 
 	az.updateRegionZonesMap(zones)
+
+	return nil
 }
 
 func (az *Cloud) updateRegionZonesMap(zones map[string][]string) {
 	az.refreshZonesLock.Lock()
 	defer az.refreshZonesLock.Unlock()
 
+	if az.regionZonesMap == nil {
+		az.regionZonesMap = make(map[string][]string)
+	}
+
 	for region, z := range zones {
 		az.regionZonesMap[region] = z
 	}


Test patch:
diff --git a/pkg/provider/azure_config_test.go b/pkg/provider/azure_config_test.go
index 558427fd1..b9b087157 100644
--- a/pkg/provider/azure_config_test.go
+++ b/pkg/provider/azure_config_test.go
@@ -20,6 +20,9 @@ import (
 	"context"
 	"testing"
 
+	"github.com/golang/mock/gomock"
+	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/zoneclient/mockzoneclient"
+
 	"github.com/Azure/go-autorest/autorest/to"
 	"github.com/stretchr/testify/assert"
 
@@ -185,6 +188,9 @@ func TestGetConfigFromSecret(t *testing.T) {
 }
 
 func TestInitializeCloudFromSecret(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
 	emptyConfig := &Config{}
 	unknownConfigTypeConfig := getTestConfig()
 	unknownConfigTypeConfig.CloudConfigType = "UnknownConfigType"
@@ -271,6 +277,10 @@ func TestInitializeCloudFromSecret(t *testing.T) {
 				assert.NoError(t, err, test.name)
 			}
 
+			mockZoneClient := mockzoneclient.NewMockInterface(ctrl)
+			mockZoneClient.EXPECT().GetZones(gomock.Any(), gomock.Any()).Return(map[string][]string{"eastus": {"1", "2", "3"}}, nil).MaxTimes(1)
+			az.ZoneClient = mockZoneClient
+
 			err := az.InitializeCloudFromSecret()
 			assert.Equal(t, test.expectErr, err != nil)
 		})
diff --git a/pkg/provider/azure_routes_test.go b/pkg/provider/azure_routes_test.go
index 0b13c563d..b2b1132ec 100644
--- a/pkg/provider/azure_routes_test.go
+++ b/pkg/provider/azure_routes_test.go
@@ -79,7 +79,7 @@ func TestDeleteRoute(t *testing.T) {
 			Routes: &[]network.Route{},
 		},
 	}
-	routeTableClient.EXPECT().Get(gomock.Any(), cloud.RouteTableResourceGroup, cloud.RouteTableName, "").Return(routeTables, nil)
+	routeTableClient.EXPECT().Get(gomock.Any(), cloud.RouteTableResourceGroup, cloud.RouteTableName, "").Return(routeTables, nil).AnyTimes()
 	routeTableClient.EXPECT().CreateOrUpdate(gomock.Any(), cloud.RouteTableResourceGroup, cloud.RouteTableName, routeTablesAfterDeletion, "").Return(nil)
 	err := cloud.DeleteRoute(context.TODO(), "cluster", &route)
 	if err != nil {
@@ -664,3 +664,81 @@ func TestListRoutes(t *testing.T) {
 		}
 	}
 }
+
+func TestCleanupOutdatedRoutes(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	for _, testCase := range []struct {
+		description                          string
+		existingRoutes, expectedRoutes       []network.Route
+		existingNodeNames                    sets.String
+		expectedChanged, enableIPV6DualStack bool
+	}{
+		{
+			description: "cleanupOutdatedRoutes should delete outdated non-dualstack routes when dualstack is enabled",
+			existingRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			expectedRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+			},
+			existingNodeNames:   sets.NewString("aks-node1-vmss000000"),
+			enableIPV6DualStack: true,
+			expectedChanged:     true,
+		},
+		{
+			description: "cleanupOutdatedRoutes should delete outdated dualstack routes when dualstack is disabled",
+			existingRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			expectedRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			existingNodeNames: sets.NewString("aks-node1-vmss000000"),
+			expectedChanged:   true,
+		},
+		{
+			description: "cleanupOutdatedRoutes should not delete unmanaged routes when dualstack is enabled",
+			existingRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			expectedRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			existingNodeNames:   sets.NewString("aks-node1-vmss000001"),
+			enableIPV6DualStack: true,
+		},
+		{
+			description: "cleanupOutdatedRoutes should not delete unmanaged routes when dualstack is disabled",
+			existingRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			expectedRoutes: []network.Route{
+				{Name: to.StringPtr("aks-node1-vmss000000____xxx")},
+				{Name: to.StringPtr("aks-node1-vmss000000")},
+			},
+			existingNodeNames: sets.NewString("aks-node1-vmss000001"),
+		},
+	} {
+		t.Run(testCase.description, func(t *testing.T) {
+			cloud := &Cloud{
+				ipv6DualStackEnabled: testCase.enableIPV6DualStack,
+				nodeNames:            testCase.existingNodeNames,
+			}
+
+			d := &delayedRouteUpdater{
+				az: cloud,
+			}
+
+			routes, changed := d.cleanupOutdatedRoutes(testCase.existingRoutes)
+			assert.Equal(t, testCase.expectedChanged, changed)
+			assert.Equal(t, testCase.expectedRoutes, routes)
+		})
+	}
+}
diff --git a/pkg/provider/azure_test.go b/pkg/provider/azure_test.go
index 39741389a..a4997096a 100644
--- a/pkg/provider/azure_test.go
+++ b/pkg/provider/azure_test.go
@@ -46,6 +46,7 @@ import (
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/securitygroupclient/mocksecuritygroupclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/subnetclient/mocksubnetclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmclient/mockvmclient"
+	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/zoneclient/mockzoneclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/consts"
 	"sigs.k8s.io/cloud-provider-azure/pkg/retry"
 )
@@ -2123,12 +2124,28 @@ func validateConfig(t *testing.T, config string) { //nolint
 }
 
 func getCloudFromConfig(t *testing.T, config string) *Cloud {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
 	configReader := strings.NewReader(config)
-	azureCloud, err := NewCloudWithoutFeatureGates(configReader)
-	if err != nil {
-		t.Error(err)
+	c, err := parseConfig(configReader)
+	assert.NoError(t, err)
+
+	az := &Cloud{
+		nodeNames:          sets.NewString(),
+		nodeZones:          map[string]sets.String{},
+		nodeResourceGroups: map[string]string{},
+		unmanagedNodes:     sets.NewString(),
+		routeCIDRs:         map[string]string{},
+		ZoneClient:         mockzoneclient.NewMockInterface(ctrl),
 	}
-	return azureCloud
+	mockZoneClient := az.ZoneClient.(*mockzoneclient.MockInterface)
+	mockZoneClient.EXPECT().GetZones(gomock.Any(), gomock.Any()).Return(map[string][]string{"eastus": {"1", "2", "3"}}, nil)
+
+	err = az.InitializeCloudFromConfig(c, false, true)
+	assert.NoError(t, err)
+
+	return az
 }
 
 // TODO include checks for other appropriate default config parameters
@@ -3295,14 +3312,18 @@ func TestInitializeCloudFromConfig(t *testing.T) {
 	defer ctrl.Finish()
 	az := GetTestCloud(ctrl)
 
-	err := az.InitializeCloudFromConfig(nil, false)
+	mockZoneClient := mockzoneclient.NewMockInterface(ctrl)
+	mockZoneClient.EXPECT().GetZones(gomock.Any(), gomock.Any()).Return(map[string][]string{"eastus": {"1", "2", "3"}}, nil).AnyTimes()
+	az.ZoneClient = mockZoneClient
+
+	err := az.InitializeCloudFromConfig(nil, false, true)
 	assert.Equal(t, fmt.Errorf("InitializeCloudFromConfig: cannot initialize from nil config"), err)
 
 	config := Config{
 		DisableAvailabilitySetNodes: true,
 		VMType:                      consts.VMTypeStandard,
 	}
-	err = az.InitializeCloudFromConfig(&config, false)
+	err = az.InitializeCloudFromConfig(&config, false, true)
 	expectedErr := fmt.Errorf("disableAvailabilitySetNodes true is only supported when vmType is 'vmss'")
 	assert.Equal(t, expectedErr, err)
 
@@ -3312,7 +3333,7 @@ func TestInitializeCloudFromConfig(t *testing.T) {
 		},
 		CloudConfigType: cloudConfigTypeFile,
 	}
-	err = az.InitializeCloudFromConfig(&config, false)
+	err = az.InitializeCloudFromConfig(&config, false, true)
 	expectedErr = fmt.Errorf("useInstanceMetadata must be enabled without Azure credentials")
 	assert.Equal(t, expectedErr, err)
 }