Pre-select cluster if exists, and suppress widgets and outputs on creation of Cluster Object, and bug fixes

Author: ChristianZaccaria

src/codeflare_sdk/cluster/widgets.py

@@ -15,6 +15,8 @@
 """
 The widgets sub-module contains the ui widgets created using the ipywidgets package.
 """
+import contextlib
+import io
 import os
 from time import sleep
 import time
@@ -113,38 +115,42 @@ def view_clusters(namespace: str = None):
     df = _fetch_cluster_data(namespace)
     if df.empty:
         print(f"No clusters found in the {namespace} namespace.")
-    else:
-        classification_widget = widgets.ToggleButtons(
-            options=df["Name"].tolist(), value=None,
-            description='Select an existing cluster:',
-        )
+        return
 
-        classification_widget.observe(lambda selection_change: _on_cluster_click(selection_change, raycluster_data_output, namespace, classification_widget), names="value")
+    classification_widget = widgets.ToggleButtons(
+        options=df["Name"].tolist(), value=df["Name"].tolist()[0],
+        description='Select an existing cluster:',
+    )
+    # Setting the initial value to trigger the event handler to display the cluster details.
+    initial_value = classification_widget.value
+    _on_cluster_click({"new": initial_value}, raycluster_data_output, namespace, classification_widget)
+    classification_widget.observe(lambda selection_change: _on_cluster_click(selection_change, raycluster_data_output, namespace, classification_widget), names="value")
 
-        delete_button = widgets.Button(
-                    description='Delete Cluster',
-                    icon='trash',
-                    tooltip="Delete the selected cluster"
-                )
-        delete_button.on_click(lambda b: _on_delete_button_click(b, classification_widget, df, raycluster_data_output, user_output, delete_button, list_jobs_button, ray_dashboard_button))
+    # UI table buttons
+    delete_button = widgets.Button(
+                description='Delete Cluster',
+                icon='trash',
+                tooltip="Delete the selected cluster"
+            )
+    delete_button.on_click(lambda b: _on_delete_button_click(b, classification_widget, df, raycluster_data_output, user_output, delete_button, list_jobs_button, ray_dashboard_button))
 
-        list_jobs_button = widgets.Button(
-                    description='View Jobs',
-                    icon='suitcase',
-                    tooltip="Open the Ray Job Dashboard"
-                )
-        list_jobs_button.on_click(lambda b: _on_list_jobs_button_click(b, classification_widget, df, user_output, url_output))
+    list_jobs_button = widgets.Button(
+                description='View Jobs',
+                icon='suitcase',
+                tooltip="Open the Ray Job Dashboard"
+            )
+    list_jobs_button.on_click(lambda b: _on_list_jobs_button_click(b, classification_widget, df, user_output, url_output))
 
-        ray_dashboard_button = widgets.Button(
-                    description='Open Ray Dashboard',
-                    icon='dashboard',
-                    tooltip="Open the Ray Dashboard in a new tab",
-                    layout=widgets.Layout(width='auto'),
-                )
-        ray_dashboard_button.on_click(lambda b: _on_ray_dashboard_button_click(b, classification_widget, df, user_output, url_output))
+    ray_dashboard_button = widgets.Button(
+                description='Open Ray Dashboard',
+                icon='dashboard',
+                tooltip="Open the Ray Dashboard in a new tab",
+                layout=widgets.Layout(width='auto'),
+            )
+    ray_dashboard_button.on_click(lambda b: _on_ray_dashboard_button_click(b, classification_widget, df, user_output, url_output))
 
-        display(widgets.VBox([classification_widget, raycluster_data_output]))
-        display(widgets.HBox([delete_button, list_jobs_button, ray_dashboard_button]), url_output, user_output)
+    display(widgets.VBox([classification_widget, raycluster_data_output]))
+    display(widgets.HBox([delete_button, list_jobs_button, ray_dashboard_button]), url_output, user_output)
 
 # Handles the event when a cluster is selected from the toggle buttons, updating the output with cluster details.
 def _on_cluster_click(selection_change, raycluster_data_output: widgets.Output, namespace: str, classification_widget: widgets.ToggleButtons):
@@ -184,7 +190,9 @@ def _on_ray_dashboard_button_click(b, classification_widget: widgets.ToggleButto
     cluster_name = classification_widget.value
     namespace = df[df["Name"]==classification_widget.value]["Namespace"].values[0]
 
-    cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
+    # Suppress from Cluster Object initialisation widgets and outputs
+    with widgets.Output(), contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+        cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
     dashboard_url = cluster.cluster_dashboard_uri()
 
     with user_output:
@@ -198,7 +206,9 @@ def _on_list_jobs_button_click(b, classification_widget: widgets.ToggleButtons,
     cluster_name = classification_widget.value
     namespace = df[df["Name"]==classification_widget.value]["Namespace"].values[0]
 
-    cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
+    # Suppress from Cluster Object initialisation widgets and outputs
+    with widgets.Output(), contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+        cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
     dashboard_url = cluster.cluster_dashboard_uri()
 
     with user_output:

tests/unit_test.py

@@ -2953,7 +2953,6 @@ def test_cluster_up_down_buttons(mocker):
 @patch.dict("os.environ", {}, clear=True)  # Mock environment with no variables
 def test_is_notebook_false():
     from codeflare_sdk.cluster.widgets import is_notebook
-
     assert is_notebook() is False
 
 
@@ -2962,9 +2961,9 @@ def test_is_notebook_false():
 )  # Mock Jupyter environment variable
 def test_is_notebook_true():
     from codeflare_sdk.cluster.widgets import is_notebook
-
     assert is_notebook() is True
 
+
 @patch.dict("os.environ", {"JPY_SESSION_NAME": "example-test"})  # Mock Jupyter environment variable
 def test_view_clusters(mocker):
     # Mock Kubernetes API responses
@@ -2974,6 +2973,10 @@ def test_view_clusters(mocker):
         return_value={"items": []}
     )
 
+    # Return empty dataframe when no clusters are found
+    mocker.patch("codeflare_sdk.cluster.cluster.list_all_clusters", return_value=[])
+    df = _fetch_cluster_data(namespace="default")
+    assert df.empty
 
     test_df=pd.DataFrame({
         "Name": ["test-cluster"],
@@ -3029,10 +3032,6 @@ def test_view_clusters(mocker):
         # Assert that the toggle options are set correctly
         mock_toggle.observe.assert_called()
 
-        # Simulate clicking the delete button
-        _on_delete_button_click(None, mock_toggle, test_df, mock_output, mock_output,
-                                mock_delete_button, mock_list_jobs_button, mock_ray_dashboard_button)
-
         # Simulate clicking the list jobs button
         _on_list_jobs_button_click(None, mock_toggle, test_df, mock_output, mock_output)
         mock_javascript.assert_called()
@@ -3041,26 +3040,17 @@ def test_view_clusters(mocker):
         _on_ray_dashboard_button_click(None, mock_toggle, test_df, mock_output, mock_output)
         mock_javascript.assert_called()
 
+        # Simulate clicking the delete button
+        _on_delete_button_click(None, mock_toggle, test_df, mock_output, mock_output,
+                                mock_delete_button, mock_list_jobs_button, mock_ray_dashboard_button)
 
 
-def test_format_status():
-    # Test each possible status
-    test_cases = [
-        (RayClusterStatus.READY, '<span style="color: green;">Ready ✓</span>'),
-        (RayClusterStatus.SUSPENDED, '<span style="color: #007BFF;">Suspended ❄️</span>'),
-        (RayClusterStatus.FAILED, '<span style="color: red;">Failed ✗</span>'),
-        (RayClusterStatus.UNHEALTHY, '<span style="color: purple;">Unhealthy</span>'),
-        (RayClusterStatus.UNKNOWN, '<span style="color: purple;">Unknown</span>'),
-    ]
-
-    for status, expected_output in test_cases:
-        assert _format_status(status) == expected_output, f"Failed for status: {status}"
-
-    # Test an unrecognized status
-    unrecognized_status = 'NotAStatus'
-    assert _format_status(unrecognized_status) == 'NotAStatus', "Failed for unrecognized status"
-
 def test_fetch_cluster_data(mocker):
+    # Return empty dataframe when no clusters are found
+    mocker.patch("codeflare_sdk.cluster.cluster.list_all_clusters", return_value=[])
+    df = _fetch_cluster_data(namespace="default")
+    assert df.empty
+
     # Create mock RayCluster objects
     mock_raycluster1 = MagicMock(spec=RayCluster)
     mock_raycluster1.name = 'test-cluster-1'
@@ -3121,6 +3111,25 @@ def test_fetch_cluster_data(mocker):
     # Assert that the DataFrame matches expected
     pd.testing.assert_frame_equal(df.reset_index(drop=True), expected_df.reset_index(drop=True))
 
+
+def test_format_status():
+    # Test each possible status
+    test_cases = [
+        (RayClusterStatus.READY, '<span style="color: green;">Ready ✓</span>'),
+        (RayClusterStatus.SUSPENDED, '<span style="color: #007BFF;">Suspended ❄️</span>'),
+        (RayClusterStatus.FAILED, '<span style="color: red;">Failed ✗</span>'),
+        (RayClusterStatus.UNHEALTHY, '<span style="color: purple;">Unhealthy</span>'),
+        (RayClusterStatus.UNKNOWN, '<span style="color: purple;">Unknown</span>'),
+    ]
+
+    for status, expected_output in test_cases:
+        assert _format_status(status) == expected_output, f"Failed for status: {status}"
+
+    # Test an unrecognized status
+    unrecognized_status = 'NotAStatus'
+    assert _format_status(unrecognized_status) == 'NotAStatus', "Failed for unrecognized status"
+
+
 # Make sure to always keep this function last
 def test_cleanup():
     os.remove(f"{aw_dir}unit-test-no-kueue.yaml")