🔎 Add shell and log probe

Author: bdsoha

src/kloudkit/testshed/docker/__init__.py

@@ -2,6 +2,9 @@
 
 from kloudkit.testshed.docker.container import Container
 from kloudkit.testshed.docker.probes.http_probe import HttpProbe
+from kloudkit.testshed.docker.probes.log_probe import LogProbe
+from kloudkit.testshed.docker.probes.probe import Probe
+from kloudkit.testshed.docker.probes.shell_probe import ShellProbe
 from kloudkit.testshed.docker.volumes.inline_volume import InlineVolume
 from kloudkit.testshed.docker.volumes.remote_volume import RemoteVolume
 
@@ -12,5 +15,8 @@
   "DockerException",
   "HttpProbe",
   "InlineVolume",
+  "LogProbe",
+  "Probe",
   "RemoteVolume",
+  "ShellProbe",
 )

src/kloudkit/testshed/docker/factory.py

@@ -1,6 +1,6 @@
 from kloudkit.testshed.core.state import ShedState
 from kloudkit.testshed.docker.container import Container
-from kloudkit.testshed.docker.probes.http_probe import HttpProbe
+from kloudkit.testshed.docker.probes.probe import Probe
 from kloudkit.testshed.docker.probes.readiness_check import ReadinessCheck
 from kloudkit.testshed.docker.runtime.cleanup import Cleanup
 from kloudkit.testshed.docker.volumes.volume_manager import VolumeManager
@@ -23,7 +23,7 @@ def build(
     image: str,
     *,
     detach: bool = True,
-    probe: HttpProbe | None = None,
+    probe: Probe | None = None,
     container_class: type[Container] | None = None,
     test_name: str | None = None,
     **kwargs,

src/kloudkit/testshed/docker/probes/http_probe.py

@@ -1,14 +1,19 @@
 from dataclasses import asdict, dataclass, replace
-from typing import Self
+from typing import TYPE_CHECKING, Self
+
+from kloudkit.testshed.docker.probes.probe import Probe
+
+
+if TYPE_CHECKING:
+  from kloudkit.testshed.docker.container import Container
 
 
 @dataclass(slots=True)
-class HttpProbe:
+class HttpProbe(Probe):
   host: str = "http://localhost"
   port: int | None = None
   endpoint: str | None = None
   command: str = "curl"
-  timeout: float = 30.0
 
   @property
   def url(self) -> str:
@@ -19,6 +24,17 @@ def url(self) -> str:
 
     return "".join((self.host, port, endpoint))
 
+  def check(self, container: "Container") -> None:
+    """Single probe attempt. Raise on failure."""
+
+    container.execute([*self.command.split(" "), self.url], raises=True)
+
+  @property
+  def failure_message(self) -> str:
+    """Message shown on timeout."""
+
+    return f"URL [{self.url}] was not reachable within {self.timeout}s"
+
   def merge(self, other: "HttpProbe", *, ignore_none: bool = True) -> Self:
     """Merge two Probes."""
 

src/kloudkit/testshed/docker/probes/log_probe.py

@@ -0,0 +1,31 @@
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from kloudkit.testshed.docker.probes.probe import Probe
+
+
+if TYPE_CHECKING:
+  from kloudkit.testshed.docker.container import Container
+
+
+@dataclass(slots=True)
+class LogProbe(Probe):
+  pattern: str = ""
+
+  def check(self, container: "Container") -> None:
+    """Single probe attempt. Raise on failure."""
+
+    logs = container.logs()
+    if not re.search(self.pattern, logs):
+      raise ValueError(f"Pattern {self.pattern!r} not found in logs")
+
+  @property
+  def failure_message(self) -> str:
+    """Message shown on timeout."""
+
+    return (
+      f"Pattern {self.pattern!r} not found"
+      f" in container logs within {self.timeout}s"
+    )
+

src/kloudkit/testshed/docker/probes/probe.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+  from kloudkit.testshed.docker.container import Container
+
+
+@dataclass(slots=True)
+class Probe(ABC):
+  timeout: float = 30.0
+
+  @abstractmethod
+  def check(self, container: "Container") -> None:
+    """Single probe attempt. Raise on failure."""
+
+  @property
+  @abstractmethod
+  def failure_message(self) -> str:
+    """Message shown on timeout."""
+
+  @classmethod
+  def resolve(
+    cls,
+    default: "Probe | None",
+    user: "Probe | None",
+    port: int | None = None,
+  ) -> "Probe | None":
+    """Resolve a final probe from default + user override + port."""
+    from kloudkit.testshed.docker.probes.http_probe import HttpProbe
+
+    probe = default
+
+    if port is not None and isinstance(probe, HttpProbe):
+      probe = probe.merge(HttpProbe(port=port))
+
+    if user is ...:
+      return probe
+
+    if isinstance(probe, HttpProbe) and isinstance(user, HttpProbe):
+      return probe.merge(user)
+
+    return user

src/kloudkit/testshed/docker/probes/readiness_check.py

@@ -1,47 +1,35 @@
 import contextlib
 import time
 
-from python_on_whales.exceptions import DockerException
-
 from kloudkit.testshed.docker.container import Container
-from kloudkit.testshed.docker.probes.http_probe import HttpProbe
+from kloudkit.testshed.docker.probes.probe import Probe
 
 import pytest
 
 
 class ReadinessCheck:
   def __init__(
-    self, container: Container, probe: HttpProbe, *, container_logs=None
+    self, container: Container, probe: Probe, *, container_logs=None
   ):
     self._container: Container = container
-    self._probe: HttpProbe = probe
+    self._probe: Probe = probe
     self._container_logs = container_logs
 
-  @property
-  def command(self) -> list[str]:
-    """Full probe test command."""
-
-    return [*self._probe.command.split(" "), self._probe.url]
-
   def wait(self) -> None:
-    """Wait until a container responds on the given endpoint."""
+    """Wait until a container passes the probe check."""
 
     deadline = time.time() + self._probe.timeout
 
-    failure_message = (
-      f"URL [{self._probe.url}] was not reachable within {self._probe.timeout}s"
-    )
-
     while time.time() < deadline:
       try:
-        self._container.execute(self.command, raises=True)
+        self._probe.check(self._container)
 
         return
-      except DockerException:
+      except Exception:
         time.sleep(0.1)
 
     if self._container_logs:
       with contextlib.suppress(Exception):
         self._container_logs(self._container.logs())
 
-    pytest.fail(failure_message, pytrace=False)
+    pytest.fail(self._probe.failure_message, pytrace=False)

src/kloudkit/testshed/docker/probes/shell_probe.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from kloudkit.testshed.docker.probes.probe import Probe
+
+
+if TYPE_CHECKING:
+  from kloudkit.testshed.docker.container import Container
+
+
+@dataclass(slots=True)
+class ShellProbe(Probe):
+  command: str | list[str] = ""
+
+  def check(self, container: "Container") -> None:
+    """Single probe attempt. Raise on failure."""
+
+    cmd = self.command if isinstance(self.command, list) else self.command.split()
+    container.execute(cmd, raises=True)
+
+  @property
+  def failure_message(self) -> str:
+    """Message shown on timeout."""
+
+    return f"Command {self.command!r} did not succeed within {self.timeout}s"

src/kloudkit/testshed/fixtures/shed.py

@@ -1,6 +1,6 @@
 from kloudkit.testshed.core.state import ShedState
 from kloudkit.testshed.docker.container_config import ContainerConfig
-from kloudkit.testshed.docker.probes.http_probe import HttpProbe
+from kloudkit.testshed.docker.probes.probe import Probe
 
 import pytest
 
@@ -42,19 +42,16 @@ def shed_factory(shed_tag, docker_sidecar, shed_container_defaults):
   """Callable factory for spinning up containers with configurable defaults."""
 
   def _wrapper(**kwargs):
-    port = kwargs.pop("port", None)
-    user_probe = kwargs.pop("probe", ...)
-
-    probe = shed_container_defaults.get("probe")
-
-    if port is not None and probe:
-      probe = probe.merge(HttpProbe(port=port))
-
-    if user_probe is not ...:
-      probe = probe.merge(user_probe) if (probe and user_probe) else user_probe
+    probe = Probe.resolve(
+      default=shed_container_defaults.get("probe"),
+      user=kwargs.pop("probe", ...),
+      port=kwargs.pop("port", None),
+    )
 
     merged_config = {**shed_container_defaults, **kwargs}
-    if probe:
+    merged_config.pop("probe", None)
+
+    if probe is not None:
       merged_config["probe"] = probe
 
     return docker_sidecar(image=shed_tag, **merged_config)

tests/integration/test_docker_sidecar.py

@@ -1,4 +1,4 @@
-from kloudkit.testshed.docker import Container, HttpProbe
+from kloudkit.testshed.docker import Container, HttpProbe, LogProbe, ShellProbe
 
 
 def test_container_creation(docker_sidecar):
@@ -34,3 +34,24 @@ def test_nginx_with_probe(docker_sidecar):
 
   assert isinstance(container, Container)
   assert "nginx" in container.execute(["curl", "-s", "http://localhost:80"])
+
+
+def test_nginx_with_log_probe(docker_sidecar):
+  container = docker_sidecar(
+    image="nginx:1.27-alpine",
+    probe=LogProbe(pattern=r"start worker process", timeout=5.0),
+  )
+
+  assert isinstance(container, Container)
+  assert container.state.running
+
+
+def test_alpine_with_shell_probe(docker_sidecar):
+  container = docker_sidecar(
+    image="alpine:latest",
+    command=["sh", "-c", "sleep 1 && touch /tmp/ready && sleep 30"],
+    probe=ShellProbe(command="test -f /tmp/ready", timeout=5.0),
+  )
+
+  assert isinstance(container, Container)
+  assert container.execute(["cat", "/tmp/ready"]) == ""

tests/unit/docker/probes/test_http_probe.py

@@ -1,4 +1,7 @@
+from unittest.mock import MagicMock
+
 from kloudkit.testshed.docker.probes.http_probe import HttpProbe
+from kloudkit.testshed.docker.probes.probe import Probe
 
 import pytest
 
@@ -96,3 +99,37 @@ def test_merge_empty_probe(ignore_none):
   else:
     assert merged.host == "http://localhost"
     assert merged.port is None
+
+
+def test_inherits_from_probe():
+  probe = HttpProbe()
+
+  assert isinstance(probe, Probe)
+
+
+def test_check_executes_command():
+  container = MagicMock()
+  probe = HttpProbe(port=8080, endpoint="/health")
+
+  probe.check(container)
+
+  container.execute.assert_called_once_with(
+    ["curl", "http://localhost:8080/health"], raises=True
+  )
+
+
+def test_check_raises_on_failure():
+  container = MagicMock()
+  container.execute.side_effect = RuntimeError("connection refused")
+  probe = HttpProbe(port=8080)
+
+  with pytest.raises(RuntimeError, match="connection refused"):
+    probe.check(container)
+
+
+def test_failure_message():
+  probe = HttpProbe(port=8080, endpoint="/health", timeout=15.0)
+
+  assert probe.failure_message == (
+    "URL [http://localhost:8080/health] was not reachable within 15.0s"
+  )