address review comments

test/e2e/cluster_upgrade_runtimesdk.go

@@ -323,19 +323,33 @@ func machineSetPreflightChecksTestHandler(ctx context.Context, c client.Client,
 	// Scale up the MachineDeployment.
 	// Since the MachineDeployment is pending an upgrade at this point the topology controller will not push any changes
 	// to the MachineDeployment. Therefore, the changes made to the MachineDeployment here will not be replaced
-	// immediately.
+	// until the AfterControlPlaneUpgrade hook unblocks the upgrade.
 	*md.Spec.Replicas += 1
 	Eventually(func() error {
 		return patchHelper.Patch(ctx, md)
 	}).Should(Succeed(), "Failed to scale up the MachineDeployment %s", klog.KObj(md))
-
-	// Verify that the MachineSet is not creating the new Machine.
-	Eventually(func(g Gomega) {
-		machineSets := framework.GetMachineSetsByCluster(ctx, framework.GetMachineSetsByClusterInput{
+	// Verify the MachineDeployment updated replicas are not overridden.
+	Consistently(func(g Gomega) {
+		mds := framework.GetMachineDeploymentsByCluster(ctx, framework.GetMachineDeploymentsByClusterInput{
 			Lister:      c,
 			ClusterName: clusterRef.Name,
 			Namespace:   clusterRef.Namespace,
 		})
+		for i := range mds {
+			if mds[i].Name != md.Name {
+				continue
+			}
+			g.Expect(mds[i].Spec.Replicas).To(Equal(md.Spec.Replicas))
+		}
+	})
+
+	// Verify that the MachineSet is not creating the new Machine.
+	Eventually(func(g Gomega) {
+		machineSets := framework.GetMachineSetsByDeployment(ctx, framework.GetMachineSetsByDeploymentInput{
+			Lister:    c,
+			MDName:    md.Name,
+			Namespace: md.Namespace,
+		})
 		g.Expect(conditions.IsFalse(machineSets[0], clusterv1.MachinesCreatedCondition)).To(BeTrue())
 		machinesCreatedCondition := conditions.Get(machineSets[0], clusterv1.MachinesCreatedCondition)
 		g.Expect(machinesCreatedCondition).NotTo(BeNil())

test/framework/machinedeployment_helpers.go

@@ -90,30 +90,6 @@ func GetMachineDeploymentsByCluster(ctx context.Context, input GetMachineDeploym
 	return deployments
 }
 
-// GetMachineSetsByClusterInput is the input for GetMachineSetsByCluster.
-type GetMachineSetsByClusterInput struct {
-	Lister      Lister
-	ClusterName string
-	Namespace   string
-}
-
-// GetMachineSetsByCluster returns the MachineSets objects for a cluster.
-// Important! this method relies on labels that are created by the CAPI controllers during the first reconciliation, so
-// it is necessary to ensure this is already happened before calling it.
-func GetMachineSetsByCluster(ctx context.Context, input GetMachineSetsByClusterInput) []*clusterv1.MachineSet {
-	machineSetList := &clusterv1.MachineSetList{}
-	Eventually(func() error {
-		return input.Lister.List(ctx, machineSetList, byClusterOptions(input.ClusterName, input.Namespace)...)
-	}, retryableOperationTimeout, retryableOperationInterval).Should(Succeed(), "Failed to list MachineSets object for Cluster %s", klog.KRef(input.Namespace, input.ClusterName))
-
-	machineSets := make([]*clusterv1.MachineSet, len(machineSetList.Items))
-	for i := range machineSetList.Items {
-		Expect(machineSetList.Items[i].Spec.Replicas).ToNot(BeNil())
-		machineSets[i] = &machineSetList.Items[i]
-	}
-	return machineSets
-}
-
 // WaitForMachineDeploymentNodesToExistInput is the input for WaitForMachineDeploymentNodesToExist.
 type WaitForMachineDeploymentNodesToExistInput struct {
 	Lister            Lister

test/framework/machineset_helpers.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package framework
+
+import (
+	"context"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	. "github.com/onsi/gomega"
+	"k8s.io/klog/v2"
+
+	clusterv1 "sigs.k8s.io/cluster-api/api/v1beta1"
+)
+
+// GetMachineSetsByDeploymentInput is the input for GetMachineSetsByDeployment.
+type GetMachineSetsByDeploymentInput struct {
+	Lister    Lister
+	MDName    string
+	Namespace string
+}
+
+// GetMachineSetsByDeployment returns the MachineSets objects for a MachineDeployment.
+// Important! this method relies on labels that are created by the CAPI controllers during the first reconciliation, so
+// it is necessary to ensure this is already happened before calling it.
+func GetMachineSetsByDeployment(ctx context.Context, input GetMachineSetsByDeploymentInput) []*clusterv1.MachineSet {
+	machineSetList := &clusterv1.MachineSetList{}
+	Eventually(func() error {
+		return input.Lister.List(ctx, machineSetList, client.InNamespace(input.Namespace), client.MatchingLabels{clusterv1.MachineDeploymentNameLabel: input.MDName})
+	}, retryableOperationTimeout, retryableOperationInterval).Should(Succeed(), "Failed to list MachineSets for MachineDeployment %s", klog.KRef(input.Namespace, input.MDName))
+
+	machineSets := make([]*clusterv1.MachineSet, len(machineSetList.Items))
+	for i := range machineSetList.Items {
+		machineSets[i] = &machineSetList.Items[i]
+	}
+	return machineSets
+}