extend MessageWaiter

ovos_bus_client/client/client.py

@@ -288,9 +288,9 @@ def wait_for_response(self, message: Message,
         if isinstance(reply_type, list):
             message_type = reply_type
         elif isinstance(reply_type, str):
-            message_type = reply_type
+            message_type = [reply_type]
         elif reply_type is None:
-            message_type = message.msg_type + '.response'
+            message_type = [message.msg_type + '.response']
         waiter = MessageWaiter(self, message_type)  # Setup response handler
         # Send message and wait for its response
         self.emit(message)

ovos_bus_client/client/waiter.py

@@ -12,6 +12,7 @@
 #
 
 from threading import Event
+from typing import List, Union
 
 try:
     from mycroft_bus_client.client.waiter import MessageWaiter as _MessageWaiterBase
@@ -28,20 +29,23 @@ class MessageWaiter:
     """Wait for a single message.
 
     Encapsulate the wait for a message logic separating the setup from
-    the actual waiting act so the waiting can be setuo, actions can be
+    the actual waiting act so the waiting can be setup, actions can be
     performed and _then_ the message can be waited for.
 
-    Argunments:
+    Arguments:
         bus: Bus to check for messages on
-        message_type: message type to wait for
+        message_type: message type(s) to wait for
     """
-    def __init__(self, bus, message_type):
+    def __init__(self, bus, message_type: Union[str, List[str]]):
         self.bus = bus
+        if not isinstance(message_type, list):
+            message_type = [message_type]
         self.msg_type = message_type
         self.received_msg = None
         # Setup response handler
         self.response_event = Event()
-        self.bus.once(message_type, self._handler)
+        for msg in self.msg_type:
+            self.bus.once(msg, self._handler)
 
     def _handler(self, message):
         """Receive response data."""
@@ -61,7 +65,8 @@ def wait(self, timeout=3.0):
         if not self.response_event.is_set():
             # Clean up the event handler
             try:
-                self.bus.remove(self.msg_type, self._handler)
+                for msg in self.msg_type:
+                    self.bus.remove(msg, self._handler)
             except (ValueError, KeyError):
                 # ValueError occurs on pyee 5.0.1 removing handlers
                 # registered with once.

test/unittests/test_client.py

@@ -11,7 +11,7 @@
 # limitations under the License.
 
 import unittest
-from unittest.mock import Mock, patch
+from unittest.mock import call, Mock, patch
 
 from pyee import ExecutorEventEmitter
 
@@ -172,6 +172,17 @@ def test_message_wait_timeout(self):
 
         assert waiter.wait(0.3) is None
 
+    def test_message_converts_to_list(self):
+        bus = Mock()
+        waiter = MessageWaiter(bus, "test.message")
+        assert isinstance(waiter.msg_type, list)
+        bus.once.assert_called_with("test.message", waiter._handler)
+
+    def test_multiple_messages(self):
+        bus = Mock()
+        waiter = MessageWaiter(bus, ["test.message", "test.message2"])
+        bus.once.assert_has_calls([call("test.message", waiter._handler), call("test.message2", waiter._handler)])
+
 
 class TestMessageCollector:
     def test_message_wait_success(self):