Update test taint checks

test/e2e/aks_marketplace.go

@@ -21,11 +21,14 @@ package e2e
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v4"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/kubernetesconfiguration/armkubernetesconfiguration"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/utils/ptr"
 	infrav1 "sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
@@ -60,30 +63,78 @@ func AKSMarketplaceExtensionSpec(ctx context.Context, inputGetter func() AKSMark
 	}, amcp)
 	Expect(err).NotTo(HaveOccurred())
 
+	agentpoolsClient, err := armcontainerservice.NewAgentPoolsClient(amcp.Spec.SubscriptionID, cred, nil)
+	Expect(err).NotTo(HaveOccurred())
+
 	extensionClient, err := armkubernetesconfiguration.NewExtensionsClient(amcp.Spec.SubscriptionID, cred, nil)
 	Expect(err).NotTo(HaveOccurred())
 
-	By("Adding a taint to the Windows node pool")
-	Eventually(func(g Gomega) {
-		ammp := &infrav1.AzureManagedMachinePool{}
-		err = mgmtClient.Get(ctx, types.NamespacedName{
-			Namespace: input.Cluster.Namespace,
-			Name:      input.Cluster.Name + "-pool2",
-		}, ammp)
+	By("Deleting all node taints for windows machine pool")
+	var ammp = &infrav1.AzureManagedMachinePool{}
+	Expect(mgmtClient.Get(ctx, types.NamespacedName{
+		Namespace: input.Cluster.Namespace,
+		Name:      input.Cluster.Name + "-pool2",
+	}, ammp)).To(Succeed())
+	initialTaints := ammp.Spec.Taints
+	var expectedTaints []infrav1.Taint
+	expectedTaints = nil
+	checkTaints := func(g Gomega) {
+		var expectedTaintStrs []*string
+		if expectedTaints != nil {
+			expectedTaintStrs = make([]*string, 0, len(expectedTaints))
+			for _, taint := range expectedTaints {
+				expectedTaintStrs = append(expectedTaintStrs, ptr.To(fmt.Sprintf("%s=%s:%s", taint.Key, taint.Value, taint.Effect)))
+			}
+		}
+
+		resp, err := agentpoolsClient.Get(ctx, amcp.Spec.ResourceGroupName, amcp.Name, *ammp.Spec.Name, nil)
 		g.Expect(err).NotTo(HaveOccurred())
-		ammp.Spec.Taints = []infrav1.Taint{
-			{
-				Effect: "NoSchedule",
-				Key:    "test",
-				Value:  "test",
-			},
+		g.Expect(resp.Properties.ProvisioningState).To(Equal(ptr.To("Succeeded")))
+		actualTaintStrs := resp.AgentPool.Properties.NodeTaints
+		if expectedTaintStrs == nil {
+			g.Expect(actualTaintStrs).To(BeNil())
+		} else {
+			g.Expect(actualTaintStrs).To(Equal(expectedTaintStrs))
 		}
+	}
+	Eventually(func(g Gomega) {
+		g.Expect(mgmtClient.Get(ctx, client.ObjectKeyFromObject(ammp), ammp)).To(Succeed())
+		ammp.Spec.Taints = expectedTaints
+		g.Expect(mgmtClient.Update(ctx, ammp)).To(Succeed())
+	}, inputGetter().WaitIntervals...).Should(Succeed())
+	Eventually(checkTaints, input.WaitIntervals...).Should(Succeed())
+
+	By("Adding a taint to the Windows node pool")
+	expectedTaints = []infrav1.Taint{
+		{
+			Effect: infrav1.TaintEffect(corev1.TaintEffectNoSchedule),
+			Key:    "capz-e2e-1",
+			Value:  "test1",
+		},
+	}
+	Eventually(func(g Gomega) {
+		g.Expect(mgmtClient.Get(ctx, client.ObjectKeyFromObject(ammp), ammp)).To(Succeed())
+		ammp.Spec.Taints = expectedTaints
+		g.Expect(mgmtClient.Update(ctx, ammp)).To(Succeed())
+		Eventually(checkTaints, input.WaitIntervals...).Should(Succeed())
+	}, input.WaitIntervals...).Should(Succeed())
+
+	By("Updating taints for Windows machine pool")
+	expectedTaints = expectedTaints[:1]
+	Eventually(func(g Gomega) {
+		g.Expect(mgmtClient.Get(ctx, client.ObjectKeyFromObject(ammp), ammp)).To(Succeed())
+		ammp.Spec.Taints = expectedTaints
+		g.Expect(mgmtClient.Update(ctx, ammp)).To(Succeed())
 	}, input.WaitIntervals...).Should(Succeed())
+	Eventually(checkTaints, input.WaitIntervals...).Should(Succeed())
 
 	By("Adding an AKS Marketplace Extension to the AzureManagedControlPlane")
 	var infraControlPlane = &infrav1.AzureManagedControlPlane{}
 	Eventually(func(g Gomega) {
-		err = mgmtClient.Get(ctx, client.ObjectKey{Namespace: input.Cluster.Spec.ControlPlaneRef.Namespace, Name: input.Cluster.Spec.ControlPlaneRef.Name}, infraControlPlane)
+		err = mgmtClient.Get(ctx, client.ObjectKey{
+			Namespace: input.Cluster.Spec.ControlPlaneRef.Namespace,
+			Name:      input.Cluster.Spec.ControlPlaneRef.Name,
+		}, infraControlPlane)
 		g.Expect(err).NotTo(HaveOccurred())
 		infraControlPlane.Spec.MarketplaceExtensions = []infrav1.MarketplaceExtension{
 			{
@@ -120,14 +171,12 @@ func AKSMarketplaceExtensionSpec(ctx context.Context, inputGetter func() AKSMark
 		g.Expect(extension.Properties.ExtensionType).To(Equal(ptr.To("testtestindustryexperiencestest.azurecomps")))
 	}, input.WaitIntervals...).Should(Succeed())
 
-	By("Removing the taint from the Windows node pool")
+	By("Restoring initial taints for Windows machine pool")
+	expectedTaints = initialTaints
 	Eventually(func(g Gomega) {
-		ammp := &infrav1.AzureManagedMachinePool{}
-		err = mgmtClient.Get(ctx, types.NamespacedName{
-			Namespace: input.Cluster.Namespace,
-			Name:      input.Cluster.Name + "-pool2",
-		}, ammp)
-		g.Expect(err).NotTo(HaveOccurred())
-		ammp.Spec.Taints = []infrav1.Taint{}
+		g.Expect(mgmtClient.Get(ctx, client.ObjectKeyFromObject(ammp), ammp)).To(Succeed())
+		ammp.Spec.Taints = expectedTaints
+		g.Expect(mgmtClient.Update(ctx, ammp)).To(Succeed())
 	}, input.WaitIntervals...).Should(Succeed())
+	Eventually(checkTaints, input.WaitIntervals...).Should(Succeed())
 }