From ad664006ec507a9785cdf58adc749f5055464676 Mon Sep 17 00:00:00 2001
From: Benedikt Burger <67148916+BenediktBurger@users.noreply.github.com>
Date: Mon, 27 May 2024 18:15:11 +0200
Subject: [PATCH] Update Director (and Fake) to changes.

---
 pyleco/directors/director.py     | 39 ++++++++++++++++++++++++++------
 pyleco/test.py                   | 20 ++++++++++++----
 tests/directors/test_director.py | 16 +++++++++++++
 3 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/pyleco/directors/director.py b/pyleco/directors/director.py
index 174a82d9..8e6f9fc1 100644
--- a/pyleco/directors/director.py
+++ b/pyleco/directors/director.py
@@ -177,23 +177,48 @@ def call_action(self, action: str, *args, actor: Optional[Union[bytes, str]] = N
         return self.ask_rpc("call_action", action=action, actor=actor, **params)
 
     # Async methods: Just send, read later.
-    def send(self, actor: Optional[Union[bytes, str]] = None, data=None, **kwargs) -> bytes:
+    def send(
+        self,
+        actor: Optional[Union[bytes, str]] = None,
+        data=None,
+        additional_payload: Optional[Iterable[bytes]] = None,
+        **kwargs,
+    ) -> bytes:
         """Send a request and return the conversation_id."""
         actor = self._actor_check(actor)
         cid0 = generate_conversation_id()
-        self.communicator.send(actor, conversation_id=cid0, data=data, **kwargs)
+        self.communicator.send(
+            actor, conversation_id=cid0, data=data, additional_payload=additional_payload, **kwargs
+        )
         return cid0
 
-    def ask_rpc_async(self, method: str, actor: Optional[Union[bytes, str]] = None,
-                      **kwargs) -> bytes:
+    def ask_rpc_async(
+        self,
+        method: str,
+        actor: Optional[Union[bytes, str]] = None,
+        additional_payload: Optional[Iterable[bytes]] = None,
+        **kwargs,
+    ) -> bytes:
         """Send a rpc request, the response can be read later with :meth:`read_rpc_response`."""
         string = self.generator.build_request_str(method=method, **kwargs)
-        return self.send(actor=actor, data=string, message_type=MessageTypes.JSON)
+        return self.send(
+            actor=actor,
+            data=string,
+            message_type=MessageTypes.JSON,
+            additional_payload=additional_payload,
+        )
 
-    def read_rpc_response(self, conversation_id: Optional[bytes] = None, **kwargs) -> Any:
+    def read_rpc_response(
+        self,
+        conversation_id: Optional[bytes] = None,
+        extract_additional_payload: bool = False,
+        **kwargs,
+    ) -> Any:
         """Read the response value corresponding to a request with a certain `conversation_id`."""
         response_message = self.communicator.read_message(conversation_id=conversation_id, **kwargs)
-        return self.communicator.interpret_rpc_response(response_message=response_message)
+        return self.communicator.interpret_rpc_response(
+            response_message=response_message, extract_additional_payload=extract_additional_payload
+        )
 
     #   Actor
     def get_parameters_async(self, parameters: Union[str, Sequence[str]],
diff --git a/pyleco/test.py b/pyleco/test.py
index 8d1da5fd..edbce43c 100644
--- a/pyleco/test.py
+++ b/pyleco/test.py
@@ -23,7 +23,7 @@
 #
 
 from __future__ import annotations
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Iterable, Optional, Sequence, Union
 
 from .core.message import Message
 from .core.internal_protocols import CommunicatorProtocol
@@ -219,14 +219,26 @@ def __init__(self, remote_class, **kwargs):
         super().__init__(**kwargs)
         self.remote_class = remote_class
 
-    def ask_rpc(self, method: str, actor: Optional[Union[bytes, str]] = None, **kwargs) -> Any:
+    def ask_rpc(
+        self,
+        method: str,
+        actor: Optional[Union[bytes, str]] = None,
+        additional_payload: Optional[Iterable[bytes]] = None,
+        extract_additional_payload: bool = False,
+        **kwargs,
+    ) -> Any:
         assert hasattr(self.remote_class, method), f"Remote class does not have method '{method}'."
         self.method = method
         self.kwargs = kwargs
         return self.return_value
 
-    def ask_rpc_async(self, method: str, actor: Optional[Union[bytes, str]] = None,
-                      **kwargs) -> bytes:
+    def ask_rpc_async(
+        self,
+        method: str,
+        actor: Optional[Union[bytes, str]] = None,
+        additional_payload: Optional[Iterable[bytes]] = None,
+        **kwargs,
+    ) -> bytes:
         assert hasattr(self.remote_class, method), f"Remote class does not have method '{method}'."
         self.method = method
         self.kwargs = kwargs
diff --git a/tests/directors/test_director.py b/tests/directors/test_director.py
index 8e7a012a..3ed5bea8 100644
--- a/tests/directors/test_director.py
+++ b/tests/directors/test_director.py
@@ -136,6 +136,22 @@ def test_read_rpc_response(director: Director):
     assert director.read_rpc_response(conversation_id=cid) == 7.5
 
 
+def test_read_binary_rpc_response(director: Director):
+    director.communicator._r = [  # type: ignore
+        Message(
+            "director",
+            "actor",
+            conversation_id=cid,
+            message_type=MessageTypes.JSON,
+            data={"id": 1, "result": None, "jsonrpc": "2.0"},
+            additional_payload=[b"123"],
+        )
+    ]
+    assert director.read_rpc_response(conversation_id=cid, extract_additional_payload=True) == [
+        b"123"
+    ]
+
+
 def test_get_properties_async(director: Director):
     properties = ["a", "some"]
     cid = director.get_parameters_async(parameters=properties)