munet: add support for timeout arg to cmd_*()

Will raise `subprocess.TimeoutExpired` for both sync and async calls.

Signed-off-by: Christian Hopps <[email protected]>

munet/base.py

@@ -942,15 +942,25 @@ def _cmd_status_finish(self, p, c, ac, o, e, raises, warn):
 
     def _cmd_status(self, cmds, raises=False, warn=True, stdin=None, **kwargs):
         """Execute a command."""
+        timeout = None
+        if "timeout" in kwargs:
+            timeout = kwargs["timeout"]
+            del kwargs["timeout"]
+
         pinput, stdin = Commander._cmd_status_input(stdin)
         p, actual_cmd = self._popen("cmd_status", cmds, stdin=stdin, **kwargs)
-        o, e = p.communicate(pinput)
+        o, e = p.communicate(pinput, timeout=timeout)
         return self._cmd_status_finish(p, cmds, actual_cmd, o, e, raises, warn)
 
     async def _async_cmd_status(
         self, cmds, raises=False, warn=True, stdin=None, text=None, **kwargs
     ):
         """Execute a command."""
+        timeout = None
+        if "timeout" in kwargs:
+            timeout = kwargs["timeout"]
+            del kwargs["timeout"]
+
         pinput, stdin = Commander._cmd_status_input(stdin)
         p, actual_cmd = await self._async_popen(
             "async_cmd_status", cmds, stdin=stdin, **kwargs
@@ -963,7 +973,12 @@ async def _async_cmd_status(
 
         if encoding is not None and isinstance(pinput, str):
             pinput = pinput.encode(encoding)
-        o, e = await p.communicate(pinput)
+        try:
+            o, e = await asyncio.wait_for(p.communicate(), timeout=timeout)
+        except (TimeoutError, asyncio.TimeoutError) as error:
+            raise subprocess.TimeoutExpired(
+                cmd=actual_cmd, timeout=timeout, output=None, stderr=None
+            ) from error
         if encoding is not None:
             o = o.decode(encoding) if o is not None else o
             e = e.decode(encoding) if e is not None else e

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "munet"
-version = "0.14.8"
+version = "0.14.9"
 description = "A package to facilitate network simulations"
 authors = ["Christian Hopps <[email protected]>"]
 license = "GPL-2.0-or-later"

tests/control/test_cmds.py

@@ -126,6 +126,27 @@ async def test_cmd_nostatus(unet, host):
     assert "No such file or directory" in e
 
 
+@pytest.mark.parametrize("host", ["host1", "container1", "remote1", "hn1"])
+async def test_cmd_status_timeout(unet, host):
+    if host == "remote1":
+        await wait_remote_up(unet)
+    host = unet.hosts[host]
+
+    try:
+        host.cmd_status("sleep 10", timeout=1)
+    except subprocess.TimeoutExpired:
+        pass
+    else:
+        assert False, "No timeout raised"
+
+    try:
+        await host.async_cmd_status("sleep 10", timeout=1)
+    except subprocess.TimeoutExpired:
+        pass
+    else:
+        assert False, "No timeout raised"
+
+
 @pytest.mark.parametrize("host", ["host1", "container1", "remote1", "hn1"])
 async def test_popen(unet, host):
     if host == "remote1":