Updated unit test for user labels

project-codeflare · May 9, 2024 · 9b49b59 · 9b49b59
1 parent 134e306
commit 9b49b59
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/src/codeflare_sdk/utils/generate_yaml.py b/src/codeflare_sdk/utils/generate_yaml.py
@@ -313,7 +313,7 @@ def write_components(
     output_file_name: str,
     namespace: str,
     local_queue: Optional[str],
-    user_labels: dict,
+    labels: dict,
 ):
     # Create the directory if it doesn't exist
     directory_path = os.path.dirname(output_file_name)
@@ -323,6 +323,7 @@ def write_components(
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
     open(output_file_name, "w").close()
     lq_name = local_queue or get_default_kueue_name(namespace)
+    cluster_labels = labels
     with open(output_file_name, "a") as outfile:
         for component in components:
             if "generictemplate" in component:
@@ -335,7 +336,7 @@ def write_components(
                     ]
                     labels = component["generictemplate"]["metadata"]["labels"]
                     labels.update({"kueue.x-k8s.io/queue-name": lq_name})
-                    labels.update(user_labels)
+                    labels.update(cluster_labels)
                 outfile.write("---\n")
                 yaml.dump(
                     component["generictemplate"], outfile, default_flow_style=False
@@ -348,11 +349,12 @@ def load_components(
     name: str,
     namespace: str,
     local_queue: Optional[str],
-    user_labels: dict,
+    labels: dict,
 ):
     component_list = []
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
     lq_name = local_queue or get_default_kueue_name(namespace)
+    cluster_labels = labels
     for component in components:
         if "generictemplate" in component:
             if (
@@ -364,7 +366,7 @@ def load_components(
                 ]
                 labels = component["generictemplate"]["metadata"]["labels"]
                 labels.update({"kueue.x-k8s.io/queue-name": lq_name})
-                labels.update(user_labels)
+                labels.update(cluster_labels)
             component_list.append(component["generictemplate"])
 
     resources = "---\n" + "---\n".join(
@@ -405,7 +407,7 @@ def generate_appwrapper(
     write_to_file: bool,
     verify_tls: bool,
     local_queue: Optional[str],
-    user_labels,
+    labels,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -457,13 +459,11 @@ def generate_appwrapper(
         if mcad:
             write_user_appwrapper(user_yaml, outfile)
         else:
-            write_components(user_yaml, outfile, namespace, local_queue, user_labels)
+            write_components(user_yaml, outfile, namespace, local_queue, labels)
         return outfile
     else:
         if mcad:
             user_yaml = load_appwrapper(user_yaml, name)
         else:
-            user_yaml = load_components(
-                user_yaml, name, namespace, local_queue, user_labels
-            )
+            user_yaml = load_components(user_yaml, name, namespace, local_queue, labels)
         return user_yaml
diff --git a/tests/test-case-no-mcad.yamls b/tests/test-case-no-mcad.yamls
@@ -5,6 +5,8 @@ metadata:
   labels:
     controller-tools.k8s.io: '1.0'
     kueue.x-k8s.io/queue-name: local-queue-default
+    testlabel: test
+    testlabel2: test
   name: unit-test-cluster-ray
   namespace: ns
 spec:

diff --git a/tests/unit_test.py b/tests/unit_test.py
@@ -324,6 +324,7 @@ def test_cluster_creation_no_mcad(mocker):
     config.name = "unit-test-cluster-ray"
     config.write_to_file = True
     config.mcad = False
+    config.labels = {"testlabel": "test", "testlabel2": "test"}
     cluster = Cluster(config)
 
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"
@@ -348,6 +349,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
     config.mcad = False
     config.write_to_file = True
     config.local_queue = "local-queue-default"
+    config.labels = {"testlabel": "test", "testlabel2": "test"}
     cluster = Cluster(config)
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"
     assert cluster.app_wrapper_name == "unit-test-cluster-ray"
@@ -373,6 +375,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
         write_to_file=True,
         mcad=False,
         local_queue="local-queue-default",
+        labels={"testlabel": "test", "testlabel2": "test"},
     )
     cluster = Cluster(config)
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"