Add timeout kwarg to exec (#673)

Author: jacobtomlinson

kr8s/_exec.py

@@ -35,6 +35,7 @@ def __init__(
         stderr: BinaryIO | None = None,
         check: bool = True,
         capture_output: bool = True,
+        timeout: int | None = None,
     ) -> None:
         self._resource = resource
         self._container = container
@@ -43,6 +44,7 @@ def __init__(
         self._stdout = stdout
         self._stderr = stderr
         self._capture_output = capture_output
+        self._timeout = timeout
 
         self.args = command
         self.stdout = b""
@@ -80,7 +82,7 @@ async def run(
                     await ws.send_bytes(STDIN_CHANNEL.to_bytes() + self._stdin.read())  # type: ignore
                 await ws.send_bytes(CLOSE_CHANNEL.to_bytes() + STDIN_CHANNEL.to_bytes())  # type: ignore
             while True:
-                message = await ws.receive_bytes()
+                message = await ws.receive_bytes(timeout=self._timeout)
                 channel, message = int(message[0]), message[1:]
                 if message:
                     if channel == STDOUT_CHANNEL:

kr8s/_objects.py

@@ -506,6 +506,7 @@ async def exec(
         stderr: BinaryIO | None = None,
         check: bool = True,
         capture_output: bool = True,
+        timeout: int | None = None,
     ) -> CompletedExec:
         """Execute a command in this object."""
         return await self.async_exec(
@@ -516,6 +517,7 @@ async def exec(
             stderr=stderr,
             check=check,
             capture_output=capture_output,
+            timeout=timeout,
         )
 
     async def async_exec(
@@ -528,6 +530,7 @@ async def async_exec(
         stderr: BinaryIO | None = None,
         check: bool = True,
         capture_output: bool = True,
+        timeout: int | None = None,
     ) -> CompletedExec:
         """Execute a command in this object."""
         if not hasattr(self, "ready_pods"):
@@ -544,6 +547,7 @@ async def async_exec(
             stderr=stderr,
             check=check,
             capture_output=capture_output,
+            timeout=timeout,
         )
 
     async def async_watch(self) -> AsyncGenerator[tuple[str, Self]]:
@@ -1393,6 +1397,7 @@ async def async_exec(
         stderr: BinaryIO | None = None,
         check: bool = True,
         capture_output: bool = True,
+        timeout: int | None = None,
     ) -> CompletedExec:
         while not await self.async_ready():
             await anyio.sleep(0.1)
@@ -1406,6 +1411,7 @@ async def async_exec(
             stdin=stdin,
             check=check,
             capture_output=capture_output,
+            timeout=timeout,
         )
         async with ex.run() as process:
             await process.wait()
@@ -1421,6 +1427,7 @@ async def exec(
         stderr: BinaryIO | None = None,
         check: bool = True,
         capture_output: bool = True,
+        timeout: int | None = None,
     ) -> CompletedExec:
         """Run a command in a container and wait until it completes.
 
@@ -1441,6 +1448,8 @@ async def exec(
                 If True, raise an exception if the command fails.
             capture_output:
                 If True, store stdout and stderr from the container in an attribute.
+            timeout:
+                If set, timeout for recieving bytes.
 
         Returns:
             A :class:`kr8s._exec.CompletedExec` object.
@@ -1460,6 +1469,7 @@ async def exec(
             stderr=stderr,
             check=check,
             capture_output=capture_output,
+            timeout=timeout,
         )
 
     @classmethod

kr8s/objects.py

@@ -241,6 +241,7 @@ def exec(  # type: ignore[override]
         stderr: BinaryIO | None = None,
         check: bool = True,
         capture_output: bool = True,
+        timeout: int | None = None,
     ) -> CompletedExec:
         return as_sync_func(self.async_exec)(
             command,
@@ -250,6 +251,7 @@ def exec(  # type: ignore[override]
             stderr=stderr,
             check=check,
             capture_output=capture_output,
+            timeout=timeout,
         )
 
     def tolerate(  # type: ignore[override]

kr8s/tests/test_objects.py

@@ -1110,6 +1110,15 @@ async def test_pod_exec(ubuntu_pod):
     assert ex.returncode == 0
 
 
+async def test_pod_exec_timeout(ubuntu_pod):
+    ex = await ubuntu_pod.exec(["date"], timeout=2)
+    assert isinstance(ex, CompletedExec)
+    assert str(datetime.datetime.now().year) in ex.stdout.decode()
+    assert ex.args == ["date"]
+    assert ex.stderr == b""
+    assert ex.returncode == 0
+
+
 async def test_pod_exec_error(ubuntu_pod):
     with pytest.raises(ExecError):
         await ubuntu_pod.exec(["date", "foo"])