Added systemd notify support with ready, stopping, and watchdog notifications.

tests/test_cli.py

@@ -12,19 +12,106 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for CLI argument parsing and credential reading.
+"""Unit tests for CLI argument parsing, credentials, and systemd integration.
 
-Tests cover _parse_args and _read_credential from crabstero.__main__.
+Tests cover the CLI helpers and runtime orchestration from crabstero.__main__.
 """
 
-from typing import TYPE_CHECKING
+import asyncio
+import signal
+import socket
+import sys
+from typing import TYPE_CHECKING, Any, ClassVar, Self
 
 import pytest
 
-from crabstero.__main__ import _parse_args, _read_credential
+from crabstero import __main__ as cli
+from crabstero.__main__ import (
+    _parse_args,
+    _read_credential,
+    _sd_notify,
+    _watchdog_interval,
+    _watchdog_loop,
+)
 
 if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine
     from pathlib import Path
+    from types import TracebackType
+
+
+class _FakeCrabstero:
+    events: ClassVar[list[object]] = []
+    login_error: ClassVar[BaseException | None] = None
+    close_error: ClassVar[BaseException | None] = None
+    connect_callback: ClassVar[Callable[[], None] | None] = None
+
+    def __init__(self, **kwargs: object) -> None:
+        self._closed = False
+        self.events.append(("init", kwargs))
+
+    async def __aenter__(self) -> Self:
+        self.events.append("enter")
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.events.append("exit")
+        await self.close()
+
+    async def login(self, token: str) -> None:
+        self.events.append(("login", token))
+        if self.login_error is not None:
+            raise self.login_error
+
+    async def connect(self, **kwargs: object) -> None:
+        self.events.append(("connect", kwargs))
+        callback = type(self).connect_callback
+        if callback is not None:
+            callback()
+            await asyncio.sleep(0)
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self.events.append("close")
+        self._closed = True
+        if self.close_error is not None:
+            raise self.close_error
+
+    def is_closed(self) -> bool:
+        return self._closed
+
+
+def _run_coroutine(coro: Coroutine[Any, Any, None]) -> None:
+    asyncio.run(coro)
+
+
+def _prepare_main_test(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    login_error: BaseException | None = None,
+    close_error: BaseException | None = None,
+) -> list[object]:
+    _FakeCrabstero.events = []
+    _FakeCrabstero.login_error = login_error
+    _FakeCrabstero.close_error = close_error
+    _FakeCrabstero.connect_callback = None
+
+    def sd_notify(state: str) -> None:
+        _FakeCrabstero.events.append(("notify", state))
+
+    monkeypatch.setattr(sys, "argv", ["crabstero", "--token", "cli-token"])
+    monkeypatch.setattr(cli, "Crabstero", _FakeCrabstero)
+    monkeypatch.setattr("crabstero.__main__.uvloop.run", _run_coroutine)
+    monkeypatch.setattr(cli, "_sd_notify", sd_notify)
+    monkeypatch.setattr(cli, "_watchdog_interval", lambda: None)
+
+    return _FakeCrabstero.events
 
 
 class TestParseArgs:
@@ -117,3 +204,168 @@ class TestReadCredential:
         """Returns None when the credential file does not exist."""
         monkeypatch.setenv("CREDENTIALS_DIRECTORY", str(tmp_path))
         assert _read_credential("nonexistent") is None
+
+
+class TestSystemdNotify:
+    """Systemd notification helper behavior."""
+
+    def test_sd_notify_noop_without_socket(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Missing NOTIFY_SOCKET is a no-op."""
+        monkeypatch.delenv("NOTIFY_SOCKET", raising=False)
+        _sd_notify("READY=1")
+
+    def test_sd_notify_sends_datagram(
+        self,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """_sd_notify sends the payload to NOTIFY_SOCKET."""
+        socket_path = tmp_path / "notify.sock"
+
+        with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as server:
+            server.bind(str(socket_path))
+            server.settimeout(1)
+            monkeypatch.setenv("NOTIFY_SOCKET", str(socket_path))
+
+            _sd_notify("READY=1")
+
+            assert server.recv(1024) == b"READY=1"
+
+
+class TestWatchdog:
+    """Systemd watchdog helper behavior."""
+
+    def test_watchdog_interval_missing(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Missing WATCHDOG_USEC disables watchdog pings."""
+        monkeypatch.delenv("WATCHDOG_USEC", raising=False)
+        assert _watchdog_interval() is None
+
+    def test_watchdog_interval_valid(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """WATCHDOG_USEC is converted to half the interval in seconds."""
+        monkeypatch.setenv("WATCHDOG_USEC", "4000000")
+        assert _watchdog_interval() == 2
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            pytest.param("invalid", id="not-integer"),
+            pytest.param("0", id="zero"),
+            pytest.param("-1", id="negative"),
+        ],
+    )
+    def test_watchdog_interval_invalid(
+        self,
+        value: str,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Invalid WATCHDOG_USEC values disable watchdog pings."""
+        monkeypatch.setenv("WATCHDOG_USEC", value)
+        assert _watchdog_interval() is None
+
+    async def test_watchdog_loop_sends_until_stopped(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """_watchdog_loop sends WATCHDOG=1 and exits when stopped."""
+        notifications: list[str] = []
+        stop = asyncio.Event()
+        monkeypatch.setattr(cli, "_sd_notify", notifications.append)
+
+        task = asyncio.create_task(_watchdog_loop(60, stop))
+        await asyncio.sleep(0)
+        stop.set()
+        await asyncio.wait_for(task, timeout=1)
+
+        assert notifications == ["WATCHDOG=1"]
+
+
+class TestMainLifecycle:
+    """CLI runtime orchestration."""
+
+    def test_ready_sent_after_login_before_connect(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """READY=1 is sent after login succeeds and before connect starts."""
+        events = _prepare_main_test(monkeypatch)
+
+        cli.main()
+
+        assert events.index(("login", "cli-token")) < events.index(
+            ("notify", "READY=1"),
+        )
+        assert events.index(("notify", "READY=1")) < events.index(
+            ("connect", {"reconnect": True}),
+        )
+
+    def test_stopping_sent_after_clean_connect_return(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """STOPPING=1 is sent when the CLI run exits cleanly after readiness."""
+        events = _prepare_main_test(monkeypatch)
+
+        cli.main()
+
+        assert events.index(("connect", {"reconnect": True})) < events.index(
+            ("notify", "STOPPING=1"),
+        )
+
+    def test_ready_not_sent_when_login_fails(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """A login failure exits without reporting readiness."""
+        events = _prepare_main_test(
+            monkeypatch,
+            login_error=RuntimeError("login failed"),
+        )
+
+        with pytest.raises(RuntimeError, match="login failed"):
+            cli.main()
+
+        assert ("notify", "READY=1") not in events
+        assert ("notify", "STOPPING=1") not in events
+
+    def test_shutdown_task_exception_is_propagated(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """A completed shutdown task exception is not silently dropped."""
+        signal_handlers: dict[int, Callable[[], None]] = {}
+
+        class SignalLoop:
+            def add_signal_handler(
+                self,
+                sig: int,
+                callback: Callable[[], None],
+            ) -> None:
+                signal_handlers[sig] = callback
+
+            def remove_signal_handler(self, sig: int) -> None:
+                signal_handlers.pop(sig, None)
+
+        events = _prepare_main_test(
+            monkeypatch,
+            close_error=RuntimeError("close failed"),
+        )
+
+        def request_shutdown() -> None:
+            signal_handlers[signal.SIGTERM]()
+
+        monkeypatch.setattr(asyncio, "get_running_loop", SignalLoop)
+        _FakeCrabstero.connect_callback = request_shutdown
+
+        with pytest.raises(RuntimeError, match="close failed"):
+            cli.main()
+
+        assert events.index(("notify", "STOPPING=1")) < events.index("close")