From 2e302502001a78ac37a06d67fd477d7cfe5af728 Mon Sep 17 00:00:00 2001
From: Jason Deal <jmdeal@amazon.com>
Date: Fri, 26 Jul 2024 10:07:28 -0700
Subject: [PATCH] add functional tests

---
 .../node/termination/suite_test.go            | 74 +++++++++++++++++++
 pkg/test/storage.go                           | 25 +++++++
 pkg/utils/node/node.go                        |  6 +-
 3 files changed, 101 insertions(+), 4 deletions(-)

diff --git a/pkg/controllers/node/termination/suite_test.go b/pkg/controllers/node/termination/suite_test.go
index 26a71cfdaf..66305d3253 100644
--- a/pkg/controllers/node/termination/suite_test.go
+++ b/pkg/controllers/node/termination/suite_test.go
@@ -763,6 +763,80 @@ var _ = Describe("Termination", func() {
 			ExpectSingletonReconciled(ctx, queue)
 			ExpectDeleted(ctx, env.Client, pod)
 		})
+		Context("VolumeAttachments", func() {
+			It("should wait for volume attachments", func() {
+				va := test.VolumeAttachment(test.VolumeAttachmentOptions{
+					NodeName:   node.Name,
+					VolumeName: "foo",
+				})
+				ExpectApplied(ctx, env.Client, node, nodeClaim, nodePool, va)
+				Expect(env.Client.Delete(ctx, node)).To(Succeed())
+
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectExists(ctx, env.Client, node)
+
+				ExpectDeleted(ctx, env.Client, va)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectNotFound(ctx, env.Client, node)
+			})
+			It("should only wait for volume attachments associated with drainable pods", func() {
+				vaDrainable := test.VolumeAttachment(test.VolumeAttachmentOptions{
+					NodeName:   node.Name,
+					VolumeName: "foo",
+				})
+				vaNonDrainable := test.VolumeAttachment(test.VolumeAttachmentOptions{
+					NodeName:   node.Name,
+					VolumeName: "bar",
+				})
+				pvc := test.PersistentVolumeClaim(test.PersistentVolumeClaimOptions{
+					VolumeName: "bar",
+				})
+				pod := test.Pod(test.PodOptions{
+					ObjectMeta: metav1.ObjectMeta{
+						OwnerReferences: defaultOwnerRefs,
+					},
+					Tolerations: []corev1.Toleration{{
+						Key:      v1.DisruptionTaintKey,
+						Operator: corev1.TolerationOpExists,
+					}},
+					PersistentVolumeClaims: []string{pvc.Name},
+				})
+				ExpectApplied(ctx, env.Client, node, nodeClaim, nodePool, vaDrainable, vaNonDrainable, pod, pvc)
+				ExpectManualBinding(ctx, env.Client, pod, node)
+				Expect(env.Client.Delete(ctx, node)).To(Succeed())
+
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectExists(ctx, env.Client, node)
+
+				ExpectDeleted(ctx, env.Client, vaDrainable)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectNotFound(ctx, env.Client, node)
+			})
+			It("should wait for volume attachments until the nodeclaim's termination grace period expires", func() {
+				va := test.VolumeAttachment(test.VolumeAttachmentOptions{
+					NodeName:   node.Name,
+					VolumeName: "foo",
+				})
+				nodeClaim.Annotations = map[string]string{
+					v1.NodeClaimTerminationTimestampAnnotationKey: fakeClock.Now().Add(time.Minute).Format(time.RFC3339),
+				}
+				ExpectApplied(ctx, env.Client, node, nodeClaim, nodePool, va)
+				Expect(env.Client.Delete(ctx, node)).To(Succeed())
+
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectExists(ctx, env.Client, node)
+
+				fakeClock.Step(5 * time.Minute)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectObjectReconciled(ctx, env.Client, terminationController, node)
+				ExpectNotFound(ctx, env.Client, node)
+			})
+		})
 	})
 	Context("Metrics", func() {
 		It("should fire the terminationSummary metric when deleting nodes", func() {
diff --git a/pkg/test/storage.go b/pkg/test/storage.go
index ab0dce5ef6..96fe7c41ce 100644
--- a/pkg/test/storage.go
+++ b/pkg/test/storage.go
@@ -153,3 +153,28 @@ func StorageClass(overrides ...StorageClassOptions) *storagev1.StorageClass {
 		VolumeBindingMode: options.VolumeBindingMode,
 	}
 }
+
+type VolumeAttachmentOptions struct {
+	metav1.ObjectMeta
+	NodeName   string
+	VolumeName string
+}
+
+func VolumeAttachment(overrides ...VolumeAttachmentOptions) *storagev1.VolumeAttachment {
+	options := VolumeAttachmentOptions{}
+	for _, opts := range overrides {
+		if err := mergo.Merge(&options, opts, mergo.WithOverride); err != nil {
+			panic(fmt.Sprintf("Failed to merge options: %s", err))
+		}
+	}
+	return &storagev1.VolumeAttachment{
+		ObjectMeta: ObjectMeta(options.ObjectMeta),
+		Spec: storagev1.VolumeAttachmentSpec{
+			NodeName: options.NodeName,
+			Attacher: "fake-csi",
+			Source: storagev1.VolumeAttachmentSource{
+				PersistentVolumeName: lo.ToPtr(options.VolumeName),
+			},
+		},
+	}
+}
diff --git a/pkg/utils/node/node.go b/pkg/utils/node/node.go
index 4bc2152e5c..c60bd5e056 100644
--- a/pkg/utils/node/node.go
+++ b/pkg/utils/node/node.go
@@ -81,11 +81,9 @@ func GetProvisionablePods(ctx context.Context, kubeClient client.Client) ([]*cor
 func GetVolumeAttachments(ctx context.Context, kubeClient client.Client, node *corev1.Node) ([]*storagev1.VolumeAttachment, error) {
 	var volumeAttachmentList storagev1.VolumeAttachmentList
 	if err := kubeClient.List(ctx, &volumeAttachmentList, client.MatchingFields{"spec.nodeName": node.Name}); err != nil {
-		return nil, fmt.Errorf("listing volumeattachments, %w", err)
+		return nil, fmt.Errorf("listing volumeAttachments, %w", err)
 	}
-	return lo.FilterMap(volumeAttachmentList.Items, func(v storagev1.VolumeAttachment, _ int) (*storagev1.VolumeAttachment, bool) {
-		return &v, v.Spec.NodeName == node.Name
-	}), nil
+	return lo.ToSlicePtr(volumeAttachmentList.Items), nil
 }
 
 func GetCondition(n *corev1.Node, match corev1.NodeConditionType) corev1.NodeCondition {