From 2bf575b815ad44b8cfd414ff79e7c580ac298e89 Mon Sep 17 00:00:00 2001
From: Ignas Baranauskas <ibaranau@redhat.com>
Date: Wed, 13 Nov 2024 15:43:19 +0000
Subject: [PATCH] fix: removes label selector from old e2e tests

---
 tests/e2e/heterogeneous_clusters_kind_test.py |  2 +-
 .../e2e/heterogeneous_clusters_oauth_test.py  |  2 +-
 tests/e2e/support.py                          | 40 +++++++++++++------
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/tests/e2e/heterogeneous_clusters_kind_test.py b/tests/e2e/heterogeneous_clusters_kind_test.py
index 8f814a7c4..052fa7b80 100644
--- a/tests/e2e/heterogeneous_clusters_kind_test.py
+++ b/tests/e2e/heterogeneous_clusters_kind_test.py
@@ -25,7 +25,7 @@ def teardown_method(self):
     @pytest.mark.nvidia_gpu
     def test_heterogeneous_clusters(self):
         create_namespace(self)
-        create_kueue_resources(self, 2)
+        create_kueue_resources(self, 2, with_labels=True, with_tolerations=True)
         self.run_heterogeneous_clusters()
 
     def run_heterogeneous_clusters(
diff --git a/tests/e2e/heterogeneous_clusters_oauth_test.py b/tests/e2e/heterogeneous_clusters_oauth_test.py
index 4a7a687cf..d57cff48d 100644
--- a/tests/e2e/heterogeneous_clusters_oauth_test.py
+++ b/tests/e2e/heterogeneous_clusters_oauth_test.py
@@ -24,7 +24,7 @@ def teardown_method(self):
 
     def test_heterogeneous_clusters(self):
         create_namespace(self)
-        create_kueue_resources(self, 2)
+        create_kueue_resources(self, 2, with_labels=True, with_tolerations=True)
         self.run_heterogeneous_clusters()
 
     def run_heterogeneous_clusters(
diff --git a/tests/e2e/support.py b/tests/e2e/support.py
index 604884668..d76b460c8 100644
--- a/tests/e2e/support.py
+++ b/tests/e2e/support.py
@@ -66,12 +66,14 @@ def create_namespace(self):
         return RuntimeError(e)
 
 
-def create_new_resource_flavor(self, num_flavors):
+def create_new_resource_flavor(self, num_flavors, with_labels, with_tolerations):
     self.resource_flavors = []
     for i in range(num_flavors):
         default = i < 1
         resource_flavor = f"test-resource-flavor-{random_choice()}"
-        create_resource_flavor(self, resource_flavor, default)
+        create_resource_flavor(
+            self, resource_flavor, default, with_labels, with_tolerations
+        )
         self.resource_flavors.append(resource_flavor)
 
 
@@ -173,7 +175,9 @@ def create_cluster_queue(self, cluster_queue, flavor):
     self.cluster_queue = cluster_queue
 
 
-def create_resource_flavor(self, flavor, default=True):
+def create_resource_flavor(
+    self, flavor, default=True, with_labels=False, with_tolerations=False
+):
     worker_label, worker_value = os.getenv("WORKER_LABEL", "worker-1=true").split("=")
     control_label, control_value = os.getenv(
         "CONTROL_LABEL", "ingress-ready=true"
@@ -182,9 +186,11 @@ def create_resource_flavor(self, flavor, default=True):
         "TOLERATION_KEY", "node-role.kubernetes.io/control-plane"
     )
 
-    node_labels = (
-        {worker_label: worker_value} if default else {control_label: control_value}
-    )
+    node_labels = {}
+    if with_labels:
+        node_labels = (
+            {worker_label: worker_value} if default else {control_label: control_value}
+        )
 
     resource_flavor_json = {
         "apiVersion": "kueue.x-k8s.io/v1beta1",
@@ -192,13 +198,19 @@ def create_resource_flavor(self, flavor, default=True):
         "metadata": {"name": flavor},
         "spec": {
             "nodeLabels": node_labels,
-            "tolerations": [
+            **(
                 {
-                    "key": toleration_key,
-                    "operator": "Exists",
-                    "effect": "NoSchedule",
+                    "tolerations": [
+                        {
+                            "key": toleration_key,
+                            "operator": "Exists",
+                            "effect": "NoSchedule",
+                        }
+                    ]
                 }
-            ],
+                if with_tolerations
+                else {}
+            ),
         },
     }
 
@@ -260,9 +272,11 @@ def create_local_queue(self, cluster_queue, local_queue, is_default=True):
     self.local_queue = local_queue
 
 
-def create_kueue_resources(self, resource_ammount=1):
+def create_kueue_resources(
+    self, resource_ammount=1, with_labels=False, with_tolerations=False
+):
     print("creating Kueue resources ...")
-    create_new_resource_flavor(self, resource_ammount)
+    create_new_resource_flavor(self, resource_ammount, with_labels, with_tolerations)
     create_new_cluster_queue(self, resource_ammount)
     create_new_local_queue(self, resource_ammount)