test: update unit tests after refactor

Author: Bobbins228

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(

src/codeflare_sdk/ray/cluster/test_generate_yaml.py src/codeflare_sdk/ray/cluster/test_build_ray_cluster.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from collections import namedtuple
 import sys
-from .generate_yaml import gen_names, update_image
+from .build_ray_cluster import gen_names, update_image
 import uuid
 
 
@@ -39,7 +39,7 @@ def test_gen_names_without_name(mocker):
 def test_update_image_without_supported_python_version(mocker):
     # Mock SUPPORTED_PYTHON_VERSIONS
     mocker.patch.dict(
-        "codeflare_sdk.ray.cluster.generate_yaml.SUPPORTED_PYTHON_VERSIONS",
+        "codeflare_sdk.ray.cluster.build_ray_cluster.SUPPORTED_PYTHON_VERSIONS",
         {
             "3.9": "ray-py3.9",
             "3.11": "ray-py3.11",
@@ -55,16 +55,13 @@ def test_update_image_without_supported_python_version(mocker):
     # Mock warnings.warn to check if it gets called
     warn_mock = mocker.patch("warnings.warn")
 
-    # Create a sample spec
-    spec = {"containers": [{"image": None}]}
-
     # Call the update_image function with no image provided
-    update_image(spec, None)
+    image = update_image(None)
 
     # Assert that the warning was called with the expected message
     warn_mock.assert_called_once_with(
         "No default Ray image defined for 3.8. Please provide your own image or use one of the following python versions: 3.9, 3.11."
     )
 
-    # Assert that no image was set in the containers since the Python version is not supported
-    assert spec["containers"][0]["image"] is None
+    # Assert that no image was set since the Python version is not supported
+    assert image is None

src/codeflare_sdk/ray/cluster/test_cluster.py

@@ -25,21 +25,17 @@
     arg_check_apply_effect,
     get_local_queue,
     createClusterConfig,
-    route_list_retrieval,
     get_ray_obj,
-    get_aw_obj,
-    get_named_aw,
     get_obj_none,
     get_ray_obj_with_status,
     get_aw_obj_with_status,
 )
-from codeflare_sdk.ray.cluster.generate_yaml import (
-    is_openshift_cluster,
-    is_kind_cluster,
-)
+from codeflare_sdk.ray.cluster.cluster import _is_openshift_cluster
 from pathlib import Path
 from unittest.mock import MagicMock
 from kubernetes import client
+import yaml
+import filecmp
 import os
 
 parent = Path(__file__).resolve().parents[4]  # project directory
@@ -205,7 +201,7 @@ def test_local_client_url(mocker):
         return_value="rayclient-unit-test-cluster-localinter-ns.apps.cluster.awsroute.org",
     )
     mocker.patch(
-        "codeflare_sdk.ray.cluster.cluster.Cluster.create_app_wrapper",
+        "codeflare_sdk.ray.cluster.cluster.Cluster.create_resource",
         return_value="unit-test-cluster-localinter.yaml",
     )
 
@@ -225,120 +221,51 @@ def test_local_client_url(mocker):
 """
 
 
-def test_get_cluster_openshift(mocker):
+def test_get_cluster_no_appwrapper(mocker):
+    """
+    This test uses the "test all params" unit test file as a comparison
+    """
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
-    # Mock the client.ApisApi function to return a mock object
-    mock_api = MagicMock()
-    mock_api.get_api_versions.return_value.groups = [
-        MagicMock(versions=[MagicMock(group_version="route.openshift.io/v1")])
-    ]
-    mocker.patch("kubernetes.client.ApisApi", return_value=mock_api)
     mocker.patch(
-        "codeflare_sdk.common.kueue.kueue.local_queue_exists",
-        return_value="true",
+        "codeflare_sdk.ray.cluster.cluster._check_aw_exists",
+        return_value=False,
     )
 
-    assert is_openshift_cluster()
-
-    def custom_side_effect(group, version, namespace, plural, **kwargs):
-        if plural == "routes":
-            return route_list_retrieval("route.openshift.io", "v1", "ns", "routes")
-        elif plural == "rayclusters":
-            return get_ray_obj("ray.io", "v1", "ns", "rayclusters")
-        elif plural == "appwrappers":
-            return get_aw_obj("workload.codeflare.dev", "v1beta2", "ns", "appwrappers")
-        elif plural == "localqueues":
-            return get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues")
-
-    mocker.patch(
-        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", get_aw_obj
-    )
-
-    mocker.patch(
-        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
-        side_effect=custom_side_effect,
-    )
-    mocker.patch(
-        "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
-        return_value=get_named_aw,
-    )
-    mocker.patch(
-        "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
-        side_effect=route_list_retrieval("route.openshift.io", "v1", "ns", "routes")[
-            "items"
-        ],
-    )
-    mocker.patch(
-        "codeflare_sdk.common.kueue.kueue.local_queue_exists",
-        return_value="true",
-    )
-
-    cluster = get_cluster(
-        "test-cluster-a", "ns"
-    )  # see tests/test_cluster_yamls/support_clusters
-    cluster_config = cluster.config
-
-    assert cluster_config.name == "test-cluster-a" and cluster_config.namespace == "ns"
-    assert cluster_config.head_cpu_requests == 2 and cluster_config.head_cpu_limits == 2
-    assert (
-        cluster_config.head_memory_requests == "8G"
-        and cluster_config.head_memory_limits == "8G"
-    )
-    assert (
-        cluster_config.worker_cpu_requests == 1
-        and cluster_config.worker_cpu_limits == 1
-    )
-    assert (
-        cluster_config.worker_memory_requests == "2G"
-        and cluster_config.worker_memory_limits == "2G"
-    )
-    assert cluster_config.num_workers == 1
-    assert cluster_config.write_to_file == False
-    assert cluster_config.local_queue == "local_default_queue"
+    with open(f"{expected_clusters_dir}/ray/unit-test-all-params.yaml") as f:
+        expected_rc = yaml.load(f, Loader=yaml.FullLoader)
+        mocker.patch(
+            "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
+            return_value=expected_rc,
+        )
+        get_cluster("test-all-params", "ns", write_to_file=True)
+        assert filecmp.cmp(
+            f"{aw_dir}test-all-params.yaml",
+            f"{expected_clusters_dir}/ray/unit-test-all-params.yaml",
+            shallow=True,
+        )
 
 
-def test_get_cluster(mocker):
-    # test get_cluster for Kind Clusters
+def test_get_cluster_with_appwrapper(mocker):
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
     mocker.patch(
-        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
-        side_effect=get_ray_obj,
-    )
-    mocker.patch(
-        "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
-        side_effect=get_named_aw,
-    )
-    mocker.patch(
-        "kubernetes.client.NetworkingV1Api.list_namespaced_ingress",
-        return_value=ingress_retrieval(cluster_name="quicktest", client_ing=True),
-    )
-    mocker.patch(
-        "codeflare_sdk.common.kueue.kueue.local_queue_exists",
-        return_value="true",
+        "codeflare_sdk.ray.cluster.cluster._check_aw_exists",
+        return_value=True,
     )
-    cluster = get_cluster(
-        "test-cluster-a"
-    )  # see tests/test_cluster_yamls/support_clusters
-    cluster_config = cluster.config
 
-    assert cluster_config.name == "test-cluster-a" and cluster_config.namespace == "ns"
-    assert cluster_config.head_cpu_requests == 2 and cluster_config.head_cpu_limits == 2
-    assert (
-        cluster_config.head_memory_requests == "8G"
-        and cluster_config.head_memory_limits == "8G"
-    )
-    assert (
-        cluster_config.worker_cpu_requests == 1
-        and cluster_config.worker_cpu_limits == 1
-    )
-    assert (
-        cluster_config.worker_memory_requests == "2G"
-        and cluster_config.worker_memory_limits == "2G"
-    )
-    assert cluster_config.num_workers == 1
-    assert cluster_config.write_to_file == False
-    assert cluster_config.local_queue == "local_default_queue"
+    with open(f"{expected_clusters_dir}/appwrapper/unit-test-all-params.yaml") as f:
+        expected_aw = yaml.load(f, Loader=yaml.FullLoader)
+        mocker.patch(
+            "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
+            return_value=expected_aw,
+        )
+        get_cluster("aw-all-params", "ns", write_to_file=True)
+        assert filecmp.cmp(
+            f"{aw_dir}aw-all-params.yaml",
+            f"{expected_clusters_dir}/appwrapper/unit-test-all-params.yaml",
+            shallow=True,
+        )
 
 
 def test_wait_ready(mocker, capsys):
@@ -356,10 +283,6 @@ def test_wait_ready(mocker, capsys):
     mocker.patch(
         "codeflare_sdk.ray.cluster.cluster._ray_cluster_status", return_value=None
     )
-    mocker.patch(
-        "codeflare_sdk.common.kueue.kueue.local_queue_exists",
-        return_value="true",
-    )
     mocker.patch.object(
         client.CustomObjectsApi,
         "list_namespaced_custom_object",
@@ -381,7 +304,6 @@ def test_wait_ready(mocker, capsys):
             namespace="ns",
             write_to_file=False,
             appwrapper=True,
-            local_queue="local-queue-default",
         )
     )
     try:
@@ -460,7 +382,7 @@ def test_list_queue_rayclusters(mocker, capsys):
     ]
     mocker.patch("kubernetes.client.ApisApi", return_value=mock_api)
 
-    assert is_openshift_cluster() == True
+    assert _is_openshift_cluster() == True
     mocker.patch(
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_obj_none("ray.io", "v1", "ns", "rayclusters"),
@@ -566,7 +488,7 @@ def test_map_to_ray_cluster(mocker):
     mocker.patch("kubernetes.config.load_kube_config")
 
     mocker.patch(
-        "codeflare_sdk.ray.cluster.cluster.is_openshift_cluster", return_value=True
+        "codeflare_sdk.ray.cluster.cluster._is_openshift_cluster", return_value=True
     )
 
     mock_api_client = mocker.MagicMock(spec=client.ApiClient)
@@ -608,3 +530,9 @@ def custom_side_effect(group, version, namespace, plural, **kwargs):
 
     assert result is not None
     assert result.dashboard == rc_dashboard
+
+
+# Make sure to always keep this function last
+def test_cleanup():
+    os.remove(f"{aw_dir}test-all-params.yaml")
+    os.remove(f"{aw_dir}aw-all-params.yaml")