From 68c3bf317bbd6f3d5a9b6a4249976c9408c5628e Mon Sep 17 00:00:00 2001
From: Markus Lehtonen <markus.lehtonen@intel.com>
Date: Wed, 5 Apr 2023 14:46:39 +0300
Subject: [PATCH] test/e2e: fix node cleanup function

The node cleanup function was not removing all NFD-labels. It omitted
NFD-originated labels that used a non-default label namespace. This
patch fixes the issue by getting all NFD-managed labels from the special
annotation (nfd.node.kubernetes.io/feature-labels).

The patch also adds the ability to cleanup extended resources in a
similar way. This will be needed by future work.

Also changes the order of cleaning up CRs and the node. It is the right
order as cleaning up the CRs may still update the node.
---
 test/e2e/node_feature_discovery_test.go | 112 ++++++++++++++++--------
 1 file changed, 74 insertions(+), 38 deletions(-)

diff --git a/test/e2e/node_feature_discovery_test.go b/test/e2e/node_feature_discovery_test.go
index 7ba2d6156f..360c44f576 100644
--- a/test/e2e/node_feature_discovery_test.go
+++ b/test/e2e/node_feature_discovery_test.go
@@ -78,60 +78,96 @@ const TestTaintNs = "nfd.node.kubernetes.io"
 // cleanupNode deletes all NFD-related metadata from the Node object, i.e.
 // labels and annotations
 func cleanupNode(cs clientset.Interface) {
-	nodeList, err := cs.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
-	Expect(err).NotTo(HaveOccurred())
+	// Per-node cleanup function
+	cleanup := func(nodeName string) error {
+		node, err := cs.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
+		Expect(err).NotTo(HaveOccurred())
 
-	for _, n := range nodeList.Items {
-		var err error
-		var node *corev1.Node
-		for retry := 0; retry < 5; retry++ {
-			node, err = cs.CoreV1().Nodes().Get(context.TODO(), n.Name, metav1.GetOptions{})
-			Expect(err).NotTo(HaveOccurred())
-
-			update := false
-			// Remove labels
-			for key := range node.Labels {
-				if strings.HasPrefix(key, nfdv1alpha1.FeatureLabelNs) {
-					delete(node.Labels, key)
-					update = true
-				}
+		update := false
+		updateStatus := false
+		// Gather info about all NFD-managed node assets outside the default prefix
+		nfdLabels := map[string]struct{}{}
+		for _, name := range strings.Split(node.Annotations[nfdv1alpha1.FeatureLabelsAnnotation], ",") {
+			if strings.Contains(name, "/") {
+				nfdLabels[name] = struct{}{}
+			}
+		}
+		nfdERs := map[string]struct{}{}
+		for _, name := range strings.Split(node.Annotations[nfdv1alpha1.ExtendedResourceAnnotation], ",") {
+			if strings.Contains(name, "/") {
+				nfdERs[name] = struct{}{}
 			}
+		}
+
+		// Remove labels
+		for key := range node.Labels {
+			_, ok := nfdLabels[key]
+			if ok || strings.HasPrefix(key, nfdv1alpha1.FeatureLabelNs) {
+				delete(node.Labels, key)
+				update = true
+			}
+		}
+
+		// Remove annotations
+		for key := range node.Annotations {
+			if strings.HasPrefix(key, nfdv1alpha1.AnnotationNs) {
+				delete(node.Annotations, key)
+				update = true
+			}
+		}
 
-			// Remove annotations
-			for key := range node.Annotations {
-				if strings.HasPrefix(key, nfdv1alpha1.AnnotationNs) {
-					delete(node.Annotations, key)
+		// Remove taints
+		for _, taint := range node.Spec.Taints {
+			if strings.HasPrefix(taint.Key, TestTaintNs) {
+				newTaints, removed := taintutils.DeleteTaint(node.Spec.Taints, &taint)
+				if removed {
+					node.Spec.Taints = newTaints
 					update = true
 				}
 			}
+		}
 
-			// Remove taints
-			for _, taint := range node.Spec.Taints {
-				if strings.HasPrefix(taint.Key, TestTaintNs) {
-					newTaints, removed := taintutils.DeleteTaint(node.Spec.Taints, &taint)
-					if removed {
-						node.Spec.Taints = newTaints
-						update = true
-					}
-				}
+		// Remove extended resources
+		for key := range node.Status.Capacity {
+			// We check for FeatureLabelNs as -resource-labels can create ERs there
+			_, ok := nfdERs[string(key)]
+			if ok || strings.HasPrefix(string(key), nfdv1alpha1.FeatureLabelNs) {
+				delete(node.Status.Capacity, key)
+				delete(node.Status.Allocatable, key)
+				updateStatus = true
 			}
+		}
 
-			if !update {
-				break
+		if updateStatus {
+			By("Deleting NFD extended resources from node " + nodeName)
+			if _, err := cs.CoreV1().Nodes().UpdateStatus(context.TODO(), node, metav1.UpdateOptions{}); err != nil {
+				return err
 			}
+		}
 
+		if update {
 			By("Deleting NFD labels, annotations and taints from node " + node.Name)
-			_, err = cs.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
-			if err != nil {
-				time.Sleep(100 * time.Millisecond)
-			} else {
-				break
+			if _, err := cs.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{}); err != nil {
+				return err
 			}
+		}
+		return nil
+	}
 
+	// Cleanup all nodes
+	nodeList, err := cs.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
+	Expect(err).NotTo(HaveOccurred())
+
+	for _, n := range nodeList.Items {
+		var err error
+		for retry := 0; retry < 5; retry++ {
+			if err = cleanup(n.Name); err == nil {
+				break
+			}
+			time.Sleep(100 * time.Millisecond)
 		}
 		Expect(err).NotTo(HaveOccurred())
 	}
-
 }
 
 func cleanupCRs(cli *nfdclient.Clientset, namespace string) {
@@ -217,8 +253,8 @@ var _ = SIGDescribe("NFD master and worker", func() {
 				Expect(err).NotTo(HaveOccurred())
 
 				// Remove pre-existing stale annotations and labels etc and CRDs
-				cleanupNode(f.ClientSet)
 				cleanupCRs(nfdClient, f.Namespace.Name)
+				cleanupNode(f.ClientSet)
 
 				// Launch nfd-master
 				By("Creating nfd master pod and nfd-master service")