diff --git a/tests/test_research_scheduling.py b/tests/test_research_scheduling.py index d33dd44..e071bda 100644 --- a/tests/test_research_scheduling.py +++ b/tests/test_research_scheduling.py @@ -3,13 +3,14 @@ import asyncio import unittest from concurrent.futures import ThreadPoolExecutor -from typing import Any +from typing import Any, AsyncIterator from typosaurus_sandbox.research.cache import TTLCache from typosaurus_sandbox.research.client import RsearchClient from typosaurus_sandbox.research.config import ResearchConfig -from typosaurus_sandbox.research.envelopes import SearchResult +from typosaurus_sandbox.research.envelopes import ChatResponse, DescribeResponse, SearchResponse, SearchResult from typosaurus_sandbox.research.frontier import QueryFrontier +from typosaurus_sandbox.research.pipeline import ResearchPipeline, WorkItem SEARCH_FIXTURE: dict[str, Any] = { "query": "subject", @@ -45,6 +46,84 @@ DESCRIBE_FIXTURE: dict[str, Any] = { "timestamp": "2026-08-07T12:00:00Z", } +IMAGES_FIXTURE: dict[str, Any] = { + "query": "subject", + "source": "wikimedia", + "count": 1, + "success": True, + "error": None, + "results": [ + { + "title": "Aurora borealis over Norway", + "url": "https://example.com/image", + "description": "Photograph of the aurora borealis.", + "source": "wikimedia", + "extra": {}, + "index": 0, + } + ], +} + + +class _FakeClient: + def __init__(self) -> None: + self.config = ResearchConfig() + self.search_calls: list[tuple[str, dict[str, Any]]] = [] + self.chat_calls: list[str] = [] + self.describe_calls: list[str] = [] + self.delay_seconds = 0.0 + self.active = 0 + self.peak = 0 + + def search_cached( + self, + query: str, + *, + source: str | None = None, + count: int | None = None, + content: bool = False, + type: str | None = None, + deep: bool = False, + ai: bool = False, + cache: bool = True, + ) -> None: + return None + + def describe_cached(self, url: str) -> None: + return None + + async def search( + self, + query: str, + *, + source: str | None = None, + count: int | None = None, + content: bool = False, + type: str | None = None, + deep: bool = False, + ai: bool = False, + cache: bool = True, + ) -> SearchResponse: + self.search_calls.append((query, {"type": type, "content": content, "count": count})) + self.active += 1 + self.peak = max(self.peak, self.active) + try: + if self.delay_seconds: + await asyncio.sleep(self.delay_seconds) + if type == "images": + return SearchResponse.from_dict(IMAGES_FIXTURE) + return SearchResponse.from_dict(SEARCH_FIXTURE) + finally: + self.active -= 1 + + async def chat(self, prompt: str, *, system: str | None = None, json_mode: bool = False, cache: bool = True) -> ChatResponse: + self.chat_calls.append(prompt) + return ChatResponse.from_dict(CHAT_FIXTURE) + + async def describe(self, url: str) -> DescribeResponse: + self.describe_calls.append(url) + return DescribeResponse.from_dict(DESCRIBE_FIXTURE) + class TestFrontierScheduling(unittest.IsolatedAsyncioTestCase): @@ -360,6 +439,68 @@ class TestPipelineSingleMechanism(unittest.IsolatedAsyncioTestCase): self.assertEqual(describe_calls[1][0], "POST") +class TestResearchPipeline(unittest.IsolatedAsyncioTestCase): + + async def test_pipeline_pool_size_bounded_by_configured_concurrency(self) -> None: + client = _FakeClient() + frontier = QueryFrontier() + default_pipeline = ResearchPipeline(client, frontier) + self.assertEqual(default_pipeline.pool_size, 8) + self.assertEqual(default_pipeline.pool_size, ResearchConfig().max_concurrency) + narrow_pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=3)) + self.assertEqual(narrow_pipeline.pool_size, 3) + floor_pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=0)) + self.assertEqual(floor_pipeline.pool_size, 1) + + async def test_pipeline_drains_all_work_items_and_dedups(self) -> None: + client = _FakeClient() + frontier = QueryFrontier() + + async def items() -> AsyncIterator[WorkItem]: + yield WorkItem("web", "subject") + yield WorkItem("web", "subject") + yield WorkItem("images", "subject") + yield WorkItem("describe", "https://example.com/page") + yield WorkItem("chat", "prompt") + + pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=4)) + report = await pipeline.run(items()) + self.assertEqual(report.requests_succeeded, 5) + self.assertEqual(report.requests_failed, 0) + self.assertEqual(len(report.outcomes), 5) + self.assertTrue(all(outcome.success for outcome in report.outcomes)) + endpoints = {outcome.item.kind: outcome.endpoint for outcome in report.outcomes} + self.assertEqual(endpoints, {"web": "/search", "images": "/search", "describe": "/describe", "chat": "/chat"}) + self.assertEqual(len(client.search_calls), 3) + self.assertEqual(len(client.describe_calls), 1) + self.assertEqual(len(client.chat_calls), 1) + stats = frontier.snapshot() + self.assertEqual(stats.urls_seen, 2) + self.assertEqual(stats.urls_duplicates_skipped, 1) + self.assertEqual(stats.queries_enqueued, 6) + self.assertEqual(stats.queries_duplicates_skipped, 2) + self.assertEqual(stats.content_seen, 2) + + async def test_pipeline_concurrency_bounded_by_pool_size(self) -> None: + client = _FakeClient() + client.delay_seconds = 0.02 + frontier = QueryFrontier() + + async def items() -> AsyncIterator[WorkItem]: + for i in range(12): + yield WorkItem("web", f"subject {i}") + + pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=4)) + report = await pipeline.run(items()) + self.assertEqual(report.requests_succeeded, 12) + self.assertGreaterEqual(client.peak, 2) + self.assertLessEqual(client.peak, 4) + + if __name__ == "__main__": unittest.main() + + + +