Add back await_message for use when it makes sense

See function doc comment for when it's not appropriate.

Signed-off-by: Daniel Bluhm <[email protected]>

Author: dbluhm

aries_staticagent/static_connection.py

@@ -333,6 +333,23 @@ async def send_and_await_reply_async(
             )
             return await asyncio.wait_for(next_message, timeout)
 
+    async def await_message(
+            self,
+            *,
+            type_: str = None,
+            condition: Callable[[Message], bool] = None,
+            timeout: int = None):
+        """
+        Wait for a message.
+
+        Note that it's possible for a message to arrive just before or during
+        the setup of this function. If it's likely that a message will arrive
+        as a result of an action taken prior to calling await_message, use the
+        `next` context manager instead.
+        """
+        with self.next(type_, cond=condition) as next_message:
+            return await asyncio.wait_for(next_message, timeout)
+
     def send(self, *args, **kwargs):
         """Blocking wrapper around send_async."""
         loop = asyncio.get_event_loop()

tests/test_send.py

@@ -1,4 +1,5 @@
 """Test StaticConection send method."""
+import asyncio
 import copy
 from functools import partial
 from asyncio import wait_for
@@ -250,6 +251,14 @@ async def test_multiple_next_fulfilled_sequentially(alice, bob):
         assert first == second
 
 
+@pytest.mark.asyncio
+async def test_next_with_type_and_cond_raises_error(alice):
+    """Test value error raised when both type and condition specified."""
+    with pytest.raises(ValueError):
+        with alice.next(MESSAGE.type, lambda msg: True):
+            pass
+
+
 @pytest.mark.asyncio
 async def test_send_and_await_reply(alice_gen, bob, send):
     """Test holding and awaiting messages."""
@@ -258,6 +267,16 @@ async def test_send_and_await_reply(alice_gen, bob, send):
     assert response == RESPONSE
 
 
+@pytest.mark.asyncio
+async def test_await_message(alice, bob):
+    """Test awaiting a message."""
+    waiting_task = asyncio.ensure_future(alice.await_message())
+    await asyncio.sleep(.1)
+    await alice.handle(bob.pack(MESSAGE))
+    message = await waiting_task
+    assert message == MESSAGE
+
+
 def test_blocking_send(alice_gen, bob, send):
     """Test blocking send"""
     alice = alice_gen(send=send)