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()
 
+
+
+
+