mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
b71b659e1b
Add websockets to the dev extra so the repro harness smoke test can import its websocket client dependency in the CI test matrix. Also apply ruff formatting to the files the formatter check was rejecting so the 3.12 lint job passes.
280 lines
9.3 KiB
Python
280 lines
9.3 KiB
Python
"""Concurrency + timeout tests for MemoryHandler._ensure_initialized (Unit 1).
|
|
|
|
Covers:
|
|
- 10 concurrent first-callers trigger exactly one backend init
|
|
(singleflight via asyncio.Lock + double-check).
|
|
- Initialization timeout is surfaced via log + leaves _initialized=False so
|
|
a subsequent call can retry (fail-open contract).
|
|
- End-to-end real-backend sanity test (no monkeypatching of internals)
|
|
that exercises LocalBackend so the system-wide check is satisfied.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from headroom.proxy.memory_handler import (
|
|
STARTUP_INIT_TIMEOUT_SECONDS,
|
|
MemoryConfig,
|
|
MemoryHandler,
|
|
)
|
|
|
|
# -------------------------------------------------------------------
|
|
# Singleflight under concurrent callers
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_ensure_initialized_runs_init_once(tmp_path, monkeypatch):
|
|
hits = {"n": 0}
|
|
start_event = asyncio.Event()
|
|
release_event = asyncio.Event()
|
|
|
|
class FakeLocalBackend:
|
|
def __init__(self, config): # noqa: D401
|
|
self.config = config
|
|
|
|
async def _ensure_initialized(self) -> None:
|
|
hits["n"] += 1
|
|
start_event.set()
|
|
# Simulate a slow cold-start so concurrent callers pile up on
|
|
# the lock. Without singleflight, hits would exceed 1.
|
|
await release_event.wait()
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
import headroom.memory.backends.local as local_mod
|
|
|
|
monkeypatch.setattr(local_mod, "LocalBackend", FakeLocalBackend)
|
|
|
|
handler = MemoryHandler(
|
|
MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
|
|
)
|
|
|
|
async def caller() -> None:
|
|
await handler._ensure_initialized()
|
|
|
|
# Fire 10 concurrent first-callers.
|
|
tasks = [asyncio.create_task(caller()) for _ in range(10)]
|
|
# Let one task enter the critical section, then release it.
|
|
await start_event.wait()
|
|
release_event.set()
|
|
await asyncio.gather(*tasks)
|
|
|
|
assert hits["n"] == 1, f"expected 1 backend init under concurrency, got {hits['n']}"
|
|
assert handler._initialized is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_initialized_noop_when_disabled(tmp_path):
|
|
handler = MemoryHandler(
|
|
MemoryConfig(enabled=False, backend="local", db_path=str(tmp_path / "mem.db"))
|
|
)
|
|
await handler._ensure_initialized()
|
|
assert handler._initialized is False
|
|
assert handler._backend is None
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Timeout fail-open
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_initialized_timeout_leaves_handler_unready(tmp_path, monkeypatch):
|
|
class HangingBackend:
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
async def _ensure_initialized(self) -> None:
|
|
# Hang forever — timeout must cancel this.
|
|
await asyncio.Event().wait()
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
import headroom.memory.backends.local as local_mod
|
|
|
|
monkeypatch.setattr(local_mod, "LocalBackend", HangingBackend)
|
|
|
|
handler = MemoryHandler(
|
|
MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
|
|
)
|
|
|
|
# Attach a handler directly to the module logger — caplog has trouble
|
|
# when third-party conftest monkeys with propagation settings.
|
|
import logging as _logging
|
|
|
|
mem_logger = _logging.getLogger("headroom.proxy.memory_handler")
|
|
captured: list[_logging.LogRecord] = []
|
|
|
|
class _ListHandler(_logging.Handler):
|
|
def emit(self, record: _logging.LogRecord) -> None:
|
|
captured.append(record)
|
|
|
|
handler_log = _ListHandler(level=_logging.ERROR)
|
|
mem_logger.addHandler(handler_log)
|
|
prev_level = mem_logger.level
|
|
mem_logger.setLevel(_logging.DEBUG)
|
|
try:
|
|
# Shrink the module-level timeout to keep the test fast.
|
|
with patch("headroom.proxy.memory_handler.STARTUP_INIT_TIMEOUT_SECONDS", 0.1):
|
|
await handler._ensure_initialized()
|
|
finally:
|
|
mem_logger.removeHandler(handler_log)
|
|
mem_logger.setLevel(prev_level)
|
|
|
|
# Fail-open: no exception, _initialized=False, error logged.
|
|
assert handler._initialized is False
|
|
found_timeout_log = any("timed out" in rec.getMessage().lower() for rec in captured)
|
|
assert found_timeout_log, (
|
|
f"expected 'timed out' log record; got: {[(r.levelname, r.getMessage()) for r in captured]}"
|
|
)
|
|
|
|
# Confirm the default constant is unchanged (sanity).
|
|
assert STARTUP_INIT_TIMEOUT_SECONDS == 30.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
|
|
tmp_path, monkeypatch
|
|
):
|
|
"""If wait_for fires while _init_backend_locked has already set
|
|
``self._backend`` but before ``self._initialized = True``, the timeout
|
|
handler must null ``_backend``. Otherwise callers doing
|
|
``if self.memory_handler._backend:`` see a truthy-but-broken backend.
|
|
"""
|
|
|
|
class SlowBackend:
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
async def _ensure_initialized(self) -> None:
|
|
# Hang long enough to blow the 0.01s timeout below.
|
|
await asyncio.sleep(5.0)
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
import headroom.memory.backends.local as local_mod
|
|
|
|
monkeypatch.setattr(local_mod, "LocalBackend", SlowBackend)
|
|
|
|
handler = MemoryHandler(
|
|
MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
|
|
)
|
|
|
|
with patch("headroom.proxy.memory_handler.STARTUP_INIT_TIMEOUT_SECONDS", 0.01):
|
|
await handler._ensure_initialized()
|
|
|
|
# Both must be consistent after timeout.
|
|
assert handler._initialized is False
|
|
assert handler._backend is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_initialized_cancellation_propagates_and_resets_state(tmp_path, monkeypatch):
|
|
"""External cancellation of an in-flight ``_ensure_initialized`` must
|
|
propagate (CancelledError is BaseException — not a swallowable error)
|
|
and leave the handler in a clean state."""
|
|
|
|
class HangingBackend:
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
async def _ensure_initialized(self) -> None:
|
|
await asyncio.Event().wait()
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
import headroom.memory.backends.local as local_mod
|
|
|
|
monkeypatch.setattr(local_mod, "LocalBackend", HangingBackend)
|
|
|
|
handler = MemoryHandler(
|
|
MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
|
|
)
|
|
|
|
task = asyncio.create_task(handler._ensure_initialized())
|
|
# Give the task a tick to enter _init_backend_locked and assign _backend.
|
|
await asyncio.sleep(0.05)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
assert handler._initialized is False
|
|
assert handler._backend is None
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Real backend init (no monkeypatching) — integration smoke test
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_localbackend_initializes_via_public_entrypoint(tmp_path):
|
|
"""End-to-end sanity check: the public ``ensure_initialized`` path works
|
|
against a real LocalBackend. This catches regressions where the new
|
|
asyncio.Lock layering (Unit 1) breaks the actual init chain."""
|
|
pytest.importorskip("sqlite3") # bundled with Python but be explicit
|
|
# Skip on environments without an embedder backend available.
|
|
try:
|
|
import onnxruntime # noqa: F401
|
|
except Exception:
|
|
try:
|
|
import sentence_transformers # noqa: F401
|
|
except Exception: # pragma: no cover - env-dependent
|
|
pytest.skip("No embedder backend available in this environment")
|
|
|
|
handler = MemoryHandler(
|
|
MemoryConfig(
|
|
enabled=True,
|
|
backend="local",
|
|
db_path=str(tmp_path / "mem.db"),
|
|
)
|
|
)
|
|
|
|
# Concurrent callers should still produce one initialized backend.
|
|
await asyncio.gather(*[handler.ensure_initialized() for _ in range(5)])
|
|
assert handler._initialized is True
|
|
assert handler._backend is not None
|
|
|
|
# warmup_embedder is best-effort; on a real backend it should succeed.
|
|
warmed = await handler.warmup_embedder()
|
|
assert warmed is True
|
|
await handler.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_warmup_embedder_returns_false_without_backend():
|
|
handler = MemoryHandler(MemoryConfig(enabled=True, backend="local"))
|
|
# _initialized=False, _backend=None → no-op, no crash.
|
|
result = await handler.warmup_embedder()
|
|
assert result is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_warmup_embedder_swallows_exceptions():
|
|
handler = MemoryHandler(MemoryConfig(enabled=True, backend="local"))
|
|
handler._initialized = True
|
|
|
|
class HM:
|
|
class _E:
|
|
async def embed(self, _text: Any) -> Any:
|
|
raise RuntimeError("synthetic embedder failure")
|
|
|
|
_embedder = _E()
|
|
|
|
class Backend:
|
|
_hierarchical_memory = HM()
|
|
|
|
handler._backend = Backend()
|
|
# Must not raise; returns False.
|
|
result = await handler.warmup_embedder()
|
|
assert result is False
|