Fix flaky disable dry-run test

Author: jonathan-innis

test/pkg/environment/common/expectations.go

@@ -636,6 +636,7 @@ func (env *Environment) EventuallyExpectUniqueNodeNames(selector labels.Selector
 
 func (env *Environment) eventuallyExpectScaleDown() {
 	GinkgoHelper()
+	By(fmt.Sprintf("Expecting the cluster to scale down to %d nodes", env.StartingNodeCount))
 	Eventually(func(g Gomega) {
 		// expect the current node count to be what it was when the test started
 		g.Expect(env.Monitor.NodeCount()).To(Equal(env.StartingNodeCount))

test/suites/integration/nodeclass_test.go

@@ -15,6 +15,7 @@ limitations under the License.
 package integration_test
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -23,13 +24,14 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/iam"
 	"github.com/awslabs/operatorpkg/status"
 	"github.com/google/uuid"
-	appsv1 "k8s.io/api/apps/v1"
+	"github.com/samber/lo"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	v1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1"
+	awserrors "github.com/aws/karpenter-provider-aws/pkg/errors"
 
 	. "github.com/awslabs/operatorpkg/test/expectations"
 	. "github.com/onsi/ginkgo/v2"
@@ -41,6 +43,22 @@ var _ = Describe("NodeClass IAM Permissions", func() {
 		roleName   string
 		policyName string
 	)
+	AfterEach(func(ctx context.Context) {
+		By("deleting the deny policy")
+		_, err := env.IAMAPI.DeleteRolePolicy(ctx, &iam.DeleteRolePolicyInput{
+			RoleName:   aws.String(roleName),
+			PolicyName: aws.String(policyName),
+		})
+		Expect(awserrors.IgnoreNotFound(err)).To(BeNil())
+		env.ExpectSettingsOverridden(corev1.EnvVar{Name: "DISABLE_DRY_RUN", Value: "false"}) // re-enable dry-run
+
+		By("waiting for NodeClass to report ValidationSucceeded=true")
+		// Ensure that we don't start the next test until the policy has been deleted and IAM returns that we have permissions back
+		EventuallyExpectWithDryRunCacheInvalidation(func(g Gomega, currGeneration int64) {
+			g.Expect(env.Client.Get(ctx, client.ObjectKeyFromObject(nodeClass), nodeClass)).To(Succeed())
+			g.Expect(nodeClass.StatusConditions().IsTrue(v1.ConditionTypeValidationSucceeded) && nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).ObservedGeneration == currGeneration).To(BeTrue())
+		})
+	}, GracePeriod(time.Minute*15))
 	DescribeTable("IAM Permission Failure Tests",
 		func(action string, expectedMessage string) {
 			policyDoc := fmt.Sprintf(`{
@@ -54,41 +72,32 @@ var _ = Describe("NodeClass IAM Permissions", func() {
 					]
 				}`, action)
 
-			deployment := &appsv1.Deployment{}
-			Expect(env.Client.Get(env.Context, types.NamespacedName{
-				Namespace: "kube-system",
-				Name:      "karpenter",
-			}, deployment)).To(Succeed())
-
+			pods := env.ExpectKarpenterPods()
 			sa := &corev1.ServiceAccount{}
 			Expect(env.Client.Get(env.Context, types.NamespacedName{
 				Namespace: "kube-system",
-				Name:      deployment.Spec.Template.Spec.ServiceAccountName,
+				Name:      pods[0].Spec.ServiceAccountName,
 			}, sa)).To(Succeed())
 
 			roleName = strings.Split(sa.Annotations["eks.amazonaws.com/role-arn"], "/")[1]
 			policyName = fmt.Sprintf("TestPolicy-%s", uuid.New().String())
 
+			By(fmt.Sprintf("adding the deny policy for %s", action))
 			_, err := env.IAMAPI.PutRolePolicy(env.Context, &iam.PutRolePolicyInput{
 				RoleName:       aws.String(roleName),
 				PolicyName:     aws.String(policyName),
 				PolicyDocument: aws.String(policyDoc),
 			})
 			Expect(err).To(BeNil())
-			DeferCleanup(func() {
-				_, err := env.IAMAPI.DeleteRolePolicy(env.Context, &iam.DeleteRolePolicyInput{
-					RoleName:   aws.String(roleName),
-					PolicyName: aws.String(policyName),
-				})
-				Expect(err).To(BeNil())
-			})
 
 			env.ExpectCreated(nodeClass)
-			Eventually(func(g Gomega) {
+
+			By("waiting for NodeClass to report ValidationSucceeded=False")
+			EventuallyExpectWithDryRunCacheInvalidation(func(g Gomega, currGeneration int64) {
 				g.Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(nodeClass), nodeClass)).To(Succeed())
-				g.Expect(nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).IsFalse()).To(BeTrue())
+				g.Expect(nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).IsFalse() && nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).ObservedGeneration == currGeneration).To(BeTrue())
 				g.Expect(nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).Reason).To(Equal(expectedMessage))
-			}).Should(Succeed())
+			})
 			ExpectStatusConditions(env, env.Client, 1*time.Minute, nodeClass, status.Condition{Type: status.ConditionReady, Status: metav1.ConditionFalse, Message: "ValidationSucceeded=False"})
 		},
 		Entry("should fail when CreateFleet is denied",
@@ -121,39 +130,55 @@ var _ = Describe("NodeClass IAM Permissions", func() {
 				}
 			]
 		}`
-		deployment := &appsv1.Deployment{}
-		Expect(env.Client.Get(env.Context, types.NamespacedName{
-			Namespace: "kube-system",
-			Name:      "karpenter",
-		}, deployment)).To(Succeed())
-
+		pods := env.ExpectKarpenterPods()
 		sa := &corev1.ServiceAccount{}
 		Expect(env.Client.Get(env.Context, types.NamespacedName{
 			Namespace: "kube-system",
-			Name:      deployment.Spec.Template.Spec.ServiceAccountName,
+			Name:      pods[0].Spec.ServiceAccountName,
 		}, sa)).To(Succeed())
 
 		roleName = strings.Split(sa.Annotations["eks.amazonaws.com/role-arn"], "/")[1]
 		policyName = fmt.Sprintf("TestPolicy-%s", uuid.New().String())
 
+		By("adding the deny policy")
 		_, err := env.IAMAPI.PutRolePolicy(env.Context, &iam.PutRolePolicyInput{
 			RoleName:       aws.String(roleName),
 			PolicyName:     aws.String(policyName),
 			PolicyDocument: aws.String(policyDoc),
 		})
 		Expect(err).To(BeNil())
-		DeferCleanup(func() {
-			_, err := env.IAMAPI.DeleteRolePolicy(env.Context, &iam.DeleteRolePolicyInput{
-				RoleName:   aws.String(roleName),
-				PolicyName: aws.String(policyName),
-			})
-			Expect(err).To(BeNil())
+
+		// Ensure that we properly validate that we see the new policy
+		env.ExpectCreated(nodeClass)
+
+		By("waiting for NodeClass to report ValidationSucceeded=False")
+		EventuallyExpectWithDryRunCacheInvalidation(func(g Gomega, currGeneration int64) {
+			g.Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(nodeClass), nodeClass)).To(Succeed())
+			g.Expect(nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).IsFalse() && nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).ObservedGeneration == currGeneration).To(BeTrue())
 		})
+
+		// Disable dry-run and now validation should succeed
 		env.ExpectSettingsOverridden(corev1.EnvVar{Name: "DISABLE_DRY_RUN", Value: "true"})
-		env.ExpectCreated(nodeClass)
 		Eventually(func(g Gomega) {
 			g.Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(nodeClass), nodeClass)).To(Succeed())
 			g.Expect(nodeClass.StatusConditions().Get(v1.ConditionTypeValidationSucceeded).IsTrue()).To(BeTrue())
 		}).Should(Succeed())
 	})
 })
+
+func EventuallyExpectWithDryRunCacheInvalidation(f func(g Gomega, currentGeneration int64)) {
+	lastUpdate := time.Time{}
+	lastGeneration := int64(0)
+
+	Eventually(func(g Gomega) {
+		// Force the cache to be reset by updating MetadataOptions
+		if time.Since(lastUpdate) > time.Second*15 {
+			stored := nodeClass.DeepCopy()
+			nodeClass.Spec.MetadataOptions.HTTPPutResponseHopLimit = lo.ToPtr(lo.FromPtr(nodeClass.Spec.MetadataOptions.HTTPPutResponseHopLimit) + 1)
+			g.Expect(env.Client.Patch(env.Context, nodeClass, client.MergeFrom(stored))).To(Succeed())
+			lastUpdate = time.Now()
+			lastGeneration = nodeClass.Generation
+		}
+		f(g, lastGeneration)
+	}).Should(Succeed())
+})