Add unit test

Author: ec2-bot

pkg/node/node.go

@@ -145,7 +145,7 @@ func (n Node) CordonAndDrain(nodeName string, reason string, recorder recorderIn
 	}
 	if n.nthConfig.UseAPIServerCacheToListPods {
 		if pods != nil {
-			pods = n.filterOutDaemonSetPods(pods)
+			pods = n.FilterOutDaemonSetPods(pods)
 			err = n.drainHelper.DeleteOrEvictPods(pods.Items)
 		}
 	} else {
@@ -649,8 +649,8 @@ func (n Node) fetchAllPods(nodeName string) (*corev1.PodList, error) {
 	return n.drainHelper.Client.CoreV1().Pods("").List(context.TODO(), listOptions)
 }
 
-// filterOutDaemonSetPods filters a list of pods to exclude DaemonSet pods when IgnoreDaemonSets is enabled
-func (n *Node) filterOutDaemonSetPods(pods *corev1.PodList) *corev1.PodList {
+// FilterOutDaemonSetPods filters a list of pods to exclude DaemonSet pods when IgnoreDaemonSets is enabled
+func (n *Node) FilterOutDaemonSetPods(pods *corev1.PodList) *corev1.PodList {
 	if !n.nthConfig.IgnoreDaemonSets {
 		return pods
 	}

pkg/node/node_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/aws/aws-node-termination-handler/pkg/node"
 	h "github.com/aws/aws-node-termination-handler/pkg/test"
 	"github.com/aws/aws-node-termination-handler/pkg/uptime"
+	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes/fake"
@@ -377,3 +378,43 @@ func TestUncordonIfRebootedTimeParseFailure(t *testing.T) {
 	err = tNode.UncordonIfRebooted(nodeName)
 	h.Assert(t, err != nil, "Failed to return error on UncordonIfReboted failure to parse time")
 }
+
+func TestFilterOutDaemonSetPods(t *testing.T) {
+	tNode, err := newNode(config.Config{IgnoreDaemonSets: true}, fake.NewSimpleClientset())
+	h.Ok(t, err)
+
+	mockPodList := &corev1.PodList{
+		Items: []corev1.Pod{
+			{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "mock-daemon-pod",
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							Kind: "DaemonSet",
+							Name: "daemon-1",
+						},
+					},
+				},
+			},
+			{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "mock-replica-pod",
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							Kind: "ReplicaSet",
+							Name: "replica-1",
+						},
+					},
+				},
+			},
+			{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "mock-regular-pod",
+				},
+			},
+		},
+	}
+
+	filteredMockPodList := tNode.FilterOutDaemonSetPods(mockPodList)
+	h.Equals(t, 2, len(filteredMockPodList.Items))
+}