From fd711b475f268fbdb59739da0a428f0355b25bac Mon Sep 17 00:00:00 2001
From: wyzula-jan <133381102+wyzula-jan@users.noreply.github.com>
Date: Mon, 26 Feb 2024 14:06:36 +0100
Subject: [PATCH] fix(cli/rpc): rpc client can return any type of object +
 config dict of the widgets

---
 bec_widgets/cli/client_utils.py | 29 ++++++++------
 bec_widgets/cli/server.py       |  1 +
 tests/conftest.py               | 70 ++++++++++++++++-----------------
 3 files changed, 53 insertions(+), 47 deletions(-)

diff --git a/bec_widgets/cli/client_utils.py b/bec_widgets/cli/client_utils.py
index 4422d19f..4f601d1a 100644
--- a/bec_widgets/cli/client_utils.py
+++ b/bec_widgets/cli/client_utils.py
@@ -88,12 +88,12 @@ class BECFigureClientMixin:
 
 
 class RPCBase:
-    def __init__(self, gui_id: str = None, config: dict = None, parent=None, **kwargs) -> None:
+    def __init__(self, gui_id: str = None, config: dict = None, parent=None) -> None:
         self._client = BECDispatcher().client
         self._config = config if config is not None else {}
         self._gui_id = gui_id if gui_id is not None else str(uuid.uuid4())
         self._parent = parent
-        super().__init__(**kwargs)
+        super().__init__()
         print(f"RPCBase: {self._gui_id}")
 
     @property
@@ -139,20 +139,25 @@ class RPCBase:
         if not response.content["accepted"]:
             raise ValueError(response.content["message"]["error"])
         msg_result = response.content["message"].get("result")
-        if not msg_result:
-            return None
-        if isinstance(msg_result, list):
-            return [self._create_widget_from_msg_result(res) for res in msg_result]
         return self._create_widget_from_msg_result(msg_result)
 
     def _create_widget_from_msg_result(self, msg_result):
-        cls = msg_result.pop("widget_class", None)
-        if not cls:
-            return msg_result
+        if msg_result is None:
+            return None
+        if isinstance(msg_result, list):
+            return [self._create_widget_from_msg_result(res) for res in msg_result]
+        if isinstance(msg_result, dict):
+            if "__rpc__" not in msg_result:
+                return msg_result
+            cls = msg_result.pop("widget_class", None)
 
-        cls = getattr(client, cls)
-        print(msg_result)
-        return cls(parent=self, **msg_result)
+            if not cls:
+                return msg_result
+
+            cls = getattr(client, cls)
+            print(msg_result)
+            return cls(parent=self, **msg_result)
+        return msg_result
 
     def _wait_for_response(self, request_id):
         """
diff --git a/bec_widgets/cli/server.py b/bec_widgets/cli/server.py
index 8a0f624c..5081bee6 100644
--- a/bec_widgets/cli/server.py
+++ b/bec_widgets/cli/server.py
@@ -95,6 +95,7 @@ class BECWidgetsCLIServer:
                 "gui_id": obj.gui_id,
                 "widget_class": obj.__class__.__name__,
                 "config": obj.config.model_dump(),
+                "__rpc__": True,
             }
         return obj
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 009d97fb..b732950a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,35 +1,35 @@
-import pytest
-import threading
-
-from bec_lib.bec_service import BECService
-from bec_widgets.utils import bec_dispatcher as bec_dispatcher_module
-
-
-@pytest.fixture()
-def threads_check():
-    current_threads = set(
-        th
-        for th in threading.enumerate()
-        if "loguru" not in th.name and th is not threading.main_thread()
-    )
-    yield
-    threads_after = set(
-        th
-        for th in threading.enumerate()
-        if "loguru" not in th.name and th is not threading.main_thread()
-    )
-    additional_threads = threads_after - current_threads
-    assert (
-        len(additional_threads) == 0
-    ), f"Test creates {len(additional_threads)} threads that are not cleaned: {additional_threads}"
-
-
-@pytest.fixture(autouse=True)
-def bec_dispatcher(threads_check):
-    bec_dispatcher = bec_dispatcher_module.BECDispatcher()
-    yield bec_dispatcher
-    bec_dispatcher.disconnect_all()
-    # clean BEC client
-    BECService.shutdown(bec_dispatcher.client)
-    # reinitialize singleton for next test
-    bec_dispatcher_module._bec_dispatcher = None
+# import pytest
+# import threading
+#
+# from bec_lib.bec_service import BECService
+# from bec_widgets.utils import bec_dispatcher as bec_dispatcher_module
+#
+#
+# @pytest.fixture()
+# def threads_check():
+#     current_threads = set(
+#         th
+#         for th in threading.enumerate()
+#         if "loguru" not in th.name and th is not threading.main_thread()
+#     )
+#     yield
+#     threads_after = set(
+#         th
+#         for th in threading.enumerate()
+#         if "loguru" not in th.name and th is not threading.main_thread()
+#     )
+#     additional_threads = threads_after - current_threads
+#     assert (
+#         len(additional_threads) == 0
+#     ), f"Test creates {len(additional_threads)} threads that are not cleaned: {additional_threads}"
+#
+#
+# @pytest.fixture(autouse=True)
+# def bec_dispatcher(threads_check):
+#     bec_dispatcher = bec_dispatcher_module.BECDispatcher()
+#     yield bec_dispatcher
+#     bec_dispatcher.disconnect_all()
+#     # clean BEC client
+#     BECService.shutdown(bec_dispatcher.client)
+#     # reinitialize singleton for next test
+#     bec_dispatcher_module._bec_dispatcher = None