munet: add new readline and waitline functions

These functions can be used with popen processes during tests to
wait for output (or timeout)

Signed-off-by: Christian Hopps <chopps@labn.net>

Author: choppsv1

munet/testing/util.py

@@ -8,12 +8,17 @@
 """Utility functions useful when using munet testing functionailty in pytest."""
 import asyncio
 import datetime
+import fcntl
 import functools
 import logging
+import os
+import re
+import select
 import sys
 import time
 
 from ..base import BaseMunet
+from ..base import Timeout
 from ..cli import async_cli
 
 
@@ -23,6 +28,7 @@
 
 
 async def async_pause_test(desc=""):
+    """Pause the running of a test offering options for CLI or PDB."""
     isatty = sys.stdout.isatty()
     if not isatty:
         desc = f" for {desc}" if desc else ""
@@ -49,11 +55,12 @@ async def async_pause_test(desc=""):
 
 
 def pause_test(desc=""):
+    """Pause the running of a test offering options for CLI or PDB."""
     asyncio.run(async_pause_test(desc))
 
 
 def retry(retry_timeout, initial_wait=0, retry_sleep=2, expected=True):
-    """decorator: retry while functions return is not None or raises an exception.
+    """Retry decorated function until it returns None, raises an exception, or timeout.
 
     * `retry_timeout`: Retry for at least this many seconds; after waiting
                        initial_wait seconds
@@ -116,3 +123,91 @@ def func_retry(*args, **kwargs):
         return func_retry
 
     return _retry
+
+
+def readline(f, timeout=None):
+    """Read a line or timeout.
+
+    This function will take over the file object, the file object should not be used
+    outside of calling this function once you begin.
+
+    Return: A line, remaining buffer if EOF (subsequent calls will return ""), or None
+    for timeout.
+    """
+    fd = f.fileno()
+    if not hasattr(f, "munet_non_block_set"):
+        flags = fcntl.fcntl(fd, fcntl.F_GETFL)
+        fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
+        f.munet_non_block_set = True
+        f.munet_lines = []
+        f.munet_buf = ""
+
+    if f.munet_lines:
+        return f.munet_lines.pop(0)
+
+    timeout = Timeout(timeout)
+    remaining = timeout.remaining()
+    while remaining > 0:
+        ready, _, _ = select.select([fd], [], [], remaining)
+        if not ready:
+            return None
+
+        c = f.read()
+        if c is None:
+            logging.error("munet readline: unexpected None during read")
+            return None
+
+        if not c:
+            logging.debug("munet readline: got eof")
+            c = f.munet_buf
+            f.munet_buf = ""
+            return c
+
+        f.munet_buf += c
+        while "\n" in f.munet_buf:
+            a, f.munet_buf = f.munet_buf.split("\n", 1)
+            f.munet_lines.append(a + "\n")
+
+        if f.munet_lines:
+            return f.munet_lines.pop(0)
+
+        remaining = timeout.remaining()
+    return None
+
+
+def waitline(f, regex, timeout=120):
+    """Match a regex within lines from a file with a timeout.
+
+    This function will take over the file object (by calling `readline` above), the file
+    object should not be used outside of calling these functions once you begin.
+
+    Return: the match object or None.
+    """
+    timeo = Timeout(timeout)
+    while not timeo.is_expired():
+        line = readline(f, timeo.remaining())
+        if line is None:
+            break
+
+        if line == "":
+            logging.warning("waitline: got eof while matching '%s'", regex)
+            return None
+
+        assert line[-1] == "\n"
+        line = line[:-1]
+        if not line:
+            continue
+
+        logging.debug("waitline: searching: '%s' for '%s'", line, regex)
+        m = re.search(regex, line)
+        if m:
+            logging.debug("waitline: matched '%s'", m.group(0))
+            return m
+
+    logging.warning(
+        "Timeout while getting output matching '%s' within %ss (actual %ss)",
+        regex,
+        timeout,
+        timeo.elapsed(),
+    )
+    return None

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "munet"
-version = "0.15.2"
+version = "0.15.3"
 description = "A package to facilitate network simulations"
 authors = ["Christian Hopps <chopps@labn.net>"]
 license = "GPL-2.0-or-later"