typosaurus-sandbox/tests/test_research_integration.py

# retoor <retoor@molodetz.nl>

import asyncio
import json
import unittest
from typing import Any, AsyncIterator
from unittest import mock

from typosaurus_sandbox.research.client import RsearchClient
from typosaurus_sandbox.research.config import ResearchConfig
from typosaurus_sandbox.research.frontier import QueryFrontier
from typosaurus_sandbox.research.pipeline import PipelineReport, ResearchPipeline, WorkItem

PROBE_SUBJECT = "python asyncio"
RSEARCH_BASE_URL = "https://rsearch.app.molodetz.nl"
RUN_TIMEOUT_SECONDS = 60.0

SEARCH_FIXTURE: dict[str, Any] = {
    "query": PROBE_SUBJECT,
    "source": "duckduckgo",
    "count": 2,
    "success": True,
    "error": None,
    "results": [
        {
            "title": "asyncio documentation",
            "url": "https://docs.python.org/3/library/asyncio.html",
            "description": "Asynchronous I/O event loop.",
            "source": "docs.python.org",
            "extra": {},
            "index": 0,
            "content": "The asyncio module provides infrastructure for writing single-threaded concurrent code.",
        },
        {
            "title": "asyncio in Python",
            "url": "https://example.com/asyncio",
            "description": "Tutorial on asyncio.",
            "source": "example.com",
            "extra": {},
            "index": 1,
            "content": "A tutorial covering the asyncio event loop and coroutines.",
        },
    ],
}


class _FakeResponse:
    def __init__(self, status: int, body: bytes) -> None:
        self.status = status
        self._body = body

    def __enter__(self) -> "_FakeResponse":
        return self

    def __exit__(self, *args: object) -> None:
        return None

    def read(self) -> bytes:
        return self._body


class TestBoundedOfflineProbe(unittest.TestCase):

    def test_bounded_probe_runs_against_mocked_transport_only(self) -> None:
        config = ResearchConfig(
            base_url=RSEARCH_BASE_URL,
            max_concurrency=2,
            default_count=2,
            request_timeout_seconds=30.0,
        )
        self.assertEqual(config.base_url, RSEARCH_BASE_URL)
        client = RsearchClient(config)
        frontier = QueryFrontier(PROBE_SUBJECT)
        requested: list[str] = []

        def fake_urlopen(request: Any, timeout: float | None = None) -> _FakeResponse:
            requested.append(request.get_full_url())
            return _FakeResponse(200, json.dumps(SEARCH_FIXTURE).encode())

        with mock.patch("urllib.request.urlopen", side_effect=fake_urlopen):
            first = asyncio.run(self._bounded_run(client, frontier))
            first_request_count = len(requested)
            second = asyncio.run(self._bounded_run(client, frontier))
            second_request_count = len(requested)

        self.assertGreaterEqual(first.requests_succeeded, 1)
        self.assertGreaterEqual(first.urls_found, 1)
        self.assertGreaterEqual(first.contents_seen, 1)
        self.assertFalse(any(outcome.cache_hit for outcome in first.outcomes))
        stats = frontier.snapshot()
        self.assertGreaterEqual(stats.urls_seen, 1)
        self.assertGreaterEqual(stats.content_seen, 1)
        self.assertGreaterEqual(first_request_count, 1)
        for url in requested:
            self.assertTrue(url.startswith(RSEARCH_BASE_URL), url)
        self.assertTrue(any("/search" in url for url in requested))
        self.assertEqual(second.requests_succeeded, 1)
        self.assertTrue(any(outcome.cache_hit for outcome in second.outcomes))
        self.assertEqual(second_request_count, first_request_count)

    async def _bounded_run(self, client: RsearchClient, frontier: QueryFrontier) -> PipelineReport:
        async def items() -> AsyncIterator[WorkItem]:
            yield WorkItem("web", PROBE_SUBJECT)

        pipeline = ResearchPipeline(client, frontier)
        return await asyncio.wait_for(pipeline.run(items()), timeout=RUN_TIMEOUT_SECONDS)


if __name__ == "__main__":
    unittest.main()