cover run-batch, serve, and _stop_pool with tests

This commit is contained in:
Yiorgis Gozadinos 2026-05-29 17:43:58 +03:00
parent 5ba7838b71
commit 6f2a40c676
No known key found for this signature in database
3 changed files with 116 additions and 5 deletions

View file

@ -166,7 +166,7 @@ class IngesterApp:
for sig in (signal.SIGINT, signal.SIGTERM):
try:
loop.add_signal_handler(sig, stop_event.set)
except NotImplementedError:
except NotImplementedError: # pragma: no cover - Windows only
# Windows; signal handlers unavailable in asyncio.
pass

View file

@ -0,0 +1,37 @@
"""haiku-ingester run-batch CLI: echoes the batch report and exits non-zero
when any job dead-letters. IngesterApp is patched so no real ingestion runs."""
from unittest.mock import AsyncMock
from typer.testing import CliRunner
from haiku.rag.ingester.app import BatchReport
from haiku.rag.ingester.cli import _cli as cli
runner = CliRunner()
def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
fake = AsyncMock()
fake.run_batch.return_value = report
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
return fake
def test_run_batch_reports_and_exits_zero(monkeypatch):
fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch)
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
assert result.exit_code == 0
assert "3 succeeded, 0 dead" in result.output
fake.run_batch.assert_awaited_once()
def test_run_batch_exits_nonzero_when_dead(monkeypatch):
_fake_app(BatchReport(succeeded=1, dead=2), monkeypatch)
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
assert result.exit_code == 1
assert "2 dead" in result.output

View file

@ -1,7 +1,7 @@
"""IngesterApp.run_batch: one discover sweep across configured sources,
drain the queue, then exit. The document store (HaikuRAG) is patched out
the behavior under test is the sweep -> queue -> worker -> drain
orchestration and orphan pruning, not embedding."""
"""IngesterApp lifecycle: run_batch (one sweep -> drain -> exit) and serve
(pollers + workers + API until shutdown). The document store (HaikuRAG) is
patched out the behavior under test is the orchestration and orphan
pruning, not embedding."""
import asyncio
from contextlib import asynccontextmanager
@ -23,6 +23,7 @@ from haiku.rag.config import (
WorkerConfig,
)
from haiku.rag.ingester.app import IngesterApp
from haiku.rag.ingester.workers.pool import WorkerPool
from haiku.rag.store.models.document import Document
@ -203,3 +204,76 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
assert report.succeeded == 0
assert report.dead == 0
client.create_document_from_source.assert_not_awaited()
async def _wait_until(predicate, *, timeout: float = 5.0):
deadline = asyncio.get_running_loop().time() + timeout
while not predicate():
if asyncio.get_running_loop().time() > deadline:
raise AssertionError("condition not reached within timeout")
await asyncio.sleep(0.02)
@pytest.mark.asyncio
@pytest.mark.parametrize("api", [True, False])
async def test_serve_starts_workers_pollers_and_shuts_down(tmp_path, use_client, api):
"""serve brings up pollers, workers and (when enabled) the HTTP API, then
tears them all down. Shutdown is driven here by cancelling the serve task,
which runs the same drain-and-close path as a SIGINT/SIGTERM."""
use_client(_mock_client())
config = _config(tmp_path)
config.ingester.api = APIConfig(enabled=api, host="127.0.0.1", port=0)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
task = asyncio.create_task(app.serve(api=api))
try:
await _wait_until(
lambda: (
app._pool is not None
and app._pool.live_workers > 0
and app._pollers is not None
and app._pollers.live_pollers > 0
)
)
finally:
task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=5.0)
# Workers and pollers are stopped after the shutdown path runs.
assert app._pool is not None and app._pollers is not None
assert app._pool.live_workers == 0
assert app._pollers.live_pollers == 0
class _SlowPool(WorkerPool):
"""A WorkerPool whose stop never finishes within the grace, to exercise
_stop_pool's timeout path. The real wiring is bypassed since _stop_pool
only calls stop() and drain_pending_releases()."""
def __init__(self) -> None:
self.released = 0
async def stop(self) -> None:
await asyncio.sleep(1.0)
async def drain_pending_releases(self, timeout: float = 2.0) -> int:
self.released += 1
return 2
@pytest.mark.asyncio
async def test_stop_pool_warns_when_shutdown_grace_elapses(tmp_path, caplog):
"""When a worker doesn't stop within the shutdown grace, _stop_pool logs a
warning and still drains any pending cancel-cleanup releases."""
config = _config(tmp_path, shutdown_grace_s=0.01)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
pool = _SlowPool()
app._pool = pool
with caplog.at_level("WARNING", logger="haiku.rag.ingester.app"):
await app._stop_pool()
assert pool.released == 1
assert "Shutdown grace" in caplog.text