diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 4c5dad19..a74b8287 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -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 diff --git a/tests/ingester/test_cli.py b/tests/ingester/test_cli.py new file mode 100644 index 00000000..06a67505 --- /dev/null +++ b/tests/ingester/test_cli.py @@ -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 diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py index 1979121b..f4126d3a 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -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