From 431e1c38406e67bca1f4a554179940e87380cfe0 Mon Sep 17 00:00:00 2001
From: Amanuel Engeda <aengeda@amazon.com>
Date: Tue, 23 May 2023 11:42:37 -0700
Subject: [PATCH] Add e2e test for Node Requirement Drift

---
 test/suites/drift/suite_test.go               | 51 ++++++++++++++++++-
 .../en/preview/concepts/deprovisioning.md     |  2 +-
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/test/suites/drift/suite_test.go b/test/suites/drift/suite_test.go
index 7c86d3adb1c5..a6cb35397ea5 100644
--- a/test/suites/drift/suite_test.go
+++ b/test/suites/drift/suite_test.go
@@ -270,6 +270,50 @@ var _ = Describe("Drift", Label("AWS"), func() {
 		env.ExpectUpdated(pod)
 		env.EventuallyExpectNotFound(pod, node)
 	})
+	DescribeTable("NodeRequirement Drift", func(provisionerOption test.ProvisionerOptions) {
+		nodeTemplate := awstest.AWSNodeTemplate(v1alpha1.AWSNodeTemplateSpec{AWS: v1alpha1.AWS{
+			SecurityGroupSelector: map[string]string{"karpenter.sh/discovery": settings.FromContext(env.Context).ClusterName},
+			SubnetSelector:        map[string]string{"karpenter.sh/discovery": settings.FromContext(env.Context).ClusterName},
+		}})
+		provisioner := test.Provisioner(test.ProvisionerOptions{
+			Requirements: []v1.NodeSelectorRequirement{{Key: v1alpha5.LabelCapacityType, Operator: v1.NodeSelectorOpIn, Values: []string{v1alpha5.CapacityTypeOnDemand}}},
+			ProviderRef:  &v1alpha5.MachineTemplateRef{Name: nodeTemplate.Name},
+		})
+		// Add a do-not-evict pod so that we can check node metadata before we deprovision
+		pod := test.Pod(test.PodOptions{
+			ObjectMeta: metav1.ObjectMeta{
+				Annotations: map[string]string{
+					v1alpha5.DoNotEvictPodAnnotationKey: "true",
+				},
+			},
+		})
+
+		provisionerOption.ObjectMeta = provisioner.ObjectMeta
+		updatedProvisioner := test.Provisioner(
+			test.ProvisionerOptions{ProviderRef: &v1alpha5.MachineTemplateRef{Name: nodeTemplate.Name}},
+			provisionerOption,
+		)
+
+		env.ExpectCreated(pod, nodeTemplate, provisioner)
+		machine := env.EventuallyExpectCreatedMachineCount("==", 1)[0]
+		node := env.EventuallyExpectCreatedNodeCount("==", 1)[0]
+		env.EventuallyExpectHealthy(pod)
+
+		env.ExpectCreatedOrUpdated(updatedProvisioner)
+
+		By("validating the drifted status condition has propagated")
+		EventuallyWithOffset(1, func(g Gomega) {
+			g.Expect(env.Client.Get(env, client.ObjectKeyFromObject(machine), machine)).To(Succeed())
+			g.Expect(machine.StatusConditions().GetCondition(v1alpha5.MachineDrifted)).ToNot(BeNil())
+			g.Expect(machine.StatusConditions().GetCondition(v1alpha5.MachineDrifted).IsTrue()).To(BeTrue())
+		}).Should(Succeed())
+
+		delete(pod.Annotations, v1alpha5.DoNotEvictPodAnnotationKey)
+		env.ExpectUpdated(pod)
+		env.EventuallyExpectNotFound(pod, node)
+	},
+		Entry("NodeRequirement Drift", test.ProvisionerOptions{Requirements: []v1.NodeSelectorRequirement{{Key: v1alpha5.LabelCapacityType, Operator: v1.NodeSelectorOpIn, Values: []string{v1alpha5.CapacityTypeSpot}}}}),
+	)
 	Describe("Provisioner Drift", func() {
 		var pod *v1.Pod
 		var nodeTemplate *v1alpha1.AWSNodeTemplate
@@ -279,7 +323,10 @@ var _ = Describe("Drift", Label("AWS"), func() {
 				SecurityGroupSelector: map[string]string{"karpenter.sh/discovery": settings.FromContext(env.Context).ClusterName},
 				SubnetSelector:        map[string]string{"karpenter.sh/discovery": settings.FromContext(env.Context).ClusterName},
 			}})
-			provisioner = test.Provisioner(test.ProvisionerOptions{ProviderRef: &v1alpha5.MachineTemplateRef{Name: nodeTemplate.Name}})
+			provisioner = test.Provisioner(test.ProvisionerOptions{
+				Requirements: []v1.NodeSelectorRequirement{{Key: v1alpha5.LabelCapacityType, Operator: v1.NodeSelectorOpIn, Values: []string{v1alpha5.CapacityTypeOnDemand}}},
+				ProviderRef:  &v1alpha5.MachineTemplateRef{Name: nodeTemplate.Name},
+			})
 			// Add a do-not-evict pod so that we can check node metadata before we deprovision
 			pod = test.Pod(test.PodOptions{
 				ObjectMeta: metav1.ObjectMeta{
@@ -289,7 +336,7 @@ var _ = Describe("Drift", Label("AWS"), func() {
 				},
 			})
 		})
-		DescribeTable("provisioner static drift", func(fieldName string, provisionerOption test.ProvisionerOptions) {
+		DescribeTable("provisioner drift", func(fieldName string, provisionerOption test.ProvisionerOptions) {
 			provisionerOption.ObjectMeta = provisioner.ObjectMeta
 			updatedProvisioner := test.Provisioner(
 				test.ProvisionerOptions{ProviderRef: &v1alpha5.MachineTemplateRef{Name: nodeTemplate.Name}},
diff --git a/website/content/en/preview/concepts/deprovisioning.md b/website/content/en/preview/concepts/deprovisioning.md
index 47a765a6e322..6510a2449c79 100644
--- a/website/content/en/preview/concepts/deprovisioning.md
+++ b/website/content/en/preview/concepts/deprovisioning.md
@@ -158,7 +158,7 @@ Read the [Drift Design](https://github.com/aws/karpenter-core/pull/366/files) fo
 | Taints                     |    x   |         |            |      x      |
 | Labels                     |    x   |         |            |      x      |
 | Annotations                |    x   |         |            |      x      |
-| Node Requirements          |        |    x    |            |             |
+| Node Requirements          |        |    x    |            |      x      |
 | Kubelet Configuration      |    x   |         |            |      x      |
 | Weight                     |        |         |      x     |     NA      |
 | Limits                     |        |         |      x     |     NA      |