test: Add e2e testing for deprecated AMIs (#6990)

aws · Oct 14, 2024 · 15e94a8 · 15e94a8
1 parent 8a00cd0
commit 15e94a8
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 16 deletions.
diff --git a/test/pkg/environment/aws/expectations.go b/test/pkg/environment/aws/expectations.go
@@ -362,6 +362,55 @@ func (env *Environment) GetAMIBySSMPath(ssmPath string) string {
 	return *parameter.Parameter.Value
 }
 
+func (env *Environment) GetDeprecatedAMI(amiID string, amifamily string) string {
+	out, err := env.EC2API.DescribeImages(&ec2.DescribeImagesInput{
+		Filters: []*ec2.Filter{
+			{
+				Name:   lo.ToPtr(fmt.Sprintf("tag:%s", coretest.DiscoveryLabel)),
+				Values: []*string{lo.ToPtr(env.K8sVersion())},
+			},
+			{
+				Name:   lo.ToPtr("tag:amiFamily"),
+				Values: []*string{lo.ToPtr(amifamily)},
+			},
+		},
+		IncludeDeprecated: lo.ToPtr(true),
+	})
+	Expect(err).To(BeNil())
+	if len(out.Images) == 1 {
+		return lo.FromPtr(out.Images[0].ImageId)
+	}
+
+	input := &ec2.CopyImageInput{
+		SourceImageId: lo.ToPtr(amiID),
+		Name:          lo.ToPtr(fmt.Sprintf("deprecated-%s-%s-%s", amiID, amifamily, env.K8sVersion())),
+		SourceRegion:  lo.ToPtr(env.Region),
+		TagSpecifications: []*ec2.TagSpecification{
+			{ResourceType: lo.ToPtr(ec2.ResourceTypeImage), Tags: []*ec2.Tag{
+				{
+					Key:   lo.ToPtr(coretest.DiscoveryLabel),
+					Value: lo.ToPtr(env.K8sVersion()),
+				},
+				{
+					Key:   lo.ToPtr("amiFamily"),
+					Value: lo.ToPtr(amifamily),
+				},
+			}},
+		},
+	}
+	output, err := env.EC2API.CopyImage(input)
+	Expect(err).To(BeNil())
+
+	deprecated, err := env.EC2API.EnableImageDeprecationWithContext(env.Context, &ec2.EnableImageDeprecationInput{
+		ImageId:     output.ImageId,
+		DeprecateAt: lo.ToPtr(time.Now()),
+	})
+	Expect(err).To(BeNil())
+	Expect(lo.FromPtr(deprecated.Return)).To(BeTrue())
+
+	return lo.FromPtr(output.ImageId)
+}
+
 func (env *Environment) EventuallyExpectRunInstances(instanceInput *ec2.RunInstancesInput) *ec2.Reservation {
 	GinkgoHelper()
 	// implement IMDSv2

diff --git a/test/suites/ami/suite_test.go b/test/suites/ami/suite_test.go
@@ -71,8 +71,10 @@ var _ = AfterEach(func() { env.AfterEach() })
 
 var _ = Describe("AMI", func() {
 	var customAMI string
+	var deprecatedAMI string
 	BeforeEach(func() {
 		customAMI = env.GetAMIBySSMPath(fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/x86_64/standard/recommended/image_id", env.K8sVersion()))
+		deprecatedAMI = env.GetDeprecatedAMI(customAMI, "AL2023")
 	})
 
 	It("should use the AMI defined by the AMI Selector Terms", func() {
@@ -156,6 +158,39 @@ var _ = Describe("AMI", func() {
 
 		env.ExpectInstance(pod.Spec.NodeName).To(HaveField("ImageId", HaveValue(Equal(customAMI))))
 	})
+	It("should support launching nodes with a deprecated ami", func() {
+		nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
+		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{
+			{
+				ID: deprecatedAMI,
+			},
+		}
+		pod := coretest.Pod()
+
+		env.ExpectCreated(pod, nodeClass, nodePool)
+		env.EventuallyExpectHealthy(pod)
+		env.ExpectCreatedNodeCount("==", 1)
+
+		env.ExpectInstance(pod.Spec.NodeName).To(HaveField("ImageId", HaveValue(Equal(deprecatedAMI))))
+	})
+	It("should prioritize launch with non-deprecated AMIs", func() {
+		nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
+		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{
+			{
+				ID: deprecatedAMI,
+			},
+			{
+				ID: customAMI,
+			},
+		}
+		pod := coretest.Pod()
+
+		env.ExpectCreated(pod, nodeClass, nodePool)
+		env.EventuallyExpectHealthy(pod)
+		env.ExpectCreatedNodeCount("==", 1)
+
+		env.ExpectInstance(pod.Spec.NodeName).To(HaveField("ImageId", HaveValue(Equal(customAMI))))
+	})
 
 	Context("AMIFamily", func() {
 		DescribeTable(

diff --git a/test/suites/drift/suite_test.go b/test/suites/drift/suite_test.go
@@ -48,6 +48,7 @@ import (
 
 var env *aws.Environment
 var amdAMI string
+var deprecatedAMI string
 var nodeClass *v1.EC2NodeClass
 var nodePool *karpv1.NodePool
 
@@ -76,6 +77,7 @@ var _ = Describe("Drift", func() {
 	var numPods int
 	BeforeEach(func() {
 		amdAMI = env.GetAMIBySSMPath(fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/x86_64/standard/recommended/image_id", env.K8sVersion()))
+		deprecatedAMI = env.GetDeprecatedAMI(amdAMI, "AL2023")
 		numPods = 1
 		// Add pods with a do-not-disrupt annotation so that we can check node metadata before we disrupt
 		dep = coretest.Deployment(coretest.DeploymentOptions{
@@ -363,18 +365,17 @@ var _ = Describe("Drift", func() {
 		env.EventuallyExpectNotFound(pod, nodeClaim, node)
 		env.EventuallyExpectHealthyPodCount(selector, numPods)
 	})
-	It("should return drifted if the AMI no longer matches the existing NodeClaims instance type", func() {
-		armAMI := env.GetAMIBySSMPath(fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/arm64/standard/recommended/image_id", env.K8sVersion()))
+	It("should disrupt nodes for deprecated AMIs to non-deprecated AMIs", func() {
 		nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
-		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: armAMI}}
+		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: deprecatedAMI}}
 
 		env.ExpectCreated(dep, nodeClass, nodePool)
 		pod := env.EventuallyExpectHealthyPodCount(selector, numPods)[0]
 		env.ExpectCreatedNodeCount("==", 1)
 
 		nodeClaim := env.EventuallyExpectCreatedNodeClaimCount("==", 1)[0]
 		node := env.EventuallyExpectNodeCount("==", 1)[0]
-		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: amdAMI}}
+		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: amdAMI}, {ID: deprecatedAMI}}
 		env.ExpectCreatedOrUpdated(nodeClass)
 
 		env.EventuallyExpectDrifted(nodeClaim)
@@ -383,26 +384,32 @@ var _ = Describe("Drift", func() {
 		env.ExpectUpdated(pod)
 		env.EventuallyExpectNotFound(pod, nodeClaim, node)
 		env.EventuallyExpectHealthyPodCount(selector, numPods)
-	})
-	It("should not disrupt nodes that have drifted without the featureGate enabled", func() {
-		env.ExpectSettingsOverridden(corev1.EnvVar{Name: "FEATURE_GATES", Value: "Drift=false"})
 
-		oldCustomAMI := env.GetAMIBySSMPath(fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/x86_64/standard/recommended/image_id", env.K8sVersionWithOffset(1)))
+		// validate the AMI id matches the non-deprecated AMI
+		pod = env.EventuallyExpectHealthyPodCount(selector, numPods)[0]
+		env.ExpectInstance(pod.Spec.NodeName).To(HaveField("ImageId", HaveValue(Equal(amdAMI))))
+
+	})
+	It("should return drifted if the AMI no longer matches the existing NodeClaims instance type", func() {
+		armAMI := env.GetAMIBySSMPath(fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/arm64/standard/recommended/image_id", env.K8sVersion()))
 		nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
-		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: oldCustomAMI}}
+		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: armAMI}}
 
 		env.ExpectCreated(dep, nodeClass, nodePool)
-		env.EventuallyExpectHealthyPodCount(selector, numPods)
+		pod := env.EventuallyExpectHealthyPodCount(selector, numPods)[0]
 		env.ExpectCreatedNodeCount("==", 1)
 
-		node := env.Monitor.CreatedNodes()[0]
+		nodeClaim := env.EventuallyExpectCreatedNodeClaimCount("==", 1)[0]
+		node := env.EventuallyExpectNodeCount("==", 1)[0]
 		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: amdAMI}}
-		env.ExpectUpdated(nodeClass)
+		env.ExpectCreatedOrUpdated(nodeClass)
 
-		// We should consistently get the same node existing for a minute
-		Consistently(func(g Gomega) {
-			g.Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), &corev1.Node{})).To(Succeed())
-		}).WithTimeout(time.Minute).Should(Succeed())
+		env.EventuallyExpectDrifted(nodeClaim)
+
+		delete(pod.Annotations, karpv1.DoNotDisruptAnnotationKey)
+		env.ExpectUpdated(pod)
+		env.EventuallyExpectNotFound(pod, nodeClaim, node)
+		env.EventuallyExpectHealthyPodCount(selector, numPods)
 	})
 	It("should disrupt nodes that have drifted due to securitygroup", func() {
 		By("getting the cluster vpc id")