Refactor ray creation

src/codeflare_sdk/common/kueue/test_kueue.py

@@ -51,8 +51,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
     config.write_to_file = True
     config.local_queue = "local-queue-default"
     cluster = Cluster(config)
-    assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-kueue.yaml"
-    assert cluster.app_wrapper_name == "unit-test-cluster-kueue"
+    assert cluster.resource_yaml == f"{aw_dir}unit-test-cluster-kueue.yaml"
     assert filecmp.cmp(
         f"{aw_dir}unit-test-cluster-kueue.yaml",
         f"{parent}/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml",
@@ -65,10 +64,9 @@ def test_cluster_creation_no_aw_local_queue(mocker):
     config.write_to_file = False
     cluster = Cluster(config)
 
-    test_rc = yaml.load(cluster.app_wrapper_yaml, Loader=yaml.FullLoader)
     with open(f"{parent}/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml") as f:
         expected_rc = yaml.load(f, Loader=yaml.FullLoader)
-        assert test_rc == expected_rc
+        assert cluster.resource_yaml == expected_rc
 
 
 def test_aw_creation_local_queue(mocker):
@@ -87,8 +85,7 @@ def test_aw_creation_local_queue(mocker):
     config.write_to_file = True
     config.local_queue = "local-queue-default"
     cluster = Cluster(config)
-    assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-aw-kueue.yaml"
-    assert cluster.app_wrapper_name == "unit-test-aw-kueue"
+    assert cluster.resource_yaml == f"{aw_dir}unit-test-aw-kueue.yaml"
     assert filecmp.cmp(
         f"{aw_dir}unit-test-aw-kueue.yaml",
         f"{parent}/tests/test_cluster_yamls/kueue/aw_kueue.yaml",
@@ -102,10 +99,9 @@ def test_aw_creation_local_queue(mocker):
     config.write_to_file = False
     cluster = Cluster(config)
 
-    test_rc = yaml.load(cluster.app_wrapper_yaml, Loader=yaml.FullLoader)
     with open(f"{parent}/tests/test_cluster_yamls/kueue/aw_kueue.yaml") as f:
         expected_rc = yaml.load(f, Loader=yaml.FullLoader)
-        assert test_rc == expected_rc
+        assert cluster.resource_yaml == expected_rc
 
 
 def test_get_local_queue_exists_fail(mocker):

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -20,6 +20,7 @@
 import yaml
 from pathlib import Path
 from kubernetes import client
+from unittest.mock import patch
 
 parent = Path(__file__).resolve().parents[4]  # project directory
 aw_dir = os.path.expanduser("~/.codeflare/resources/")
@@ -381,3 +382,38 @@ def mocked_ingress(port, cluster_name="unit-test-cluster", annotations: dict = N
         ),
     )
     return mock_ingress
+
+
+@patch.dict("os.environ", {"NB_PREFIX": "test-prefix"})
+def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Cluster:
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
+
+    config = ClusterConfiguration(
+        name=cluster_name,
+        namespace="ns",
+        head_cpu_requests=4,
+        head_cpu_limits=8,
+        head_memory_requests=12,
+        head_memory_limits=16,
+        head_extended_resource_requests={"nvidia.com/gpu": 1, "intel.com/gpu": 2},
+        worker_cpu_requests=4,
+        worker_cpu_limits=8,
+        num_workers=10,
+        worker_memory_requests=12,
+        worker_memory_limits=16,
+        appwrapper=is_appwrapper,
+        envs={"key1": "value1", "key2": "value2"},
+        image="example/ray:tag",
+        image_pull_secrets=["secret1", "secret2"],
+        write_to_file=True,
+        verify_tls=True,
+        labels={"key1": "value1", "key2": "value2"},
+        worker_extended_resource_requests={"nvidia.com/gpu": 1},
+        extended_resource_mapping={"example.com/gpu": "GPU", "intel.com/gpu": "TPU"},
+        overwrite_default_resource_mapping=True,
+        local_queue="local-queue-default",
+    )
+    return Cluster(config)

src/codeflare_sdk/ray/appwrapper/test_status.py

@@ -19,6 +19,7 @@
 )
 from codeflare_sdk.ray.appwrapper import AppWrapper, AppWrapperStatus
 from codeflare_sdk.ray.cluster.status import CodeFlareClusterStatus
+from codeflare_sdk.common.utils.unit_test_support import get_local_queue
 import os
 
 aw_dir = os.path.expanduser("~/.codeflare/resources/")
@@ -28,8 +29,8 @@ def test_cluster_status(mocker):
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
     mocker.patch(
-        "codeflare_sdk.common.kueue.kueue.local_queue_exists",
-        return_value="true",
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
     fake_aw = AppWrapper("test", AppWrapperStatus.FAILED)
 
@@ -39,7 +40,7 @@ def test_cluster_status(mocker):
             namespace="ns",
             write_to_file=True,
             appwrapper=True,
-            local_queue="local_default_queue",
+            local_queue="local-queue-default",
         )
     )
     mocker.patch(