tests/auth_log: Add method to fetch the next relevant message from the messaging bus

Signed-off-by: Joseph Sutton <josephsutton@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

diff --git a/python/samba/tests/auth_log_base.py b/python/samba/tests/auth_log_base.py
index 9c600ff1f395454d806ef0d80ef81ad9a5e00761..c257f2fbafbf5c3256ac86f8f88a53605219ba36 100644 (file)
--- a/python/samba/tests/auth_log_base.py
+++ b/python/samba/tests/auth_log_base.py
@@ -28,6 +28,10 @@ import os
 import re
 
 
+class NoMessageException(Exception):
+    pass
+
+
 class AuthLogTestBase(samba.tests.TestCase):
 
     @classmethod
@@ -127,6 +131,40 @@ class AuthLogTestBase(samba.tests.TestCase):
         self.connection = None
         return list(filter(self.isRemote, self.context["messages"]))
 
+    def nextMessage(self, msgFilter=None):
+        """Return the next relevant message, or throw a NoMessageException."""
+        def is_relevant(msg):
+            if not self.isRemote(msg):
+                return False
+
+            if msgFilter is None:
+                return True
+
+            return msgFilter(msg)
+
+        messages = self.context['messages']
+
+        while True:
+            timeout = 2
+            until = time.time() + timeout
+
+            while not messages:
+                # Fetch a new message from the messaging bus.
+
+                current = time.time()
+                if until < current:
+                    break
+
+                self.msg_ctx.loop_once(until - current)
+
+            if not messages:
+                raise NoMessageException('timed out looking for a message')
+
+            # Grab the next message from the queue.
+            msg = messages.pop(0)
+            if is_relevant(msg):
+                return msg
+
     # Discard any previously queued messages.
     @classmethod
     def discardMessages(cls):