perf: bootstrap Caido concurrently with the scan start (#1143)

Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>

tests/test_caido_bootstrap.py

@@ -0,0 +1,81 @@
+"""A bootstrap that dies mid-setup must not leave its transport behind.
+
+The bootstrap now runs concurrently with the scan start, so teardown can
+cancel it at any await — including inside ``Client.connect()``, where the
+client exists but no caller will ever see it to close it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+import types
+from typing import Any
+
+import pytest
+
+from strix.runtime.caido_bootstrap import bootstrap_caido
+
+
+class _FakeExecResult:
+    stderr = b""
+    exit_code = 0
+
+    def __init__(self, stdout: str) -> None:
+        self.stdout = stdout
+
+    def ok(self) -> bool:
+        return True
+
+
+class _FakeSession:
+    async def exec(self, *_args: Any, **_kwargs: Any) -> _FakeExecResult:
+        return _FakeExecResult('{"data":{"loginAsGuest":{"token":{"accessToken":"t"}}}}')
+
+
+class _FakeClient:
+    def __init__(self, connect_error: BaseException) -> None:
+        self.connect_error = connect_error
+        self.closed = False
+
+    async def connect(self) -> None:
+        raise self.connect_error
+
+    async def aclose(self) -> None:
+        self.closed = True
+
+
+async def _bootstrap_expecting(
+    monkeypatch: pytest.MonkeyPatch, error: BaseException
+) -> _FakeClient:
+    """Run a bootstrap whose ``connect()`` fails with ``error``."""
+    client = _FakeClient(error)
+    # The SDK is imported inside bootstrap_caido (it is slow to import), so the
+    # fakes are injected as the modules it imports.
+    sdk = types.ModuleType("caido_sdk_client")
+    sdk.Client = lambda *_a, **_k: client  # type: ignore[attr-defined]
+    sdk.TokenAuthOptions = lambda token: token  # type: ignore[attr-defined]
+    sdk_types = types.ModuleType("caido_sdk_client.types")
+    sdk_types.CreateProjectOptions = lambda **_k: None  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "caido_sdk_client", sdk)
+    monkeypatch.setitem(sys.modules, "caido_sdk_client.types", sdk_types)
+
+    with pytest.raises(type(error)):
+        await bootstrap_caido(
+            _FakeSession(),  # type: ignore[arg-type]
+            host_url="http://host",
+            container_url="http://container",
+        )
+    return client
+
+
+async def test_cancellation_during_connect_closes_the_client(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client = await _bootstrap_expecting(monkeypatch, asyncio.CancelledError())
+    assert client.closed
+
+
+async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None:
+    client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener"))
+    assert client.closed

tests/test_caido_handle.py

@@ -0,0 +1,103 @@
+"""Tests for the concurrent Caido bootstrap handle."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+
+from strix.runtime.caido_handle import CaidoBootstrapHandle
+
+
+class _FakeClient:
+    def __init__(self) -> None:
+        self.closed = False
+
+    async def aclose(self) -> None:
+        self.closed = True
+
+
+def _handle(coro: Any) -> CaidoBootstrapHandle:
+    return CaidoBootstrapHandle(asyncio.ensure_future(coro))
+
+
+async def test_get_waits_for_the_bootstrap() -> None:
+    client = _FakeClient()
+    started = asyncio.Event()
+
+    async def _bootstrap() -> Any:
+        started.set()
+        await asyncio.sleep(0.01)
+        return client
+
+    handle = _handle(_bootstrap())
+    await started.wait()
+    assert handle.peek() is None
+    assert await handle.get() is client
+    assert handle.peek() is client
+
+
+async def test_get_reraises_bootstrap_failure_to_every_caller() -> None:
+    async def _bootstrap() -> Any:
+        raise RuntimeError("caido never came up")
+
+    handle = _handle(_bootstrap())
+    for _ in range(2):
+        with pytest.raises(RuntimeError, match="caido never came up"):
+            await handle.get()
+    assert handle.peek() is None
+
+
+async def test_caller_cancellation_does_not_cancel_the_shared_bootstrap() -> None:
+    client = _FakeClient()
+
+    async def _bootstrap() -> Any:
+        await asyncio.sleep(0.05)
+        return client
+
+    handle = _handle(_bootstrap())
+
+    with pytest.raises(TimeoutError):
+        await asyncio.wait_for(handle.get(), timeout=0.01)
+
+    assert await handle.get() is client
+
+
+async def test_aclose_closes_a_finished_client() -> None:
+    client = _FakeClient()
+
+    async def _bootstrap() -> Any:
+        return client
+
+    handle = _handle(_bootstrap())
+    await handle.get()
+    await handle.aclose()
+    assert client.closed is True
+
+
+async def test_aclose_cancels_an_in_flight_bootstrap() -> None:
+    cancelled = asyncio.Event()
+
+    async def _bootstrap() -> Any:
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            cancelled.set()
+            raise
+        return _FakeClient()
+
+    handle = _handle(_bootstrap())
+    await asyncio.sleep(0)
+    await handle.aclose()
+    assert cancelled.is_set()
+
+
+async def test_aclose_swallows_a_failed_bootstrap() -> None:
+    async def _bootstrap() -> Any:
+        raise RuntimeError("boom")
+
+    handle = _handle(_bootstrap())
+    with pytest.raises(RuntimeError, match="boom"):
+        await handle.get()
+    await handle.aclose()

tests/test_proxy_client.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast
 
 import pytest
 
+from strix.runtime.caido_handle import CaidoBootstrapHandle
 from strix.tools.proxy import caido_api, tools
 
 
@@ -198,12 +199,31 @@ class _Ctx:
         self.context = context
 
 
-def test_ctx_client_returns_client_when_present() -> None:
+async def test_ctx_client_returns_client_when_present() -> None:
     client = _FakeClient("host")
-    got = tools._ctx_client(cast("Any", _Ctx({"caido_client": client})))
+    got = await tools._ctx_client(cast("Any", _Ctx({"caido_client": client})))
     assert got is client
 
 
-def test_ctx_client_returns_none_without_client() -> None:
-    assert tools._ctx_client(cast("Any", _Ctx({}))) is None
-    assert tools._ctx_client(cast("Any", _Ctx(None))) is None
+async def test_ctx_client_returns_none_without_client() -> None:
+    assert await tools._ctx_client(cast("Any", _Ctx({}))) is None
+    assert await tools._ctx_client(cast("Any", _Ctx(None))) is None
+
+
+async def test_ctx_client_resolves_bootstrap_handle() -> None:
+    client = _FakeClient("host")
+
+    async def _bootstrap() -> Any:
+        return client
+
+    handle = CaidoBootstrapHandle(asyncio.ensure_future(_bootstrap()))
+    got = await tools._ctx_client(cast("Any", _Ctx({"caido_client": handle})))
+    assert got is client
+
+
+async def test_ctx_client_degrades_when_bootstrap_failed() -> None:
+    async def _bootstrap() -> Any:
+        raise RuntimeError("caido never came up")
+
+    handle = CaidoBootstrapHandle(asyncio.ensure_future(_bootstrap()))
+    assert await tools._ctx_client(cast("Any", _Ctx({"caido_client": handle}))) is None