Adding unit tests for apply

Author: akram

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -27,10 +27,15 @@
 
 
 def createClusterConfig():
+    config = createClusterConfigWithNumWorkers()
+    return config
+
+
+def createClusterConfigWithNumWorkers(num_workers=2):
     config = ClusterConfiguration(
         name="unit-test-cluster",
         namespace="ns",
-        num_workers=2,
+        num_workers=num_workers,
         worker_cpu_requests=3,
         worker_cpu_limits=4,
         worker_memory_requests=5,
@@ -41,13 +46,18 @@ def createClusterConfig():
     return config
 
 
-def createClusterWithConfig(mocker):
+def createClusterWithConfigAndNumWorkers(mocker, num_workers=2):
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
     mocker.patch(
         "kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
         return_value={"spec": {"domain": "apps.cluster.awsroute.org"}},
     )
-    cluster = Cluster(createClusterConfig())
+    cluster = Cluster(createClusterConfigWithNumWorkers(num_workers))
+    return cluster
+
+
+def createClusterWithConfig(mocker):
+    cluster = createClusterWithConfigAndNumWorkers(mocker)
     return cluster
 
 

src/codeflare_sdk/ray/cluster/cluster.py

@@ -144,7 +144,7 @@ def up(self):
         the Kueue localqueue.
         """
         # TODO: Add deprecation message in favor of apply()
-        # print( "WARNING: The up() is planned for deprecation in favor of apply().")        
+        # print( "WARNING: The up() is planned for deprecation in favor of apply().")
 
         # check if RayCluster CustomResourceDefinition exists if not throw RuntimeError
         self._throw_for_no_raycluster()
@@ -182,7 +182,6 @@ def up(self):
         except Exception as e:  # pragma: no cover
             return _kube_api_error_handling(e)
 
-
     def apply(self, force=False):
         """
         Applies the Cluster yaml using server-side apply.
@@ -197,15 +196,14 @@ def apply(self, force=False):
         try:
             # Get the RayCluster custom resource definition
             api = dynamic_client.resources.get(
-                api_version="ray.io/v1",
-                kind="RayCluster"
+                api_version="ray.io/v1", kind="RayCluster"
             )
         except Exception as e:
             raise RuntimeError("Failed to get RayCluster resource: " + str(e))
 
         # Read the YAML file and parse it into a dictionary
         try:
-            with open(self.resource_yaml, 'r') as f:
+            with open(self.resource_yaml, "r") as f:
                 resource_body = yaml.safe_load(f)
         except FileNotFoundError:
             raise RuntimeError(f"Resource YAML file '{self.resource_yaml}' not found.")
@@ -224,7 +222,7 @@ def apply(self, force=False):
                 name=resource_name,
                 namespace=self.config.namespace,
                 field_manager="cluster-manager",
-                force_conflicts=force  # Allow forcing conflicts if needed
+                force_conflicts=force,  # Allow forcing conflicts if needed
             )
             print(f"Cluster '{self.config.name}' applied successfully.")
         except ApiException as e:
@@ -234,7 +232,9 @@ def apply(self, force=False):
                     "To force the patch, set 'force=True' in the apply() method."
                 )
             elif e.status == 404:
-                print(f"Namespace '{self.config.namespace}' or resource '{resource_name}' not found. Verify the namespace or CRD.")
+                print(
+                    f"Namespace '{self.config.namespace}' or resource '{resource_name}' not found. Verify the namespace or CRD."
+                )
             else:
                 raise RuntimeError(f"Failed to apply cluster: {e.reason}")
 

src/codeflare_sdk/ray/cluster/test_cluster.py

@@ -20,6 +20,7 @@
 )
 from codeflare_sdk.common.utils.unit_test_support import (
     createClusterWithConfig,
+    createClusterWithConfigAndNumWorkers,
     arg_check_del_effect,
     ingress_retrieval,
     arg_check_apply_effect,
@@ -67,11 +68,49 @@ def test_cluster_up_down(mocker):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    cluster = cluster = createClusterWithConfig(mocker)
+    cluster = createClusterWithConfig(mocker)
     cluster.up()
     cluster.down()
 
 
+def test_cluster_apply_scale_up_scale_down(mocker):
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster._throw_for_no_raycluster")
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
+        return_value={"spec": {"domain": ""}},
+    )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.create_namespaced_custom_object",
+        side_effect=arg_check_apply_effect,
+    )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.delete_namespaced_custom_object",
+        side_effect=arg_check_del_effect,
+    )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_cluster_custom_object",
+        return_value={"items": []},
+    )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
+    initial_num_workers = 1
+    scaled_up_num_workers = 2
+    cluster = createClusterWithConfigAndNumWorkers(mocker, initial_num_workers)
+    cluster.apply()
+    cluster.wait_ready(timeout=5)
+    cluster = createClusterWithConfigAndNumWorkers(mocker, scaled_up_num_workers)
+    cluster.apply()
+    cluster.wait_ready(timeout=5)
+    cluster = createClusterWithConfigAndNumWorkers(mocker, initial_num_workers)
+    cluster.apply()
+    cluster.wait_ready(timeout=5)
+    cluster.down()
+
+
 def test_cluster_up_down_no_mcad(mocker):
     mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster._throw_for_no_raycluster")
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
@@ -117,7 +156,7 @@ def test_cluster_uris(mocker):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    cluster = cluster = createClusterWithConfig(mocker)
+    cluster = createClusterWithConfig(mocker)
     mocker.patch(
         "kubernetes.client.NetworkingV1Api.list_namespaced_ingress",
         return_value=ingress_retrieval(
@@ -159,7 +198,7 @@ def ray_addr(self, *args):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    cluster = cluster = createClusterWithConfig(mocker)
+    cluster = createClusterWithConfig(mocker)
     mocker.patch(
         "ray.job_submission.JobSubmissionClient._check_connection_and_version_with_url",
         return_value="None",

tests/e2e/cluster_apply_test.py

@@ -82,7 +82,9 @@ def test_cluster_apply(self):
         # Wait for the updated cluster to be ready
         cluster.wait_ready()
         updated_status = cluster.status()
-        assert updated_status["ready"], f"Cluster {cluster_name} is not ready after update: {updated_status}"
+        assert updated_status[
+            "ready"
+        ], f"Cluster {cluster_name} is not ready after update: {updated_status}"
 
         # Verify the cluster is updated
         updated_ray_cluster = get_ray_cluster(cluster_name, namespace)
@@ -152,4 +154,3 @@ def test_apply_invalid_update(self):
 
         # Clean up
         cluster.down()
-

tests/e2e/support.py

@@ -10,6 +10,21 @@
     _kube_api_error_handling,
 )
 
+def get_ray_cluster(cluster_name, namespace):
+    api = client.CustomObjectsApi()
+    try:
+        return api.get_namespaced_custom_object(
+            group="ray.io",
+            version="v1",
+            namespace=namespace,
+            plural="rayclusters",
+            name=cluster_name,
+        )
+    except client.exceptions.ApiException as e:
+        if e.status == 404:
+            return None
+        raise
+
 
 def get_ray_image():
     default_ray_image = "quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06"