From 113938e71a6dacba37164069e2c795cc9db168d4 Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Tue, 22 Apr 2025 16:15:39 +0200
Subject: [PATCH] test: fix rpc widgets e2e test

---
 tests/end-2-end/test_rpc_widgets_e2e.py | 75 +++++++++++++++++++++++--
 1 file changed, 69 insertions(+), 6 deletions(-)

diff --git a/tests/end-2-end/test_rpc_widgets_e2e.py b/tests/end-2-end/test_rpc_widgets_e2e.py
index 712f29c0..4ed1fcb2 100644
--- a/tests/end-2-end/test_rpc_widgets_e2e.py
+++ b/tests/end-2-end/test_rpc_widgets_e2e.py
@@ -14,10 +14,21 @@ def wait_for_namespace_change(
     parent_widget: RPCBase | RPCReference,
     object_name: str,
     widget_gui_id: str,
-    timeout: int = 10000,
+    timeout: float = 10000,
     exists: bool = True,
 ):
-    """Utility method to wait for the namespace to be created in the widget."""
+    """
+    Utility method to wait for the namespace to be created in the widget.
+
+    Args:
+        qtbot: The qtbot fixture.
+        gui: The client_utils.BECGuiClient 'gui' object from the CLI.
+        parent_widget: The widget that creates a new widget.
+        object_name: The name of the widget that was created. Must appear as attribute in namespace of parent.
+        widget_gui_id: The gui_id of the created widget.
+        timeout: The timeout in milliseconds for the qtbot to wait for changes to appear.
+        exists: If True, wait for the object to be created. If False, wait for the object to be removed.
+    """
     # GUI object is not registered in the registry (yet)
     if parent_widget is gui:
 
@@ -67,6 +78,16 @@ def test_available_widgets(qtbot, connected_client_gui_obj):
     """This test checks that all widgets that are available via gui.available_widgets can be created and removed."""
     gui = connected_client_gui_obj
     dock_area = gui.bec
+    # Number of top level widgets, should be 4
+    top_level_widgets_count = 4
+    assert len(gui._server_registry) == top_level_widgets_count
+    # Number of widgets with parent_id == None, should be 2
+    widgets = [
+        widget for widget in gui._server_registry.values() if widget["config"]["parent_id"] is None
+    ]
+    assert len(widgets) == 2
+
+    # Test all relevant widgets
     for object_name in gui.available_widgets.__dict__:
         # Skip private attributes
         if object_name.startswith("_"):
@@ -74,15 +95,57 @@ def test_available_widgets(qtbot, connected_client_gui_obj):
         # Skip VSCode widget as Code server is not available in the Docker image
         if object_name == "VSCodeEditor":
             continue
-        # Create widget the widget
+
+        #############################
+        ######### Add widget ########
+        #############################
+
+        # Create widget the widget and wait for the widget to be registered in the ipython registry
         dock, widget = create_widget(
             qtbot, gui, dock_area, getattr(gui.available_widgets, object_name)
         )
-        # The create_widget method already waits for the widget to be created
-        # and added to the ipython registry. We can here assert if the dock_area
-        # has the dock and the widget
+        # Check that the widget is indeed registered on the server and the client
         assert gui._ipython_registry.get(widget._gui_id, None) is not None
+        assert gui._server_registry.get(widget._gui_id, None) is not None
+        # Check that namespace was updated
         assert hasattr(dock_area, dock.object_name)
         assert hasattr(dock, widget.object_name)
+
+        # Check that no additional top level widgets were created without a parent_id
+        widgets = [
+            widget
+            for widget in gui._server_registry.values()
+            if widget["config"]["parent_id"] is None
+        ]
+        assert len(widgets) == 2
+
+        #############################
+        ####### Remove widget #######
+        #############################
+
         # Now we remove the widget again
+        dock_name = dock.object_name
+        dock_id = dock._gui_id
+        widget_id = widget._gui_id
         dock_area.delete(dock.object_name)
+        # Wait for namespace to change
+        wait_for_namespace_change(qtbot, gui, dock_area, dock_name, dock_id, exists=False)
+        # Assert that dock and widget are removed from the ipython registry and the namespace
+        assert hasattr(dock_area, dock_name) is False
+        # Client registry
+        assert gui._ipython_registry.get(dock_id, None) is None
+        assert gui._ipython_registry.get(widget_id, None) is None
+        # Server registry
+        assert gui._server_registry.get(dock_id, None) is None
+        assert gui._server_registry.get(widget_id, None) is None
+
+        # Check that the number of top level widgets is still the same. As the cleanup is done by the
+        # qt event loop, we need to wait for the qtbot to finish the cleanup
+        qtbot.waitUntil(lambda: len(gui._server_registry) == top_level_widgets_count)
+        # Number of widgets with parent_id == None, should be 2
+        widgets = [
+            widget
+            for widget in gui._server_registry.values()
+            if widget["config"]["parent_id"] is None
+        ]
+        assert len(widgets) == 2