test: test ui up and down buttons

Signed-off-by: Bobbins228 <[email protected]>

Author: Bobbins228

tests/unit_test.py

@@ -82,11 +82,14 @@
 import ray
 import pytest
 import yaml
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 from pytest_mock import MockerFixture
 from ray.job_submission import JobSubmissionClient
 from codeflare_sdk.job.ray_jobs import RayJobClient
 
+import ipywidgets as widgets
+from IPython.display import display
+
 # For mocking openshift client results
 fake_res = openshift.Result("fake")
 
@@ -2873,6 +2876,50 @@ def test_cluster_config_deprecation_conversion(mocker):
     assert config.worker_cpu_limits == 2
 
 
+"""
+    Ipywidgets tests
+"""
+
+
+def test_cluster_up_down_buttons(mocker):
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
+        return_value={"spec": {"domain": "apps.cluster.awsroute.org"}},
+    )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
+    os.environ["JPY_SESSION_NAME"] = "example-test"
+    cluster = Cluster(createClusterConfig())
+
+    with patch("ipywidgets.Button") as MockButton, patch("ipywidgets.Output"), patch(
+        "ipywidgets.HBox"
+    ), patch.object(cluster, "up") as mock_up, patch.object(
+        cluster, "down"
+    ) as mock_down:
+        # Create mock button instances
+        mock_up_button = MagicMock()
+        mock_down_button = MagicMock()
+
+        # Ensure the mock Button class returns the mock button instances in sequence
+        MockButton.side_effect = [mock_down_button, mock_up_button]
+
+        # Call the method under test
+        cluster._cluster_up_down_buttons()
+
+        # Simulate the button clicks by calling the mock on_click handlers
+        mock_up_button.on_click.call_args[0][0](None)  # Simulate clicking "Cluster Up"
+        mock_down_button.on_click.call_args[0][0](
+            None
+        )  # Simulate clicking "Cluster Down"
+
+        # Check if the `up` and `down` methods were called
+        mock_up.assert_called_once()
+        mock_down.assert_called_once()
+
+
 # Make sure to always keep this function last
 def test_cleanup():
     os.remove(f"{aw_dir}unit-test-no-kueue.yaml")