From 417d344f5adccf97cdd25282ab8212c916626087 Mon Sep 17 00:00:00 2001
From: Christian Hopps <chopps@labn.net>
Date: Thu, 6 Jun 2024 08:28:37 +0000
Subject: [PATCH] munet: add support for `timeout` arg to `cmd_*()`

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

Signed-off-by: Christian Hopps <chopps@labn.net>
---
 munet/base.py              | 19 +++++++++++++++++--
 pyproject.toml             |  2 +-
 tests/control/test_cmds.py | 21 +++++++++++++++++++++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/munet/base.py b/munet/base.py
index 8336c26..e77eb15 100644
--- a/munet/base.py
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index 76a34fc..adf259f 100644
--- a/pyproject.toml
+++ b/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 <chopps@labn.net>"]
 license = "GPL-2.0-or-later"
diff --git a/tests/control/test_cmds.py b/tests/control/test_cmds.py
index 8ed4588..15e69d5 100644
--- a/tests/control/test_cmds.py
+++ b/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":