Updated tests and load_components

Author: Bobbins228

src/codeflare_sdk/utils/generate_yaml.py

@@ -678,11 +678,23 @@ def write_components(
     print(f"Written to: {output_file_name}")
 
 
-def load_components(user_yaml: dict, name: str):
+def load_components(
+    user_yaml: dict, name: str, namespace: str, local_queue: Optional[str]
+):
     component_list = []
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
+    lq_name = local_queue or get_default_kueue_name(namespace)
     for component in components:
         if "generictemplate" in component:
+            if (
+                "workload.codeflare.dev/appwrapper"
+                in component["generictemplate"]["metadata"]["labels"]
+            ):
+                del component["generictemplate"]["metadata"]["labels"][
+                    "workload.codeflare.dev/appwrapper"
+                ]
+                labels = component["generictemplate"]["metadata"]["labels"]
+                labels.update({"kueue.x-k8s.io/queue-name": lq_name})
             component_list.append(component["generictemplate"])
 
     resources = "---\n" + "---\n".join(
@@ -790,11 +802,11 @@ def generate_appwrapper(
         if mcad:
             write_user_appwrapper(user_yaml, outfile)
         else:
-            write_components(user_yaml, outfile, local_queue)
+            write_components(user_yaml, outfile, namespace, local_queue)
         return outfile
     else:
         if mcad:
             user_yaml = load_appwrapper(user_yaml, name)
         else:
-            user_yaml = load_components(user_yaml, name)
+            user_yaml = load_components(user_yaml, name, namespace, local_queue)
         return user_yaml

tests/unit_test.py

@@ -333,7 +333,6 @@ def get_local_queue(group, version, namespace, plural):
 
 
 def test_cluster_creation_no_mcad(mocker):
-    # With written resources
     # Create Ray Cluster with no local queue specified
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch(
@@ -359,6 +358,7 @@ def test_cluster_creation_no_mcad(mocker):
 
 
 def test_cluster_creation_no_mcad_local_queue(mocker):
+    # With written resources
     # Create Ray Cluster with local queue specified
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch(
@@ -395,6 +395,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
         image="quay.io/project-codeflare/ray:latest-py39-cu118",
         write_to_file=False,
         mcad=False,
+        local_queue="local-queue-default",
     )
     cluster = Cluster(config)
     test_resources = []