From 3d1e539d2bdecc9763a5d5be86f53da0638e8189 Mon Sep 17 00:00:00 2001
From: mathieui <mathieui@mathieui.net>
Date: Fri, 4 Dec 2020 22:59:34 +0100
Subject: [PATCH] XMLStream: Add a wait_until coroutine

It will set a disposable handler on an event and wait on it with a
specific timeout. Useful for integration tests without callback hell.
---
 slixmpp/xmlstream/xmlstream.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py
index af494903..066d84df 100644
--- a/slixmpp/xmlstream/xmlstream.py
+++ b/slixmpp/xmlstream/xmlstream.py
@@ -12,7 +12,7 @@
     :license: MIT, see LICENSE for more details
 """
 
-from typing import Optional, Set, Callable
+from typing import Optional, Set, Callable, Any
 
 import functools
 import logging
@@ -1130,3 +1130,18 @@ class XMLStream(asyncio.BaseProtocol):
         :param exception: An unhandled exception object.
         """
         pass
+
+    async def wait_until(self, event: str, timeout=30) -> Any:
+        """Utility method to wake on the next firing of an event.
+        (Registers a disposable handler on it)
+
+        :param str event: Event to wait on.
+        :param int timeout: Timeout
+        """
+        fut = asyncio.Future()
+        self.add_event_handler(
+            event,
+            fut.set_result,
+            disposable=True,
+        )
+        return await asyncio.wait_for(fut, timeout)