spb/ai-atlas

1"""Shared test helpers. Connector tests never hit the network: they replay fixtures through `extract_from_fixture`."""2from __future__ import annotations34from datetime import UTC, datetime5from pathlib import Path67import pytest89from aiatlas.sdk.connector import BaseConnector, RunContext10from aiatlas.sdk.facts import Facts, Target11from aiatlas.sdk.fetch import file_result1213FIXTURES = Path(__file__).parent / "fixtures"141516def fixture_path(*parts: str) -> Path:17    return FIXTURES.joinpath(*parts)181920async def extract_from_fixture(connector: BaseConnector, target: Target, path: Path, *, content_type: str = "text/html") -> Facts:21    """Parse a saved response and run the connector's extract() — no database, no network."""22    ctx = RunContext(run_id="test", connector=connector, fetcher=None, started_at=datetime.now(UTC))  # type: ignore[arg-type]23    res = file_result(str(path), url=target.url, content_type=content_type)24    parsed = connector.parse(target, res)25    facts = await connector.extract(ctx, target, res, parsed)26    return facts or Facts()272829@pytest.fixture30def fixtures() -> Path:31    return FIXTURES323334def claims_of(facts: Facts, entity_name: str) -> dict[str, object]:35    """{property: value} for one entity (from Facts.claims and EntityRef.attributes)."""36    out: dict[str, object] = {}37    for e in facts.entities:38        if e.name == entity_name:39            out.update(e.attributes)40    for c in facts.claims:41        if c.entity.name == entity_name:42            out[c.property] = c.value43    return out444546def entity_names(facts: Facts, entity_type: str) -> set[str]:47    return {e.name for e in facts.entities if e.entity_type == entity_type}48