From 039256c1f480014848039b218d6b3e632f390659 Mon Sep 17 00:00:00 2001
From: Konstantinos Angelopoulos <konstantinos.angelopoulos@sap.com>
Date: Mon, 25 Nov 2024 11:37:18 +0100
Subject: [PATCH 1/3] implement migration

---
 .gitignore                                    |   3 +-
 Makefile                                      |   1 +
 hack/api-reference/api.md                     |  17 ++-
 pkg/admission/validator/shoot.go              |   4 +-
 pkg/apis/azure/helper/helper.go               |  14 +-
 pkg/apis/azure/helper/helper_test.go          |  33 ++---
 pkg/apis/azure/types_infrastructure.go        |   6 +-
 .../azure/v1alpha1/types_infrastructure.go    |   6 +-
 .../azure/v1alpha1/zz_generated.conversion.go |   2 +
 pkg/apis/azure/validation/infrastructure.go   |  43 ++----
 .../azure/validation/infrastructure_test.go   | 138 ++++++------------
 pkg/azure/client/loadbalancers.go             |   9 ++
 pkg/azure/client/types.go                     |   1 +
 pkg/azure/types.go                            |   2 +
 .../infrastructure/infraflow/ensurer.go       | 135 ++++++++++++++++-
 .../infrastructure/infraflow/flow_context.go  |  24 ++-
 .../infrastructure/infraflow/infra_adapter.go |  33 +++--
 .../infrastructure/infraflow/migrations.go    |  97 ++++++++++++
 .../infraflow/shared/whiteboard.go            |   7 +-
 .../worker/machine_dependencies_test.go       |   3 -
 pkg/internal/infrastructure/terraform_test.go |   1 +
 .../infrastructure/infrastructure_test.go     |  48 ------
 22 files changed, 399 insertions(+), 228 deletions(-)
 create mode 100644 pkg/controller/infrastructure/infraflow/migrations.go

diff --git a/.gitignore b/.gitignore
index a17b5f0d3..0d7066afe 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,6 +5,7 @@
 **/dev
 /bin
 hack/tools/bin
+/.run
 
 *.coverprofile
 *.html
@@ -21,4 +22,4 @@ TODO
 
 # GitGuardian
 .cache_ggshield
-.gitguardian.yaml
\ No newline at end of file
+.gitguardian.yaml
diff --git a/Makefile b/Makefile
index c1567ee89..4711695ce 100644
--- a/Makefile
+++ b/Makefile
@@ -79,6 +79,7 @@ start-admission:
 		./cmd/$(EXTENSION_PREFIX)-$(ADMISSION_NAME) \
 		--webhook-config-server-host=0.0.0.0 \
 		--webhook-config-server-port=$(WEBHOOK_CONFIG_PORT) \
+		--leader-election-namespace=garden \
         --webhook-config-mode=$(WEBHOOK_CONFIG_MODE) \
         $(WEBHOOK_PARAM)
 
diff --git a/hack/api-reference/api.md b/hack/api-reference/api.md
index 14e87b318..59b255ee1 100644
--- a/hack/api-reference/api.md
+++ b/hack/api-reference/api.md
@@ -1021,7 +1021,20 @@ ResourceGroup
 </em>
 </td>
 <td>
-<p>AvailabilitySets is a list of created availability sets</p>
+<p>AvailabilitySets is a list of created availability sets
+Deprecated: Will be removed in future versions.</p>
+</td>
+</tr>
+<tr>
+<td>
+<code>migratingToVMO</code></br>
+<em>
+bool
+</em>
+</td>
+<td>
+<p>MigratingToVMO indicates whether the infrastructure controller has prepared the migration from Availability set.
+Deprecated: Will be removed in future versions.</p>
 </td>
 </tr>
 <tr>
@@ -1034,7 +1047,7 @@ ResourceGroup
 </em>
 </td>
 <td>
-<p>AvailabilitySets is a list of created route tables</p>
+<p>RouteTables is a list of created route tables</p>
 </td>
 </tr>
 <tr>
diff --git a/pkg/admission/validator/shoot.go b/pkg/admission/validator/shoot.go
index eb0370a5d..b5cca03cb 100644
--- a/pkg/admission/validator/shoot.go
+++ b/pkg/admission/validator/shoot.go
@@ -21,7 +21,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 
 	api "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure"
-	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
 	azurevalidation "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/validation"
 )
 
@@ -114,7 +113,7 @@ func (s *shoot) validateShoot(shoot *core.Shoot, oldInfraConfig, infraConfig *ap
 		// Cloudprofile validation
 		allErrs = append(allErrs, azurevalidation.ValidateInfrastructureConfigAgainstCloudProfile(oldInfraConfig, infraConfig, shoot.Spec.Region, cloudProfileSpec, infraConfigPath)...)
 		// Provider validation
-		allErrs = append(allErrs, azurevalidation.ValidateInfrastructureConfig(infraConfig, shoot.Spec.Networking, helper.HasShootVmoAlphaAnnotation(shoot.Annotations), infraConfigPath)...)
+		allErrs = append(allErrs, azurevalidation.ValidateInfrastructureConfig(infraConfig, shoot, infraConfigPath)...)
 	}
 	if cpConfig != nil {
 		allErrs = append(allErrs, azurevalidation.ValidateControlPlaneConfig(cpConfig, shoot.Spec.Kubernetes.Version, cpConfigPath)...)
@@ -169,7 +168,6 @@ func (s *shoot) validateUpdate(oldShoot, shoot *core.Shoot, cloudProfileSpec *ga
 		allErrs = append(allErrs, azurevalidation.ValidateInfrastructureConfigUpdate(oldInfraConfig, infraConfig, metaDataPath)...)
 	}
 
-	allErrs = append(allErrs, azurevalidation.ValidateVmoConfigUpdate(helper.HasShootVmoAlphaAnnotation(oldShoot.Annotations), helper.HasShootVmoAlphaAnnotation(shoot.Annotations), metaDataPath)...)
 	allErrs = append(allErrs, azurevalidation.ValidateWorkersUpdate(oldShoot.Spec.Provider.Workers, shoot.Spec.Provider.Workers, workersPath)...)
 
 	allErrs = append(allErrs, s.validateShoot(shoot, oldInfraConfig, infraConfig, cloudProfileSpec, cpConfig)...)
diff --git a/pkg/apis/azure/helper/helper.go b/pkg/apis/azure/helper/helper.go
index 706646ae4..e444ef626 100644
--- a/pkg/apis/azure/helper/helper.go
+++ b/pkg/apis/azure/helper/helper.go
@@ -125,9 +125,19 @@ func FindImageFromCloudProfile(cloudProfileConfig *api.CloudProfileConfig, image
 	return nil, fmt.Errorf("no machine image found with name %q, architecture %q and version %q", imageName, *architecture, imageVersion)
 }
 
-// IsVmoRequired determines if VMO is required.
+// IsVmoRequired determines if VMO is required. It is different from the condition in the infrastructure as this one depends on whether the infra controller
+// has finished migrating the Availability sets.
 func IsVmoRequired(infrastructureStatus *api.InfrastructureStatus) bool {
-	return !infrastructureStatus.Zoned && len(infrastructureStatus.AvailabilitySets) == 0
+	return !infrastructureStatus.Zoned && (len(infrastructureStatus.AvailabilitySets) == 0 || infrastructureStatus.MigratingToVMO)
+}
+
+// HasShootVmoMigrationAnnotation determines if the passed Shoot annotations contain instruction to use VMO.
+func HasShootVmoMigrationAnnotation(shootAnnotations map[string]string) bool {
+	value, exists := shootAnnotations[azure.ShootVmoMigrationAnnotation]
+	if exists && value == "true" {
+		return true
+	}
+	return false
 }
 
 // HasShootVmoAlphaAnnotation determines if the passed Shoot annotations contain instruction to use VMO.
diff --git a/pkg/apis/azure/helper/helper_test.go b/pkg/apis/azure/helper/helper_test.go
index 84c44a322..fad78292d 100644
--- a/pkg/apis/azure/helper/helper_test.go
+++ b/pkg/apis/azure/helper/helper_test.go
@@ -11,7 +11,6 @@ import (
 
 	api "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure"
 	. "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
-	"github.com/gardener/gardener-extension-provider-azure/pkg/azure"
 )
 
 var (
@@ -151,40 +150,30 @@ var _ = Describe("Helper", func() {
 			"ubuntu", "1", ptr.To("foo"), &api.MachineImage{Name: "ubuntu", Version: "1", Image: api.Image{SharedGalleryImageID: &profileSharedImageId}, Architecture: ptr.To("foo")}),
 	)
 
-	DescribeTable("#IsVmoRequired",
-		func(zoned bool, availabilitySet *api.AvailabilitySet, expectedVmoRequired bool) {
+	DescribeTable("#IsVmoRequiredForInfrastructure",
+		func(zoned bool, availabilitySet *api.AvailabilitySet, migrateToVMO bool, expectedVmoRequired bool) {
 			var infrastructureStatus = &api.InfrastructureStatus{
 				Zoned: zoned,
 			}
 			if availabilitySet != nil {
 				infrastructureStatus.AvailabilitySets = append(infrastructureStatus.AvailabilitySets, *availabilitySet)
 			}
+			infrastructureStatus.MigratingToVMO = migrateToVMO
 
 			Expect(IsVmoRequired(infrastructureStatus)).To(Equal(expectedVmoRequired))
 		},
-		Entry("should require a VMO", false, nil, true),
-		Entry("should not require VMO for zoned cluster", true, nil, false),
+		Entry("should require a VMO", false, nil, false, true),
+		Entry("should not require VMO for zoned cluster", true, nil, false, false),
 		Entry("should not require VMO for a cluster with primary availabilityset (non zoned)", false, &api.AvailabilitySet{
 			ID:      "/my/azure/availabilityset/id",
 			Name:    "my-availabilityset",
 			Purpose: api.PurposeNodes,
-		}, false),
-	)
-
-	DescribeTable("#HasShootVmoAlphaAnnotation",
-		func(hasVmoAnnotaion, hasCorrectVmoAnnotationValue, expectedResult bool) {
-			var annotations = map[string]string{}
-			if hasVmoAnnotaion {
-				annotations[azure.ShootVmoUsageAnnotation] = "some-arbitrary-value"
-			}
-			if hasCorrectVmoAnnotationValue {
-				annotations[azure.ShootVmoUsageAnnotation] = "true"
-			}
-			Expect(HasShootVmoAlphaAnnotation(annotations)).To(Equal(expectedResult))
-		},
-		Entry("should return true as shoot annotations contain vmo alpha annotation with value true", true, true, true),
-		Entry("should return false as shoot annotations contain vmo alpha annotation with wrong value", true, false, false),
-		Entry("should return false as shoot annotations do not contain vmo alpha annotation", false, false, false),
+		}, false, false),
+		Entry("should require VMO for a cluster with primary availabilityset with migration to VMO", false, &api.AvailabilitySet{
+			ID:      "/my/azure/availabilityset/id",
+			Name:    "my-availabilityset",
+			Purpose: api.PurposeNodes,
+		}, true, true),
 	)
 })
 
diff --git a/pkg/apis/azure/types_infrastructure.go b/pkg/apis/azure/types_infrastructure.go
index f5dcf2d0f..92cf2561a 100644
--- a/pkg/apis/azure/types_infrastructure.go
+++ b/pkg/apis/azure/types_infrastructure.go
@@ -105,8 +105,12 @@ type InfrastructureStatus struct {
 	// ResourceGroup is azure resource group
 	ResourceGroup ResourceGroup
 	// AvailabilitySets is a list of created availability sets
+	// Deprecated: Will be removed in future versions.
 	AvailabilitySets []AvailabilitySet
-	// AvailabilitySets is a list of created route tables
+	// MigratingToVMO indicates whether the infrastructure controller has prepared the migration from Availability set.
+	// Deprecated: Will be removed in future versions.
+	MigratingToVMO bool
+	// RouteTables is a list of created route tables
 	RouteTables []RouteTable
 	// SecurityGroups is a list of created security groups
 	SecurityGroups []SecurityGroup
diff --git a/pkg/apis/azure/v1alpha1/types_infrastructure.go b/pkg/apis/azure/v1alpha1/types_infrastructure.go
index 137abdc5d..9742a97e4 100644
--- a/pkg/apis/azure/v1alpha1/types_infrastructure.go
+++ b/pkg/apis/azure/v1alpha1/types_infrastructure.go
@@ -118,8 +118,12 @@ type InfrastructureStatus struct {
 	// ResourceGroup is azure resource group
 	ResourceGroup ResourceGroup `json:"resourceGroup"`
 	// AvailabilitySets is a list of created availability sets
+	// Deprecated: Will be removed in future versions.
 	AvailabilitySets []AvailabilitySet `json:"availabilitySets"`
-	// AvailabilitySets is a list of created route tables
+	// MigratingToVMO indicates whether the infrastructure controller has prepared the migration from Availability set.
+	// Deprecated: Will be removed in future versions.
+	MigratingToVMO bool `json:"migratingToVMO,omitempty"`
+	// RouteTables is a list of created route tables
 	RouteTables []RouteTable `json:"routeTables"`
 	// SecurityGroups is a list of created security groups
 	SecurityGroups []SecurityGroup `json:"securityGroups"`
diff --git a/pkg/apis/azure/v1alpha1/zz_generated.conversion.go b/pkg/apis/azure/v1alpha1/zz_generated.conversion.go
index 204552b79..50eb768d7 100644
--- a/pkg/apis/azure/v1alpha1/zz_generated.conversion.go
+++ b/pkg/apis/azure/v1alpha1/zz_generated.conversion.go
@@ -758,6 +758,7 @@ func autoConvert_v1alpha1_InfrastructureStatus_To_azure_InfrastructureStatus(in
 		return err
 	}
 	out.AvailabilitySets = *(*[]azure.AvailabilitySet)(unsafe.Pointer(&in.AvailabilitySets))
+	out.MigratingToVMO = in.MigratingToVMO
 	out.RouteTables = *(*[]azure.RouteTable)(unsafe.Pointer(&in.RouteTables))
 	out.SecurityGroups = *(*[]azure.SecurityGroup)(unsafe.Pointer(&in.SecurityGroups))
 	out.Identity = (*azure.IdentityStatus)(unsafe.Pointer(in.Identity))
@@ -778,6 +779,7 @@ func autoConvert_azure_InfrastructureStatus_To_v1alpha1_InfrastructureStatus(in
 		return err
 	}
 	out.AvailabilitySets = *(*[]AvailabilitySet)(unsafe.Pointer(&in.AvailabilitySets))
+	out.MigratingToVMO = in.MigratingToVMO
 	out.RouteTables = *(*[]RouteTable)(unsafe.Pointer(&in.RouteTables))
 	out.SecurityGroups = *(*[]SecurityGroup)(unsafe.Pointer(&in.SecurityGroups))
 	out.Identity = (*IdentityStatus)(unsafe.Pointer(in.Identity))
diff --git a/pkg/apis/azure/validation/infrastructure.go b/pkg/apis/azure/validation/infrastructure.go
index 29ccc4a61..0f08d3018 100644
--- a/pkg/apis/azure/validation/infrastructure.go
+++ b/pkg/apis/azure/validation/infrastructure.go
@@ -16,7 +16,7 @@ import (
 
 	apisazure "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
-	"github.com/gardener/gardener-extension-provider-azure/pkg/azure"
+	azuretypes "github.com/gardener/gardener-extension-provider-azure/pkg/azure"
 )
 
 const (
@@ -75,8 +75,9 @@ func validateInfrastructureConfigZones(oldInfra, infra *apisazure.Infrastructure
 }
 
 // ValidateInfrastructureConfig validates a InfrastructureConfig object.
-func ValidateInfrastructureConfig(infra *apisazure.InfrastructureConfig, networking *core.Networking, hasVmoAlphaAnnotation bool, fldPath *field.Path) field.ErrorList {
+func ValidateInfrastructureConfig(infra *apisazure.InfrastructureConfig, shoot *core.Shoot, fldPath *field.Path) field.ErrorList {
 	allErrs := field.ErrorList{}
+	networking := shoot.Spec.Networking
 
 	var (
 		nodes, pods, services             cidrvalidation.CIDR
@@ -105,11 +106,7 @@ func ValidateInfrastructureConfig(infra *apisazure.InfrastructureConfig, network
 		allErrs = append(allErrs, field.Invalid(fldPath.Child("resourceGroup"), infra.ResourceGroup, "specifying an existing resource group is not supported yet"))
 	}
 
-	if infra.Zoned && hasVmoAlphaAnnotation {
-		allErrs = append(allErrs, field.Invalid(fldPath.Child("zoned"), infra.Zoned, fmt.Sprintf("specifying a zoned cluster and having the %q annotation is not allowed", azure.ShootVmoUsageAnnotation)))
-	}
-
-	allErrs = append(allErrs, validateNetworkConfig(infra, nodes, pods, services, hasVmoAlphaAnnotation, fldPath)...)
+	allErrs = append(allErrs, validateNetworkConfig(shoot, infra, nodes, pods, services, fldPath)...)
 
 	if infra.Identity != nil && (infra.Identity.Name == "" || infra.Identity.ResourceGroup == "") {
 		allErrs = append(allErrs, field.Invalid(fldPath.Child("identity"), infra.Identity, "specifying an identity requires the name of the identity and the resource group which hosts the identity"))
@@ -119,11 +116,11 @@ func ValidateInfrastructureConfig(infra *apisazure.InfrastructureConfig, network
 }
 
 func validateNetworkConfig(
+	shoot *core.Shoot,
 	infra *apisazure.InfrastructureConfig,
 	nodes cidrvalidation.CIDR,
 	pods cidrvalidation.CIDR,
 	services cidrvalidation.CIDR,
-	hasVmoAlphaAnnotation bool,
 	fldPath *field.Path,
 ) field.ErrorList {
 
@@ -164,7 +161,7 @@ func validateNetworkConfig(
 			allErrs = append(allErrs, nodes.ValidateSubset(workerCIDR)...)
 		}
 
-		allErrs = append(allErrs, validateNatGatewayConfig(config.NatGateway, infra.Zoned, hasVmoAlphaAnnotation, networksPath.Child("natGateway"))...)
+		allErrs = append(allErrs, validateNatGatewayConfig(config.NatGateway, helper.HasShootVmoMigrationAnnotation(shoot.GetAnnotations()), networksPath.Child("natGateway"))...)
 		return allErrs
 	}
 
@@ -172,10 +169,10 @@ func validateNetworkConfig(
 	if !infra.Zoned {
 		allErrs = append(allErrs, field.Forbidden(zonesPath, "cannot specify zones in an non-zonal cluster"))
 	}
-
 	if config.NatGateway != nil {
 		allErrs = append(allErrs, field.Forbidden(workersPath, "natGateway cannot be specified when workers field is missing"))
 	}
+
 	if len(config.ServiceEndpoints) > 0 {
 		allErrs = append(allErrs, field.Forbidden(workersPath, "serviceEndpoints cannot be specified when workers field is missing"))
 	}
@@ -282,7 +279,7 @@ func validateZones(zones []apisazure.Zone, nodes, pods, services cidrvalidation.
 	return allErrs
 }
 
-func validateNatGatewayConfig(natGatewayConfig *apisazure.NatGatewayConfig, zoned bool, hasVmoAlphaAnnotation bool, natGatewayPath *field.Path) field.ErrorList {
+func validateNatGatewayConfig(natGatewayConfig *apisazure.NatGatewayConfig, hasShootVmoMigrationAnnotation bool, natGatewayPath *field.Path) field.ErrorList {
 	allErrs := field.ErrorList{}
 
 	if natGatewayConfig == nil {
@@ -296,11 +293,8 @@ func validateNatGatewayConfig(natGatewayConfig *apisazure.NatGatewayConfig, zone
 		return nil
 	}
 
-	// NatGateway cannot be offered for Shoot clusters with a primary AvailabilitySet.
-	// The NatGateway is not compatible with the Basic SKU Loadbalancers which are
-	// required to use for Shoot clusters with AvailabilitySet.
-	if !zoned && !hasVmoAlphaAnnotation {
-		return append(allErrs, field.Forbidden(natGatewayPath, "NatGateway is currently only supported for zonal and VMO clusters"))
+	if hasShootVmoMigrationAnnotation {
+		allErrs = append(allErrs, field.Forbidden(natGatewayPath.Child("enabled"), fmt.Sprintf("natGateway cannot be enabled with the annotation %s", azuretypes.ShootVmoMigrationAnnotation)))
 	}
 
 	if natGatewayConfig.IdleConnectionTimeoutMinutes != nil && (*natGatewayConfig.IdleConnectionTimeoutMinutes < natGatewayMinTimeoutInMinutes || *natGatewayConfig.IdleConnectionTimeoutMinutes > natGatewayMaxTimeoutInMinutes) {
@@ -395,23 +389,6 @@ func ValidateInfrastructureConfigUpdate(oldConfig, newConfig *apisazure.Infrastr
 	return allErrs
 }
 
-// ValidateVmoConfigUpdate validates the VMO configuration on update.
-func ValidateVmoConfigUpdate(oldShootHasAlphaVmoAnnotation, newShootHasAlphaVmoAnnotation bool, metaDataPath *field.Path) field.ErrorList {
-	allErrs := field.ErrorList{}
-
-	// Check if old shoot has not the vmo alpha annotation and forbid to add it.
-	if !oldShootHasAlphaVmoAnnotation && newShootHasAlphaVmoAnnotation {
-		allErrs = append(allErrs, field.Forbidden(metaDataPath.Child("annotations"), fmt.Sprintf("not allowed to add annotation %q to an already existing shoot cluster", azure.ShootVmoUsageAnnotation)))
-	}
-
-	// Check if old shoot has the vmo alpha annotaion and forbid to remove it.
-	if oldShootHasAlphaVmoAnnotation && !newShootHasAlphaVmoAnnotation {
-		allErrs = append(allErrs, field.Forbidden(metaDataPath.Child("annotations"), fmt.Sprintf("not allowed to remove annotation %q to an already existing shoot cluster", azure.ShootVmoUsageAnnotation)))
-	}
-
-	return allErrs
-}
-
 func validateVnetConfigUpdate(oldNeworkConfig, newNetworkConfig *apisazure.NetworkConfig, networkConfigPath *field.Path) field.ErrorList {
 	allErrs := field.ErrorList{}
 	vnetPath := networkConfigPath.Child("vnet")
diff --git a/pkg/apis/azure/validation/infrastructure_test.go b/pkg/apis/azure/validation/infrastructure_test.go
index 66f7663c8..514564b53 100644
--- a/pkg/apis/azure/validation/infrastructure_test.go
+++ b/pkg/apis/azure/validation/infrastructure_test.go
@@ -5,8 +5,6 @@
 package validation_test
 
 import (
-	"fmt"
-
 	"github.com/gardener/gardener/pkg/apis/core"
 	"github.com/gardener/gardener/pkg/apis/core/v1beta1"
 	. "github.com/gardener/gardener/pkg/utils/test/matchers"
@@ -18,20 +16,19 @@ import (
 
 	apisazure "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure"
 	. "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/validation"
-	"github.com/gardener/gardener-extension-provider-azure/pkg/azure"
 )
 
 var _ = Describe("InfrastructureConfig validation", func() {
 	var (
-		infrastructureConfig  *apisazure.InfrastructureConfig
-		nodes                 string
-		resourceGroup         = "shoot--test--foo"
-		hasVmoAlphaAnnotation bool
+		infrastructureConfig *apisazure.InfrastructureConfig
+		nodes                string
+		resourceGroup        = "shoot--test--foo"
 
 		pods        = "100.96.0.0/11"
 		services    = "100.64.0.0/13"
 		vnetCIDR    = "10.0.0.0/8"
 		invalidCIDR = "invalid-cidr"
+		shoot       = core.Shoot{}
 		networking  = core.Networking{}
 
 		workers      = "10.250.3.0/24"
@@ -47,6 +44,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			Services: &services,
 			Nodes:    &nodes,
 		}
+		shoot.Spec.Networking = &networking
 		infrastructureConfig = &apisazure.InfrastructureConfig{
 			Networks: apisazure.NetworkConfig{
 				Workers: &workers,
@@ -55,7 +53,6 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				},
 			},
 		}
-		hasVmoAlphaAnnotation = false
 		ddosProtectionPlanID = "/subscriptions/test/resourceGroups/test/providers/Microsoft.Network/ddosProtectionPlans/test-ddos-protection-plan"
 	})
 
@@ -63,7 +60,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 		It("should forbid specifying a resource group configuration", func() {
 			infrastructureConfig.ResourceGroup = &apisazure.ResourceGroup{}
 
-			errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+			errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 			Expect(errorList).To(ConsistOfFields(Fields{
 				"Type":  Equal(field.ErrorTypeInvalid),
@@ -77,7 +74,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				infrastructureConfig.Networks.VNet = apisazure.VNet{
 					Name: &vnetName,
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -92,7 +89,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				infrastructureConfig.Networks.VNet = apisazure.VNet{
 					ResourceGroup: &vnetGroup,
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -107,7 +104,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					Name:          ptr.To(""),
 					ResourceGroup: ptr.To(""),
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -130,7 +127,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					ResourceGroup: &vnetGroup,
 					CIDR:          &vnetCIDR,
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -149,7 +146,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				infrastructureConfig.ResourceGroup = &apisazure.ResourceGroup{
 					Name: resourceGroup,
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -169,7 +166,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				infrastructureConfig.Networks = apisazure.NetworkConfig{
 					Workers: &workers,
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).To(HaveLen(0))
 			})
 
@@ -177,7 +174,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				networking.Nodes = ptr.To("10.250.3.0/24")
 				infrastructureConfig.ResourceGroup = nil
 				infrastructureConfig.Networks.VNet.DDosProtectionPlanID = &ddosProtectionPlanID
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).To(HaveLen(0))
 			})
 
@@ -188,7 +185,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					ResourceGroup:        &resourceGroup,
 					DDosProtectionPlanID: &ddosProtectionPlanID,
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -198,25 +195,11 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			})
 		})
 
-		Context("Zonal", func() {
-			It(fmt.Sprintf("should forbid specifying the %q annotation for a zonal cluster", azure.ShootVmoUsageAnnotation), func() {
-				infrastructureConfig.Zoned = true
-				hasVmoAlphaAnnotation = true
-
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
-
-				Expect(errorList).To(ConsistOfFields(Fields{
-					"Type":  Equal(field.ErrorTypeInvalid),
-					"Field": Equal("zoned"),
-				}))
-			})
-		})
-
 		Context("CIDR", func() {
 			It("should forbid invalid VNet CIDRs", func() {
 				infrastructureConfig.Networks.VNet.CIDR = &invalidCIDR
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":   Equal(field.ErrorTypeInvalid),
@@ -228,7 +211,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			It("should forbid invalid workers CIDR", func() {
 				infrastructureConfig.Networks.Workers = &invalidCIDR
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":   Equal(field.ErrorTypeInvalid),
@@ -241,7 +224,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				emptyStr := ""
 				infrastructureConfig.Networks.Workers = &emptyStr
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -254,7 +237,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			It("should forbid nil workers CIDR", func() {
 				infrastructureConfig.Networks.Workers = nil
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(
 					Fields{
@@ -268,7 +251,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				notOverlappingCIDR := "1.1.1.1/32"
 				infrastructureConfig.Networks.Workers = &notOverlappingCIDR
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":   Equal(field.ErrorTypeInvalid),
@@ -284,7 +267,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			It("should forbid Pod CIDR to overlap with VNet CIDR", func() {
 				networking.Pods = ptr.To("10.0.0.1/32")
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":   Equal(field.ErrorTypeInvalid),
@@ -295,7 +278,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			It("should forbid Services CIDR to overlap with VNet CIDR", func() {
 				networking.Services = ptr.To("10.0.0.1/32")
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":   Equal(field.ErrorTypeInvalid),
@@ -313,7 +296,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				infrastructureConfig.Networks.Workers = &workers
 				infrastructureConfig.Networks.VNet = apisazure.VNet{CIDR: &vpcCIDR}
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).To(HaveLen(2))
 				Expect(errorList).To(ConsistOfFields(Fields{
@@ -334,14 +317,14 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					Name:          "test-identiy",
 					ResourceGroup: "identity-resource-group",
 				}
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should return errors because no name or resource group is given", func() {
 				infrastructureConfig.Identity = &apisazure.IdentityConfig{
 					Name: "test-identiy",
 				}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":  Equal(field.ErrorTypeInvalid),
 					"Field": Equal("identity"),
@@ -357,19 +340,19 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 			It("should pass as there is no NatGateway config is provided", func() {
 				infrastructureConfig.Networks.NatGateway = nil
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should pass as the NatGateway is disabled", func() {
 				infrastructureConfig.Networks.NatGateway.Enabled = false
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should fail as NatGatway is disabled but additional config for the NatGateway is supplied", func() {
 				infrastructureConfig.Networks.NatGateway.Enabled = false
 				infrastructureConfig.Networks.NatGateway.Zone = ptr.To[int32](2)
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).To(ConsistOfFields(Fields{
 					"Type":   Equal(field.ErrorTypeInvalid),
 					"Field":  Equal("networks.natGateway"),
@@ -377,21 +360,9 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				}))
 			})
 
-			It("should fail as NatGatway is enabled but the cluster is not zonal", func() {
-				infrastructureConfig.Zoned = false
-				infrastructureConfig.Networks.NatGateway.Enabled = true
-
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
-				Expect(errorList).To(ConsistOfFields(Fields{
-					"Type":   Equal(field.ErrorTypeForbidden),
-					"Field":  Equal("networks.natGateway"),
-					"Detail": Equal("NatGateway is currently only supported for zonal and VMO clusters"),
-				}))
-			})
-
 			It("should pass as the NatGateway has a zone", func() {
 				infrastructureConfig.Networks.NatGateway.Zone = ptr.To[int32](2)
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			Context("User provided public IP", func() {
@@ -406,7 +377,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 				It("should fail as NatGateway has no zone but an external public ip", func() {
 					infrastructureConfig.Networks.NatGateway.Zone = nil
-					errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+					errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 					Expect(errorList).To(ConsistOfFields(Fields{
 						"Type":   Equal(field.ErrorTypeInvalid),
 						"Field":  Equal("networks.natGateway.zone"),
@@ -416,7 +387,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 				It("should fail as resource is in a different zone as the NatGateway", func() {
 					infrastructureConfig.Networks.NatGateway.IPAddresses[0].Zone = 2
-					errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+					errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 					Expect(errorList).To(ConsistOfFields(Fields{
 						"Type":   Equal(field.ErrorTypeInvalid),
 						"Field":  Equal("networks.natGateway.ipAddresses[0].zone"),
@@ -426,7 +397,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 				It("should fail as name is empty", func() {
 					infrastructureConfig.Networks.NatGateway.IPAddresses[0].Name = ""
-					errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+					errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 					Expect(errorList).To(ConsistOfFields(Fields{
 						"Type":   Equal(field.ErrorTypeRequired),
 						"Field":  Equal("networks.natGateway.ipAddresses[0].name"),
@@ -436,7 +407,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 				It("should fail as resource group is empty", func() {
 					infrastructureConfig.Networks.NatGateway.IPAddresses[0].ResourceGroup = ""
-					errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+					errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 					Expect(errorList).To(ConsistOfFields(Fields{
 						"Type":   Equal(field.ErrorTypeRequired),
 						"Field":  Equal("networks.natGateway.ipAddresses[0].resourceGroup"),
@@ -450,7 +421,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					var timeoutValue int32 = 0
 					infrastructureConfig.Zoned = true
 					infrastructureConfig.Networks.NatGateway = &apisazure.NatGatewayConfig{Enabled: true, IdleConnectionTimeoutMinutes: &timeoutValue}
-					errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+					errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 					Expect(errorList).To(HaveLen(1))
 					Expect(errorList).To(ConsistOfFields(Fields{
 						"Type":   Equal(field.ErrorTypeInvalid),
@@ -463,7 +434,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					var timeoutValue int32 = 121
 					infrastructureConfig.Zoned = true
 					infrastructureConfig.Networks.NatGateway = &apisazure.NatGatewayConfig{Enabled: true, IdleConnectionTimeoutMinutes: &timeoutValue}
-					errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+					errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 					Expect(errorList).To(HaveLen(1))
 					Expect(errorList).To(ConsistOfFields(Fields{
 						"Type":   Equal(field.ErrorTypeInvalid),
@@ -476,7 +447,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 					var timeoutValue int32 = 120
 					infrastructureConfig.Zoned = true
 					infrastructureConfig.Networks.NatGateway = &apisazure.NatGatewayConfig{Enabled: true, IdleConnectionTimeoutMinutes: &timeoutValue}
-					Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+					Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 				})
 			})
 		})
@@ -511,13 +482,13 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			})
 
 			It("should succeed", func() {
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should succeed for nil service and pod CIDR", func() {
 				networking.Pods = nil
 				networking.Services = nil
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should succeed with NAT Gateway", func() {
@@ -527,7 +498,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				infrastructureConfig.Networks.Zones[1].NatGateway = &apisazure.ZonedNatGatewayConfig{
 					Enabled: true,
 				}
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should succeed with NAT Gateway and  public IPs", func() {
@@ -540,12 +511,12 @@ var _ = Describe("InfrastructureConfig validation", func() {
 						},
 					},
 				}
-				Expect(ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)).To(BeEmpty())
+				Expect(ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)).To(BeEmpty())
 			})
 
 			It("should forbid non canonical CIDRs", func() {
 				infrastructureConfig.Networks.Zones[0].CIDR = "10.250.0.1/24"
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).NotTo(BeEmpty())
 				Expect(errorList).To(HaveLen(1))
 				Expect(errorList).To(ConsistOfFields(Fields{
@@ -556,7 +527,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 			It("should forbid overlapping zone CIDRs", func() {
 				infrastructureConfig.Networks.Zones[0].CIDR = zoneCIDR1
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).NotTo(BeEmpty())
 				Expect(errorList).To(HaveLen(1))
 				Expect(errorList).To(ConsistOfFields(Fields{
@@ -568,7 +539,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 
 			It("should forbid not specifying VNet when using Zones", func() {
 				infrastructureConfig.Networks.VNet = apisazure.VNet{}
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).NotTo(BeEmpty())
 				Expect(errorList).To(HaveLen(1))
 				Expect(errorList).To(ConsistOfFields(Fields{
@@ -583,7 +554,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 				networking.Nodes = ptr.To("10.150.0.0/16")
 
 				infrastructureConfig.Networks.VNet.CIDR = &vpcCIDR
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 
 				Expect(errorList).NotTo(BeEmpty())
 				Expect(errorList).To(HaveLen(2))
@@ -601,7 +572,7 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			It("should forbid specifying zone multiple times", func() {
 				infrastructureConfig.Networks.Zones[0].Name = zoneName1
 
-				errorList := ValidateInfrastructureConfig(infrastructureConfig, &networking, hasVmoAlphaAnnotation, providerPath)
+				errorList := ValidateInfrastructureConfig(infrastructureConfig, &shoot, providerPath)
 				Expect(errorList).NotTo(BeEmpty())
 				Expect(errorList).To(HaveLen(1))
 				Expect(errorList).To(ConsistOfFields(Fields{
@@ -908,25 +879,4 @@ var _ = Describe("InfrastructureConfig validation", func() {
 			})
 		})
 	})
-
-	DescribeTable("#ValidateVmoConfigUpdate",
-		func(newHasVmoAlphaAnnotation, oldHasVmoAlphaAnnotation, expectErrors bool) {
-			var (
-				path      *field.Path
-				errorList = ValidateVmoConfigUpdate(newHasVmoAlphaAnnotation, oldHasVmoAlphaAnnotation, path)
-			)
-			if !expectErrors {
-				Expect(errorList).To(HaveLen(0))
-				return
-			}
-			Expect(errorList).To(ConsistOf(PointTo(MatchFields(IgnoreExtras, Fields{
-				"Type":  Equal(field.ErrorTypeForbidden),
-				"Field": Equal("annotations"),
-			}))))
-		},
-		Entry("should pass as old and new cluster have vmo alpha annotation", true, true, false),
-		Entry("should pass as old and new cluster don't have vmo alpha annotation", false, false, false),
-		Entry("should forbid removing the vmo alpha annotation for an already existing cluster", true, false, true),
-		Entry("should forbid adding the vmo alpha annotation to an already existing cluster", false, true, true),
-	)
 })
diff --git a/pkg/azure/client/loadbalancers.go b/pkg/azure/client/loadbalancers.go
index 39760ded4..03e1bea03 100644
--- a/pkg/azure/client/loadbalancers.go
+++ b/pkg/azure/client/loadbalancers.go
@@ -25,6 +25,15 @@ func NewLoadBalancersClient(auth internal.ClientAuth, tc azcore.TokenCredential,
 	return &LoadBalancersClient{client}, err
 }
 
+// Get gets a given virtual load balancer by name
+func (c *LoadBalancersClient) Get(ctx context.Context, resourceGroupName, name string) (*armnetwork.LoadBalancer, error) {
+	res, err := c.client.Get(ctx, resourceGroupName, name, nil)
+	if err != nil {
+		return nil, FilterNotFoundError(err)
+	}
+	return &res.LoadBalancer, err
+}
+
 // List lists all subnets of a given virtual network.
 func (c *LoadBalancersClient) List(ctx context.Context, resourceGroupName string) ([]*armnetwork.LoadBalancer, error) {
 	pager := c.client.NewListPager(resourceGroupName, nil)
diff --git a/pkg/azure/client/types.go b/pkg/azure/client/types.go
index d9e8c176e..7786b8709 100644
--- a/pkg/azure/client/types.go
+++ b/pkg/azure/client/types.go
@@ -125,6 +125,7 @@ type Subnet interface {
 
 // LoadBalancer represents an Azure LoadBalancer k8sClient.
 type LoadBalancer interface {
+	GetFunc[armnetwork.LoadBalancer]
 	ListFunc[armnetwork.LoadBalancer]
 	DeleteFunc[armnetwork.LoadBalancer]
 }
diff --git a/pkg/azure/types.go b/pkg/azure/types.go
index 075d22beb..83cfbfec7 100644
--- a/pkg/azure/types.go
+++ b/pkg/azure/types.go
@@ -14,6 +14,8 @@ const (
 
 	// ShootVmoUsageAnnotation is an annotation assigned to the Shoot resource which indicates if VMO should be used.
 	ShootVmoUsageAnnotation = "alpha.azure.provider.extensions.gardener.cloud/vmo"
+	// ShootVmoMigrationAnnotation is an annotation assigned to the Shoot resource which indicates if the availability set shoot, should be migrated to a VMO shoot.
+	ShootVmoMigrationAnnotation = "migration.azure.provider.extensions.gardener.cloud/vmo"
 
 	// NetworkLayoutZoneMigrationAnnotation is used when migrating from a single subnet network layout to a multiple subnet network layout to indicate the zone that the existing subnet should be assigned to.
 	NetworkLayoutZoneMigrationAnnotation = "migration.azure.provider.extensions.gardener.cloud/zone"
diff --git a/pkg/controller/infrastructure/infraflow/ensurer.go b/pkg/controller/infrastructure/infraflow/ensurer.go
index 0a100bd47..b9d8fc536 100644
--- a/pkg/controller/infrastructure/infraflow/ensurer.go
+++ b/pkg/controller/infrastructure/infraflow/ensurer.go
@@ -8,18 +8,23 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v5"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v4"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
+	"github.com/gardener/gardener/pkg/utils/flow"
 	"github.com/go-logr/logr"
+	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/utils/ptr"
+	k8sclient "sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/v1alpha1"
+	"github.com/gardener/gardener-extension-provider-azure/pkg/azure"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/azure/client"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/controller/infrastructure/infraflow/shared"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/internal/infrastructure"
@@ -171,7 +176,16 @@ func (fctx *FlowContext) EnsureAvailabilitySet(ctx context.Context) error {
 		return nil
 	}
 
-	avset, err := fctx.ensureAvailabilitySet(ctx, log, *avsetCfg)
+	// complete AS migration.
+	if v := fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Get("complete"); v != nil && *v == "true" {
+		err := fctx.deleteAvailabilitySet(ctx, log)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	avset, err := fctx.ensureAvailabilitySet(ctx, log)
 	if err != nil {
 		return err
 	}
@@ -183,13 +197,41 @@ func (fctx *FlowContext) EnsureAvailabilitySet(ctx context.Context) error {
 	fctx.whiteboard.GetChild(ChildKeyIDs).Set(KindAvailabilitySet.String(), *avset.ID)
 	return nil
 }
+func (fctx *FlowContext) deleteAvailabilitySet(ctx context.Context, log logr.Logger) error {
+	// try to delete the availability set. It can only  work if it does not contain any VMs.
+	asClient, err := fctx.factory.AvailabilitySet()
+	if err != nil {
+		return err
+	}
+
+	av, err := asClient.Get(ctx, fctx.adapter.AvailabilitySetConfig().ResourceGroup, fctx.adapter.AvailabilitySetConfig().Name)
+	if err != nil {
+		return err
+	}
+	if av == nil {
+		return nil
+	}
+	// if the AS contains no VMs then we attempt to delete it and complete the migration
+	if len(av.Properties.VirtualMachines) == 0 {
+		log.Info("Deleting Availability Set", "Name", *av.Name)
+		if err := asClient.Delete(ctx, fctx.adapter.ResourceGroupName(), *av.Name); err != nil {
+			return err
+		}
+		fctx.whiteboard.GetChild(ChildKeyIDs).Delete(KindAvailabilitySet.String())
+		fctx.inventory.Delete(*av.ID)
+		return nil
+	}
+	log.Info("Skipping deleting Availability Set because it still contains VMs", "Name", *av.Name)
+	return nil
+}
 
-func (fctx *FlowContext) ensureAvailabilitySet(ctx context.Context, log logr.Logger, avsetCfg AvailabilitySetConfig) (*armcompute.AvailabilitySet, error) {
+func (fctx *FlowContext) ensureAvailabilitySet(ctx context.Context, log logr.Logger) (*armcompute.AvailabilitySet, error) {
 	asClient, err := fctx.factory.AvailabilitySet()
 	if err != nil {
 		return nil, err
 	}
 
+	avsetCfg := fctx.adapter.AvailabilitySetConfig()
 	avset, err := asClient.Get(ctx, avsetCfg.ResourceGroup, avsetCfg.Name)
 	if err != nil {
 		return nil, err
@@ -699,6 +741,88 @@ func (fctx *FlowContext) EnsureManagedIdentity(ctx context.Context) (err error)
 	return err
 }
 
+// MigrateAvailabilitySet prepares an AS-based shoot to be migrated to VMSS-Flex.
+func (fctx *FlowContext) MigrateAvailabilitySet(ctx context.Context) error {
+	var (
+		log = shared.LogFromContext(ctx)
+		c   = fctx.client
+	)
+
+	// return early if the migration has already been complete
+	if v := fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Get("complete"); v != nil && *v == "true" {
+		return nil
+	}
+	// return early if the cluster does not have AS.
+	if fctx.whiteboard.GetChild(ChildKeyIDs).Get(KindAvailabilitySet.String()) == nil {
+		return nil
+	}
+
+	// IF VMOs are not needed, or the migration is already done, return early.
+	if !helper.HasShootVmoMigrationAnnotation(fctx.cluster.Shoot.GetAnnotations()) {
+		return nil
+	}
+
+	log.Info("Preparing for the migration to VMOs")
+	scaleDownDeployment := func(ctx context.Context, key k8sclient.ObjectKey) error {
+		log.Info("Scaling deployment to 0 replicas", "Name", key.String())
+		deployment := &appsv1.Deployment{}
+		if err := fctx.client.Get(ctx, k8sclient.ObjectKey{
+			Namespace: fctx.infra.Namespace,
+			Name:      azure.CloudControllerManagerName,
+		}, deployment); k8sclient.IgnoreNotFound(err) != nil {
+			return err
+		} else if k8sclient.IgnoreNotFound(err) != nil {
+			return nil
+		}
+
+		if ptr.Deref(deployment.Spec.Replicas, 1) == 0 {
+			return nil
+		}
+		patch := k8sclient.MergeFrom(deployment.DeepCopy())
+		deployment.Spec.Replicas = ptr.To(int32(0))
+		// Apply the patch on the "scale" subresource
+		if err := c.SubResource("scale").Patch(ctx, deployment, patch); err != nil {
+			return fmt.Errorf("failed to patch deployment %s with replicas: %w", key.String(), err)
+		}
+		return nil
+	}
+
+	if err := flow.Parallel(func(ctx context.Context) error {
+		return scaleDownDeployment(ctx, k8sclient.ObjectKey{Namespace: fctx.infra.Namespace, Name: azure.CloudControllerManagerName})
+	}, func(ctx context.Context) error {
+		// we want to scale CA down to avoid VMs getting created as they may claim internal subnet IPs in the case of an existing internal loadbalancer.
+		return scaleDownDeployment(ctx, k8sclient.ObjectKey{Namespace: fctx.infra.Namespace, Name: "cluster-autoscaler"})
+	}).RetryUntilTimeout(5*time.Second, defaultTimeout)(ctx); err != nil {
+		return err
+	}
+
+	log.Info("Backing-up Public IPs to be migrated")
+	// we will first back up the PIPs that we want to migrate. In the simplest case, we would only migrate PIPs in the shoot's RG. But the loadbalancer may reference PIPs from other RGs.
+	// If we delete the LB we would have no way to recover the PIPs in other RGs, hence we will back it up.
+	if err := fctx.BackupPIPsForBasicLBMigration(ctx); err != nil {
+		return err
+	}
+
+	loadbalancerClient, err := fctx.factory.LoadBalancer()
+	if err != nil {
+		return err
+	}
+	log.Info("Deleting load balancer", "Name", fctx.infra.Namespace)
+	if err := loadbalancerClient.Delete(ctx, fctx.adapter.ResourceGroupName(), fctx.adapter.TechnicalName()); err != nil {
+		return err
+	}
+	log.Info("Deleting internal load balancer", "Name", fctx.infra.Namespace)
+	if err := loadbalancerClient.Delete(ctx, fctx.adapter.ResourceGroupName(), fmt.Sprintf("%s-internal", fctx.adapter.TechnicalName())); err != nil {
+		return err
+	}
+
+	if err := fctx.UpdatePublicIPs(ctx); err != nil {
+		return err
+	}
+	fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Set("complete", "true")
+	return fctx.PersistState(ctx)
+}
+
 // GetInfrastructureStatus returns the infrastructure status.
 func (fctx *FlowContext) GetInfrastructureStatus(_ context.Context) (*v1alpha1.InfrastructureStatus, error) {
 	status := &v1alpha1.InfrastructureStatus{
@@ -754,7 +878,8 @@ func (fctx *FlowContext) GetInfrastructureStatus(_ context.Context) (*v1alpha1.I
 	}
 	status.Networks.OutboundAccessType = outboundAccessType
 
-	if cfg := fctx.adapter.AvailabilitySetConfig(); cfg != nil {
+	if fctx.whiteboard.GetChild(ChildKeyIDs).Get(KindAvailabilitySet.String()) != nil {
+		cfg := fctx.adapter.AvailabilitySetConfig()
 		status.AvailabilitySets = []v1alpha1.AvailabilitySet{
 			{
 				Purpose:            v1alpha1.PurposeNodes,
@@ -764,6 +889,9 @@ func (fctx *FlowContext) GetInfrastructureStatus(_ context.Context) (*v1alpha1.I
 				CountUpdateDomains: cfg.CountUpdateDomains,
 			},
 		}
+		if v := fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Get("complete"); v != nil && *v == "true" {
+			status.MigratingToVMO = true
+		}
 	}
 
 	if identity := fctx.cfg.Identity; identity != nil {
@@ -782,6 +910,7 @@ func (fctx *FlowContext) GetInfrastructureState() *runtime.RawExtension {
 	state := &v1alpha1.InfrastructureState{
 		TypeMeta:     helper.InfrastructureStateTypeMeta,
 		ManagedItems: fctx.inventory.ToList(),
+		Data:         fctx.whiteboard.ExportAsFlatMap(),
 	}
 
 	return &runtime.RawExtension{
diff --git a/pkg/controller/infrastructure/infraflow/flow_context.go b/pkg/controller/infrastructure/infraflow/flow_context.go
index e3667dc0f..fbc288927 100644
--- a/pkg/controller/infrastructure/infraflow/flow_context.go
+++ b/pkg/controller/infrastructure/infraflow/flow_context.go
@@ -7,6 +7,7 @@ package infraflow
 import (
 	"context"
 	"errors"
+	"sync"
 	"time"
 
 	"github.com/gardener/gardener/extensions/pkg/controller"
@@ -28,6 +29,14 @@ const (
 	defaultLongTimeout = 4 * time.Minute
 )
 
+var setSeparator = sync.OnceFunc(func() {
+	shared.Separator = "|"
+})
+
+func init() {
+	setSeparator()
+}
+
 // FlowContext is the reconciler for all managed resources
 type FlowContext struct {
 	log            logr.Logger
@@ -37,6 +46,7 @@ type FlowContext struct {
 	auth           *internal.ClientAuth
 	infra          *extensionsv1alpha1.Infrastructure
 	state          *azure.InfrastructureState
+	status         *azure.InfrastructureStatus
 	cluster        *controller.Cluster
 	whiteboard     shared.Whiteboard
 	adapter        *InfrastructureAdapter
@@ -104,6 +114,7 @@ func NewFlowContext(opts Opts) (*FlowContext, error) {
 		state:      opts.State,
 		cluster:    opts.Cluster,
 		cfg:        cfg,
+		status:     status,
 		whiteboard: wb,
 		providerAccess: &access{
 			opts.Factory,
@@ -148,8 +159,8 @@ func (fctx *FlowContext) buildReconcileGraph() *flow.Graph {
 		fctx.EnsureVirtualNetwork, shared.Timeout(defaultTimeout), shared.Dependencies(resourceGroup))
 
 	_ = fctx.AddTask(g, "ensure availability set",
-		fctx.EnsureAvailabilitySet, shared.DoIf(fctx.adapter.AvailabilitySetConfig() != nil),
-		shared.Timeout(defaultTimeout), shared.Dependencies(resourceGroup))
+		fctx.EnsureAvailabilitySet,
+		shared.Timeout(defaultTimeout), shared.Dependencies(resourceGroup), shared.DoIf(fctx.adapter.IsAvailabilitySetReconciliationRequired()))
 
 	_ = fctx.AddTask(g, "ensure managed identity",
 		fctx.EnsureManagedIdentity, shared.DoIf(fctx.cfg.Identity != nil))
@@ -165,8 +176,15 @@ func (fctx *FlowContext) buildReconcileGraph() *flow.Graph {
 	nat := fctx.AddTask(g, "ensure nats",
 		fctx.EnsureNatGateways, shared.Timeout(defaultLongTimeout), shared.Dependencies(resourceGroup, ip))
 
-	_ = fctx.AddTask(g, "ensure subnets", fctx.EnsureSubnets,
+	subnet := fctx.AddTask(g, "ensure subnets", fctx.EnsureSubnets,
 		shared.Timeout(defaultLongTimeout), shared.Dependencies(vnet, routeTable, securityGroup, nat))
+
+	// a sync point for when the "normal" reconciliation is finished. Currently, it ends with the subnet reconciliation.
+	reconciliationFinishedPoint := flow.NewTaskIDs(subnet)
+
+	_ = fctx.AddTask(g, "availability set migration", fctx.MigrateAvailabilitySet,
+		shared.Timeout(defaultLongTimeout), shared.Dependencies(reconciliationFinishedPoint),
+		shared.DoIf(!fctx.cfg.Zoned))
 	return g
 }
 
diff --git a/pkg/controller/infrastructure/infraflow/infra_adapter.go b/pkg/controller/infrastructure/infraflow/infra_adapter.go
index 886795124..98cf5158f 100644
--- a/pkg/controller/infrastructure/infraflow/infra_adapter.go
+++ b/pkg/controller/infrastructure/infraflow/infra_adapter.go
@@ -157,9 +157,23 @@ type AvailabilitySetConfig struct {
 	Location           string
 }
 
-// AvailabilitySetRequired returns true if gardener should create an availability set for the shoot.
-func (ia *InfrastructureAdapter) availabilitySetRequired() (bool, error) {
-	return infrastructure.IsPrimaryAvailabilitySetRequired(ia.infra, ia.config, ia.cluster)
+// IsAvailabilitySetReconciliationRequired returns true if gardener should create an availability set for the shoot.
+func (ia *InfrastructureAdapter) IsAvailabilitySetReconciliationRequired() bool {
+	if ia.config.Zoned {
+		return false
+	}
+	// If the infrastructureStatus already exists that mean the Infrastucture is already created.
+	if len(ia.status.AvailabilitySets) > 0 {
+		if _, err := helper.FindAvailabilitySetByPurpose(ia.status.AvailabilitySets, azure.PurposeNodes); err == nil {
+			return true
+		}
+	}
+	return false
+}
+
+// IsVmoRequiredForInfrastructure determines if VMO is required.
+func (ia *InfrastructureAdapter) IsVmoRequiredForInfrastructure() bool {
+	return !ia.config.Zoned && (len(ia.status.AvailabilitySets) == 0 || helper.HasShootVmoMigrationAnnotation(ia.cluster.Shoot.GetAnnotations()))
 }
 
 // AvailabilitySetConfig returns the configuration for the shoot's availability set.
@@ -167,18 +181,17 @@ func (ia *InfrastructureAdapter) AvailabilitySetConfig() *AvailabilitySetConfig
 	return ia.avSetConfig
 }
 
+// AvailabilitySetName is the name of the availability set of the shoot.
+func (ia *InfrastructureAdapter) AvailabilitySetName() string {
+	return fmt.Sprintf("%s-avset-workers", ia.TechnicalName())
+}
+
 // AvailabilitySetConfig returns the availability set's configuration.
 func (ia *InfrastructureAdapter) availabilitySetConfig() (*AvailabilitySetConfig, error) {
-	if ok, err := ia.availabilitySetRequired(); err != nil {
-		return nil, err
-	} else if !ok {
-		return nil, nil
-	}
-
 	asc := &AvailabilitySetConfig{
 		AzureResourceMetadata: AzureResourceMetadata{
 			ResourceGroup: ia.ResourceGroupName(),
-			Name:          fmt.Sprintf("%s-avset-workers", ia.TechnicalName()),
+			Name:          ia.AvailabilitySetName(),
 			Kind:          KindAvailabilitySet,
 		},
 	}
diff --git a/pkg/controller/infrastructure/infraflow/migrations.go b/pkg/controller/infrastructure/infraflow/migrations.go
new file mode 100644
index 000000000..d454e7eec
--- /dev/null
+++ b/pkg/controller/infrastructure/infraflow/migrations.go
@@ -0,0 +1,97 @@
+// SPDX-FileCopyrightText: 2024 SAP SE or an SAP affiliate company and Gardener contributors
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package infraflow
+
+import (
+	"context"
+	"errors"
+	"sync"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v4"
+
+	"github.com/gardener/gardener-extension-provider-azure/pkg/controller/infrastructure/infraflow/shared"
+)
+
+// BackupPIPsForBasicLBMigration saves the Public IPs attached to the basic load balancer of the shoot to the state.
+func (fctx *FlowContext) BackupPIPsForBasicLBMigration(ctx context.Context) error {
+	loadbalancerClient, err := fctx.factory.LoadBalancer()
+	if err != nil {
+		return err
+	}
+	lb, err := loadbalancerClient.Get(ctx, fctx.adapter.ResourceGroupName(), fctx.infra.Namespace)
+	if err != nil {
+		return err
+	}
+	// if we don't find the loadbalancer or some properties are missing, exit early
+	if lb == nil ||
+		lb.SKU == nil ||
+		lb.SKU.Name == nil ||
+		*lb.SKU.Name != armnetwork.LoadBalancerSKUNameBasic ||
+		lb.Properties == nil ||
+		len(lb.Properties.FrontendIPConfigurations) == 0 {
+		return nil
+	}
+	if lb.Properties == nil || len(lb.Properties.FrontendIPConfigurations) == 0 {
+		return nil
+	}
+	for _, fipc := range lb.Properties.FrontendIPConfigurations {
+		if fipc.Properties == nil {
+			continue
+		}
+		if fipc.Properties.PublicIPAddress != nil {
+			resourceID, err := arm.ParseResourceID(*fipc.Properties.PublicIPAddress.ID)
+			if err != nil {
+				return err
+			}
+			fctx.whiteboard.GetChild("migration").GetChild("basic-lb").GetChild(resourceID.ResourceGroupName).Set(resourceID.Name, "true")
+		}
+	}
+	return fctx.PersistState(ctx)
+}
+
+// UpdatePublicIPs updates the public IPs saved in the state from basic to standard SKU.
+func (fctx *FlowContext) UpdatePublicIPs(ctx context.Context) error {
+	log := shared.LogFromContext(ctx)
+	ipc, err := fctx.factory.PublicIP()
+	if err != nil {
+		return err
+	}
+	var (
+		wg   = sync.WaitGroup{}
+		errs error
+	)
+	for _, resourceGroup := range fctx.whiteboard.GetChild("migration").GetChild("basic-lb").GetChildrenKeys() {
+		for _, pipName := range fctx.whiteboard.GetChild("migration").GetChild("basic-lb").GetChild(resourceGroup).Keys() {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				pip, err := ipc.Get(ctx, resourceGroup, pipName, nil)
+				if err != nil {
+					errs = errors.Join(err)
+					return
+				}
+				if pip == nil || pip.SKU == nil || pip.SKU.Name == nil || *pip.SKU.Name != armnetwork.PublicIPAddressSKUNameBasic {
+					return
+				}
+				pip.SKU = &armnetwork.PublicIPAddressSKU{
+					Name: to.Ptr(armnetwork.PublicIPAddressSKUNameStandard),
+					Tier: to.Ptr(armnetwork.PublicIPAddressSKUTierRegional),
+				}
+				log.Info("upgrading basic PIP", "ResourceGroup", resourceGroup, "Name", pipName)
+				_, err = ipc.CreateOrUpdate(ctx, resourceGroup, pipName, *pip)
+				if err != nil {
+					errs = errors.Join(err)
+					return
+				}
+				// removing upgraded PIP from state.
+				fctx.whiteboard.GetChild("migration").GetChild("basic-lb").GetChild(resourceGroup).Delete(pipName)
+			}()
+		}
+	}
+	wg.Wait()
+	return errs
+}
diff --git a/pkg/controller/infrastructure/infraflow/shared/whiteboard.go b/pkg/controller/infrastructure/infraflow/shared/whiteboard.go
index bd75781f8..8a9e7001c 100644
--- a/pkg/controller/infrastructure/infraflow/shared/whiteboard.go
+++ b/pkg/controller/infrastructure/infraflow/shared/whiteboard.go
@@ -14,12 +14,15 @@ import (
 )
 
 const (
-	// Separator is used to on translating keys to/from flat maps
-	Separator = "/"
 	// deleted is a special value to mark a resource id as deleted
 	deleted = "<deleted>"
 )
 
+var (
+	// Separator is used to on translating keys to/from flat maps
+	Separator = "/"
+)
+
 // FlatMap is a semantic name for string map used for importing and exporting
 type FlatMap map[string]string
 
diff --git a/pkg/controller/worker/machine_dependencies_test.go b/pkg/controller/worker/machine_dependencies_test.go
index 1862a9222..8b37f1f4f 100644
--- a/pkg/controller/worker/machine_dependencies_test.go
+++ b/pkg/controller/worker/machine_dependencies_test.go
@@ -83,9 +83,6 @@ var _ = Describe("MachinesDependencies", func() {
 
 			faultDomainCount = 3
 			cluster = makeCluster("", "westeurope", nil, nil, faultDomainCount)
-			cluster.Shoot.Annotations = map[string]string{
-				azure.ShootVmoUsageAnnotation: "true",
-			}
 			infrastructureStatus = makeInfrastructureStatus(resourceGroupName, "vnet-name", "subnet-name", false, nil, nil, nil)
 			pool = extensionsv1alpha1.WorkerPool{
 				Name: "my-pool",
diff --git a/pkg/internal/infrastructure/terraform_test.go b/pkg/internal/infrastructure/terraform_test.go
index a73333e7e..f1931aab9 100644
--- a/pkg/internal/infrastructure/terraform_test.go
+++ b/pkg/internal/infrastructure/terraform_test.go
@@ -76,6 +76,7 @@ var _ = Describe("Terraform", func() {
 					},
 				},
 			},
+			Status: extensionsv1alpha1.InfrastructureStatus{},
 		}
 
 		cluster = MakeCluster("11.0.0.0/16", "12.0.0.0/16", infra.Spec.Region, countFaultDomain, countUpdateDomain)
diff --git a/test/integration/infrastructure/infrastructure_test.go b/test/integration/infrastructure/infrastructure_test.go
index 0df98ecbe..b5bdf0536 100644
--- a/test/integration/infrastructure/infrastructure_test.go
+++ b/test/integration/infrastructure/infrastructure_test.go
@@ -299,54 +299,6 @@ var _ = BeforeSuite(func() {
 })
 
 var _ = Describe("Infrastructure tests", func() {
-	Context("AvailabilitySet cluster", func() {
-		AfterEach(func() {
-			framework.RunCleanupActions()
-		})
-
-		It("should successfully create and delete AvailabilitySet cluster creating new vNet", func() {
-			providerConfig := newInfrastructureConfig(nil, nil, nil, false)
-
-			namespace, err := generateName()
-			Expect(err).ToNot(HaveOccurred())
-
-			err = runTest(ctx, log, c, clientSet, namespace, providerConfig, false, decoder)
-			Expect(err).ToNot(HaveOccurred())
-		})
-
-		It("should successfully create and delete AvailabilitySet cluster using existing vNet and existing identity", func() {
-			foreignName, err := generateName()
-			Expect(err).ToNot(HaveOccurred())
-			foreignNameVnet := foreignName + "-vnet"
-			foreignNameId := foreignName + "-id"
-
-			var cleanupHandle framework.CleanupActionHandle
-			cleanupHandle = framework.AddCleanupAction(func() {
-				Expect(ignoreAzureNotFoundError(teardownResourceGroup(ctx, clientSet, foreignName))).To(Succeed())
-				framework.RemoveCleanupAction(cleanupHandle)
-			})
-
-			Expect(prepareNewResourceGroup(ctx, log, clientSet, foreignName, *region)).To(Succeed())
-			Expect(prepareNewVNet(ctx, log, clientSet, foreignName, foreignNameVnet, *region, VNetCIDR)).To(Succeed())
-			Expect(prepareNewIdentity(ctx, log, clientSet, foreignName, foreignNameId, *region)).To(Succeed())
-
-			vnetConfig := &azurev1alpha1.VNet{
-				Name:          ptr.To(foreignNameVnet),
-				ResourceGroup: ptr.To(foreignName),
-			}
-			identityConfig := &azurev1alpha1.IdentityConfig{
-				Name:          foreignNameId,
-				ResourceGroup: foreignName,
-			}
-			providerConfig := newInfrastructureConfig(vnetConfig, nil, identityConfig, false)
-
-			namespace, err := generateName()
-			Expect(err).ToNot(HaveOccurred())
-			err = runTest(ctx, log, c, clientSet, namespace, providerConfig, false, decoder)
-			Expect(err).ToNot(HaveOccurred())
-		})
-	})
-
 	Context("Zonal cluster", func() {
 		AfterEach(func() {
 			framework.RunCleanupActions()

From ba7a8a1980245481a46f92d326d0bc2107a78df8 Mon Sep 17 00:00:00 2001
From: Konstantinos Angelopoulos <konstantinos.angelopoulos@sap.com>
Date: Fri, 29 Nov 2024 10:46:02 +0100
Subject: [PATCH 2/3] fix ccm not found error handling

---
 pkg/controller/infrastructure/infraflow/ensurer.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pkg/controller/infrastructure/infraflow/ensurer.go b/pkg/controller/infrastructure/infraflow/ensurer.go
index b9d8fc536..990bca08d 100644
--- a/pkg/controller/infrastructure/infraflow/ensurer.go
+++ b/pkg/controller/infrastructure/infraflow/ensurer.go
@@ -18,6 +18,7 @@ import (
 	"github.com/gardener/gardener/pkg/utils/flow"
 	"github.com/go-logr/logr"
 	appsv1 "k8s.io/api/apps/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/utils/ptr"
 	k8sclient "sigs.k8s.io/controller-runtime/pkg/client"
@@ -771,7 +772,7 @@ func (fctx *FlowContext) MigrateAvailabilitySet(ctx context.Context) error {
 			Name:      azure.CloudControllerManagerName,
 		}, deployment); k8sclient.IgnoreNotFound(err) != nil {
 			return err
-		} else if k8sclient.IgnoreNotFound(err) != nil {
+		} else if err != nil && apierrors.IsNotFound(err) {
 			return nil
 		}
 

From 4b3f0217a5324d6d99c3725276cf55076bd513b0 Mon Sep 17 00:00:00 2001
From: Konstantinos Angelopoulos <konstantinos.angelopoulos@sap.com>
Date: Fri, 29 Nov 2024 15:37:47 +0100
Subject: [PATCH 3/3] redundant checks and consts for strings

---
 pkg/controller/infrastructure/infraflow/const.go         | 4 ++++
 pkg/controller/infrastructure/infraflow/ensurer.go       | 8 ++++----
 pkg/controller/infrastructure/infraflow/infra_adapter.go | 5 -----
 pkg/controller/infrastructure/infraflow/migrations.go    | 3 ---
 4 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/pkg/controller/infrastructure/infraflow/const.go b/pkg/controller/infrastructure/infraflow/const.go
index ce4e02c0d..8e0cc12b0 100644
--- a/pkg/controller/infrastructure/infraflow/const.go
+++ b/pkg/controller/infrastructure/infraflow/const.go
@@ -17,4 +17,8 @@ const (
 	KeyManagedIdentityClientId = "managed_identity_client_id"
 	// KeyManagedIdentityId is a key for the MI's identity ID.
 	KeyManagedIdentityId = "managed_identity_id"
+	// ChildKeyMigration is the prefix key for data stored during migrations.
+	ChildKeyMigration = "migration"
+	// ChildKeyComplete is a key to indicate whether a task is complete.
+	ChildKeyComplete = "complete"
 )
diff --git a/pkg/controller/infrastructure/infraflow/ensurer.go b/pkg/controller/infrastructure/infraflow/ensurer.go
index 990bca08d..fc6306842 100644
--- a/pkg/controller/infrastructure/infraflow/ensurer.go
+++ b/pkg/controller/infrastructure/infraflow/ensurer.go
@@ -178,7 +178,7 @@ func (fctx *FlowContext) EnsureAvailabilitySet(ctx context.Context) error {
 	}
 
 	// complete AS migration.
-	if v := fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Get("complete"); v != nil && *v == "true" {
+	if v := fctx.whiteboard.GetChild(ChildKeyMigration).GetChild(KindAvailabilitySet.String()).Get(ChildKeyComplete); v != nil && *v == "true" {
 		err := fctx.deleteAvailabilitySet(ctx, log)
 		if err != nil {
 			return err
@@ -750,7 +750,7 @@ func (fctx *FlowContext) MigrateAvailabilitySet(ctx context.Context) error {
 	)
 
 	// return early if the migration has already been complete
-	if v := fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Get("complete"); v != nil && *v == "true" {
+	if v := fctx.whiteboard.GetChild(ChildKeyMigration).GetChild(KindAvailabilitySet.String()).Get(ChildKeyComplete); v != nil && *v == "true" {
 		return nil
 	}
 	// return early if the cluster does not have AS.
@@ -820,7 +820,7 @@ func (fctx *FlowContext) MigrateAvailabilitySet(ctx context.Context) error {
 	if err := fctx.UpdatePublicIPs(ctx); err != nil {
 		return err
 	}
-	fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Set("complete", "true")
+	fctx.whiteboard.GetChild(ChildKeyMigration).GetChild(KindAvailabilitySet.String()).Set(ChildKeyComplete, "true")
 	return fctx.PersistState(ctx)
 }
 
@@ -890,7 +890,7 @@ func (fctx *FlowContext) GetInfrastructureStatus(_ context.Context) (*v1alpha1.I
 				CountUpdateDomains: cfg.CountUpdateDomains,
 			},
 		}
-		if v := fctx.whiteboard.GetChild("migration").GetChild(KindAvailabilitySet.String()).Get("complete"); v != nil && *v == "true" {
+		if v := fctx.whiteboard.GetChild(ChildKeyMigration).GetChild(KindAvailabilitySet.String()).Get(ChildKeyComplete); v != nil && *v == "true" {
 			status.MigratingToVMO = true
 		}
 	}
diff --git a/pkg/controller/infrastructure/infraflow/infra_adapter.go b/pkg/controller/infrastructure/infraflow/infra_adapter.go
index 98cf5158f..500a7aaa9 100644
--- a/pkg/controller/infrastructure/infraflow/infra_adapter.go
+++ b/pkg/controller/infrastructure/infraflow/infra_adapter.go
@@ -171,11 +171,6 @@ func (ia *InfrastructureAdapter) IsAvailabilitySetReconciliationRequired() bool
 	return false
 }
 
-// IsVmoRequiredForInfrastructure determines if VMO is required.
-func (ia *InfrastructureAdapter) IsVmoRequiredForInfrastructure() bool {
-	return !ia.config.Zoned && (len(ia.status.AvailabilitySets) == 0 || helper.HasShootVmoMigrationAnnotation(ia.cluster.Shoot.GetAnnotations()))
-}
-
 // AvailabilitySetConfig returns the configuration for the shoot's availability set.
 func (ia *InfrastructureAdapter) AvailabilitySetConfig() *AvailabilitySetConfig {
 	return ia.avSetConfig
diff --git a/pkg/controller/infrastructure/infraflow/migrations.go b/pkg/controller/infrastructure/infraflow/migrations.go
index d454e7eec..e7cce4ad1 100644
--- a/pkg/controller/infrastructure/infraflow/migrations.go
+++ b/pkg/controller/infrastructure/infraflow/migrations.go
@@ -35,9 +35,6 @@ func (fctx *FlowContext) BackupPIPsForBasicLBMigration(ctx context.Context) erro
 		len(lb.Properties.FrontendIPConfigurations) == 0 {
 		return nil
 	}
-	if lb.Properties == nil || len(lb.Properties.FrontendIPConfigurations) == 0 {
-		return nil
-	}
 	for _, fipc := range lb.Properties.FrontendIPConfigurations {
 		if fipc.Properties == nil {
 			continue