"""Tests for the generic ``ChunkedMapper`` primitive. The mapper is the per-chunk fan-out machinery extracted from ``ChunkedReasoner``: char-budgeted slicing, parallel scheduling under a semaphore, time-bounded extraction with cancellation, progress events, and worker-failure tolerance. These tests drive it with a stubbed ``Agent[None, T]`` so the model boundary stays patched out. """ from __future__ import annotations import asyncio from dataclasses import dataclass from unittest.mock import AsyncMock, patch import pytest from pydantic import BaseModel from pydantic_ai import Agent from stirling.agents.shared.chunked_mapper import ChunkedMapper from stirling.contracts.documents import Page from stirling.services.runtime import AppRuntime @dataclass class _StubAgentResult[T]: output: T class _Extracted(BaseModel): """Tiny per-chunk extractor payload used by these tests.""" label: str def _page(n: int, text: str) -> Page: return Page(page_number=n, text=text, char_count=len(text)) def _build_mapper( runtime: AppRuntime, *, chars_per_slice: int | None = None, concurrency: int | None = None, worker_timeout_seconds: float | None = None, ) -> ChunkedMapper[_Extracted]: """Build a mapper wrapping a real ``Agent`` whose ``.run`` is patched per test.""" extractor: Agent[None, _Extracted] = Agent( model=runtime.fast_model, output_type=_Extracted, model_settings=runtime.fast_model_settings, ) return ChunkedMapper( runtime, extractor=extractor, chars_per_slice=chars_per_slice, concurrency=concurrency, worker_timeout_seconds=worker_timeout_seconds, ) class TestSlicePages: """The static helper is pure: no I/O, no scheduling.""" def test_single_slice_when_under_budget(self) -> None: pages = [_page(1, "abc"), _page(2, "def"), _page(3, "gh")] slices = ChunkedMapper.slice_pages(pages, chars_per_slice=20) assert [[p.page_number for p in s] for s in slices] == [[1, 2, 3]] def test_starts_new_slice_when_budget_exceeded(self) -> None: pages = [_page(1, "a" * 6), _page(2, "b" * 6), _page(3, "c" * 6)] slices = ChunkedMapper.slice_pages(pages, chars_per_slice=10) # 6 + 6 > 10 → break after each page assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]] def test_oversized_page_is_its_own_slice(self) -> None: """Page boundaries are never broken: an oversize page becomes its own slice.""" pages = [_page(1, "small"), _page(2, "x" * 100), _page(3, "tiny")] slices = ChunkedMapper.slice_pages(pages, chars_per_slice=10) assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]] def test_rejects_non_positive_budget(self) -> None: with pytest.raises(ValueError, match="chars_per_slice"): ChunkedMapper.slice_pages([_page(1, "x")], chars_per_slice=0) class TestFormatChunkContent: def test_renders_page_markers(self) -> None: rendered = ChunkedMapper.format_chunk_content([_page(2, "two"), _page(3, "three")]) assert "[Page 2]\ntwo" in rendered assert "[Page 3]\nthree" in rendered # Blank-line separator between pages assert "two\n\n[Page 3]" in rendered class TestMapPages: @pytest.mark.anyio async def test_single_chunk_returns_single_output(self, runtime: AppRuntime) -> None: mapper = _build_mapper(runtime, chars_per_slice=1000) pages = [_page(1, "alpha"), _page(2, "beta"), _page(3, "gamma")] canned = _Extracted(label="one") with patch.object( mapper._extractor, "run", AsyncMock(return_value=_StubAgentResult(output=canned)), ) as run_mock: outputs = await mapper.map_pages(pages, "what") assert run_mock.await_count == 1 assert len(outputs) == 1 assert outputs[0].pages == [1, 2, 3] assert outputs[0].output == canned assert outputs[0].label == "pages=1-3" @pytest.mark.anyio async def test_multi_chunk_outputs_are_in_document_order(self, runtime: AppRuntime) -> None: """Outputs are sorted by first covered page regardless of completion order.""" mapper = _build_mapper(runtime, chars_per_slice=10, concurrency=3) pages = [_page(i, "x" * 8) for i in range(1, 4)] # Each chunk's worker awaits a release event; we release in reverse # order so completion order is the inverse of slice order. release = [asyncio.Event() for _ in pages] call_index = 0 async def _gated(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]: nonlocal call_index mine = call_index call_index += 1 await release[mine].wait() return _StubAgentResult(output=_Extracted(label=f"slice-{mine + 1}")) async def _release_in_reverse() -> None: await asyncio.sleep(0) for ev in reversed(release): ev.set() await asyncio.sleep(0) await asyncio.sleep(0) with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_gated)): task = asyncio.create_task(mapper.map_pages(pages, "anything")) await _release_in_reverse() outputs = await task assert [o.pages for o in outputs] == [[1], [2], [3]] @pytest.mark.anyio async def test_worker_failure_drops_only_that_chunk(self, runtime: AppRuntime) -> None: mapper = _build_mapper(runtime, chars_per_slice=10) pages = [_page(i, "x" * 8) for i in range(1, 4)] results: list[_Extracted | BaseException] = [ _Extracted(label="a"), RuntimeError("boom"), _Extracted(label="c"), ] async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]: value = results.pop(0) if isinstance(value, BaseException): raise value return _StubAgentResult(output=value) with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_stub)): outputs = await mapper.map_pages(pages, "anything") assert len(outputs) == 2 assert {o.output.label for o in outputs} == {"a", "c"} @pytest.mark.anyio async def test_worker_timeout_drops_only_that_chunk(self, runtime: AppRuntime) -> None: mapper = _build_mapper(runtime, chars_per_slice=10, worker_timeout_seconds=0.05) pages = [_page(i, "x" * 8) for i in range(1, 4)] async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]: # Page 2 hangs forever; pages 1 and 3 return immediately. prompt = _args[0] assert isinstance(prompt, str) if "[Page 2]" in prompt: await asyncio.sleep(10) return _StubAgentResult(output=_Extracted(label="ok")) with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_stub)): outputs = await mapper.map_pages(pages, "anything") covered = sorted({p for o in outputs for p in o.pages}) assert covered == [1, 3] @pytest.mark.anyio async def test_outer_cancellation_drains_pending_tasks(self, runtime: AppRuntime) -> None: """Cancellation propagating in from upstream cancels per-chunk model calls rather than letting them keep billing tokens.""" mapper = _build_mapper(runtime, chars_per_slice=10, concurrency=5) pages = [_page(i, "x" * 8) for i in range(1, 5)] cancellations = 0 async def _hang(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]: nonlocal cancellations try: await asyncio.sleep(60) except asyncio.CancelledError: cancellations += 1 raise return _StubAgentResult(output=_Extracted(label="never")) with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_hang)): task = asyncio.create_task(mapper.map_pages(pages, "anything")) # Yield once so all four workers are blocked on their sleep. await asyncio.sleep(0) await asyncio.sleep(0) task.cancel() with pytest.raises(asyncio.CancelledError): await task assert cancellations == len(pages) @pytest.mark.anyio async def test_semaphore_caps_concurrency(self, runtime: AppRuntime) -> None: """At most ``concurrency`` workers run at once; with strictly more work items than slots the observed max is exactly the configured cap.""" concurrency = 2 mapper = _build_mapper(runtime, chars_per_slice=10, concurrency=concurrency) pages = [_page(i, "x" * 8) for i in range(1, 6)] # 5 items > 2 slots active = 0 peak = 0 async def _track(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]: nonlocal active, peak active += 1 peak = max(peak, active) # Yield enough times that other waiters get a chance to enter. for _ in range(5): await asyncio.sleep(0) active -= 1 return _StubAgentResult(output=_Extracted(label="ok")) with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_track)): outputs = await mapper.map_pages(pages, "anything") assert len(outputs) == 5 assert peak == concurrency @pytest.mark.anyio async def test_rejects_empty_pages(self, runtime: AppRuntime) -> None: mapper = _build_mapper(runtime) with pytest.raises(ValueError, match="at least one page"): await mapper.map_pages([], "anything") class TestSummaryCounts: """``summary_counts`` callback feeds the WholeDocSliceDone event's excerpts/facts counters from the consumer's extractor output shape without the mapper itself duck-typing fields on ``T``.""" @pytest.mark.anyio async def test_default_callback_emits_zero_counts(self, runtime: AppRuntime) -> None: """No callback supplied → events emit ``excerpts=0 facts=0``.""" from stirling.contracts import WholeDocSliceDone from stirling.services import reset_progress_emitter, set_progress_emitter mapper = _build_mapper(runtime, chars_per_slice=1000) pages = [_page(1, "small")] canned = _Extracted(label="ok") emitted: list[WholeDocSliceDone] = [] async def _emit(event: object) -> None: if isinstance(event, WholeDocSliceDone): emitted.append(event) token = set_progress_emitter(_emit) try: with patch.object( mapper._extractor, "run", AsyncMock(return_value=_StubAgentResult(output=canned)), ): await mapper.map_pages(pages, "q") finally: reset_progress_emitter(token) assert len(emitted) == 1 assert emitted[0].excerpts == 0 assert emitted[0].facts == 0 @pytest.mark.anyio async def test_user_callback_drives_counts(self, runtime: AppRuntime) -> None: """A supplied callback receives each chunk's typed output and its returned tuple is what the event carries.""" from stirling.contracts import WholeDocSliceDone from stirling.services import reset_progress_emitter, set_progress_emitter captured: list[_Extracted] = [] def _counts(output: _Extracted) -> tuple[int, int]: captured.append(output) return (3, 7) extractor: Agent[None, _Extracted] = Agent( model=runtime.fast_model, output_type=_Extracted, model_settings=runtime.fast_model_settings, ) mapper: ChunkedMapper[_Extracted] = ChunkedMapper( runtime, extractor=extractor, chars_per_slice=1000, summary_counts=_counts, ) canned = _Extracted(label="ok") emitted: list[WholeDocSliceDone] = [] async def _emit(event: object) -> None: if isinstance(event, WholeDocSliceDone): emitted.append(event) token = set_progress_emitter(_emit) try: with patch.object( mapper._extractor, "run", AsyncMock(return_value=_StubAgentResult(output=canned)), ): await mapper.map_pages([_page(1, "small")], "q") finally: reset_progress_emitter(token) assert len(captured) == 1 assert captured[0].label == "ok" assert emitted[0].excerpts == 3 assert emitted[0].facts == 7 class TestChunkOutputShape: @pytest.mark.anyio async def test_single_page_label(self, runtime: AppRuntime) -> None: mapper = _build_mapper(runtime, chars_per_slice=5) pages = [_page(7, "x" * 6)] # one oversize page → one slice canned = _Extracted(label="solo") with patch.object( mapper._extractor, "run", AsyncMock(return_value=_StubAgentResult(output=canned)), ): outputs = await mapper.map_pages(pages, "q") assert outputs[0].label == "pages=7"