From 6b52d7b2da910212aca1a64ed884899f7425f252 Mon Sep 17 00:00:00 2001
From: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
Date: Mon, 15 Apr 2024 18:16:30 +0200
Subject: [PATCH] E2E test: Create CleanupNFDObjects func

Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
---
 tests/e2e/common/kubernetes.go          | 33 +++++++++++++++----------
 tests/e2e/gpu-feature-discovery_test.go |  2 +-
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/tests/e2e/common/kubernetes.go b/tests/e2e/common/kubernetes.go
index 3040a2c81..3a5594205 100644
--- a/tests/e2e/common/kubernetes.go
+++ b/tests/e2e/common/kubernetes.go
@@ -165,18 +165,18 @@ func CleanupNode(ctx context.Context, cs clientset.Interface) {
 	}
 }
 
-func CleanupNodeFeatures(ctx context.Context, cli *nfdclient.Clientset, namespace string) {
-	// Clean NodeFeatureRule objects
-	nfrs, err := cli.NfdV1alpha1().NodeFeatureRules().List(ctx, metav1.ListOptions{})
+// cleanupNodeFeatures deletes all NodeFeature objects in the given namespace
+func cleanupNodeFeatures(ctx context.Context, cli *nfdclient.Clientset, namespace string) {
+	nfs, err := cli.NfdV1alpha1().NodeFeatures(namespace).List(ctx, metav1.ListOptions{})
 	if errors.IsNotFound(err) {
 		// Omitted error, nothing to do.
 	} else {
 		Expect(err).NotTo(HaveOccurred())
 
-		if len(nfrs.Items) != 0 {
-			By("Deleting NodeFeatureRule objects from the cluster")
-			for _, nfr := range nfrs.Items {
-				err = cli.NfdV1alpha1().NodeFeatureRules().Delete(ctx, nfr.Name, metav1.DeleteOptions{})
+		if len(nfs.Items) != 0 {
+			By("Deleting NodeFeature objects from namespace " + namespace)
+			for _, nf := range nfs.Items {
+				err = cli.NfdV1alpha1().NodeFeatures(namespace).Delete(ctx, nf.Name, metav1.DeleteOptions{})
 				if errors.IsNotFound(err) {
 					// Omitted error
 					continue
@@ -186,18 +186,20 @@ func CleanupNodeFeatures(ctx context.Context, cli *nfdclient.Clientset, namespac
 			}
 		}
 	}
+}
 
-	// Clean NodeFeature objects
-	nfs, err := cli.NfdV1alpha1().NodeFeatures(namespace).List(ctx, metav1.ListOptions{})
+// cleanupNodeFeatureRules deletes all NodeFeatureRule objects
+func cleanupNodeFeatureRules(ctx context.Context, cli *nfdclient.Clientset) {
+	nfrs, err := cli.NfdV1alpha1().NodeFeatureRules().List(ctx, metav1.ListOptions{})
 	if errors.IsNotFound(err) {
 		// Omitted error, nothing to do.
 	} else {
 		Expect(err).NotTo(HaveOccurred())
 
-		if len(nfs.Items) != 0 {
-			By("Deleting NodeFeature objects from namespace " + namespace)
-			for _, nf := range nfs.Items {
-				err = cli.NfdV1alpha1().NodeFeatures(namespace).Delete(ctx, nf.Name, metav1.DeleteOptions{})
+		if len(nfrs.Items) != 0 {
+			By("Deleting NodeFeatureRule objects from the cluster")
+			for _, nfr := range nfrs.Items {
+				err = cli.NfdV1alpha1().NodeFeatureRules().Delete(ctx, nfr.Name, metav1.DeleteOptions{})
 				if errors.IsNotFound(err) {
 					// Omitted error
 					continue
@@ -208,3 +210,8 @@ func CleanupNodeFeatures(ctx context.Context, cli *nfdclient.Clientset, namespac
 		}
 	}
 }
+
+func CleanupNFDObjects(ctx context.Context, cli *nfdclient.Clientset, namespace string) {
+	cleanupNodeFeatureRules(ctx, cli)
+	cleanupNodeFeatures(ctx, cli, namespace)
+}
diff --git a/tests/e2e/gpu-feature-discovery_test.go b/tests/e2e/gpu-feature-discovery_test.go
index 30e86c66e..c17d14dc6 100644
--- a/tests/e2e/gpu-feature-discovery_test.go
+++ b/tests/e2e/gpu-feature-discovery_test.go
@@ -145,7 +145,7 @@ var _ = NVDescribe("GPU Feature Discovery", func() {
 			Expect(err).NotTo(HaveOccurred())
 			// Cleanup node
 			common.CleanupNode(ctx, f.ClientSet)
-			common.CleanupNodeFeatures(ctx, nfdClient, f.Namespace.Name)
+			common.CleanupNFDObjects(ctx, nfdClient, f.Namespace.Name)
 		})
 
 		AfterAll(func(ctx context.Context) {