atlas-mcp/tests/test_atlas_contracts.py

import pytest

import app.atlas as atlas_module
from app.atlas import enrich_entity, resolve_entity
from app.models import AtlasAlias, AtlasEntity, AtlasIdentifier, AtlasProvenance
from app.type_classifier import TypeClassification


@pytest.mark.anyio
async def test_resolve_entity_returns_canonical_structure():
    entity = await resolve_entity("Trump")

    assert entity.atlas_id.startswith("atlas:")
    assert entity.canonical_label
    assert entity.aliases[0].label.lower() == "trump" or entity.aliases[0].label.lower() == "donald trump"
    assert entity.provenance
    assert entity.raw_payload["raw"] == "Trump"


@pytest.mark.anyio
async def test_enrich_entity_returns_dataset_shape():
    entity = await resolve_entity("Trump")
    result = enrich_entity(entity, constraints={"type": "person"}, depth=2)

    assert result.seed_entity.atlas_id == entity.atlas_id
    assert result.query_context == {"type": "person"}
    assert result.depth == 2
    assert result.related_entities == []


def test_internal_models_support_identity_and_provenance():
    entity = AtlasEntity(
        atlas_id="atlas:donald-trump",
        canonical_label="Donald Trump",
        entity_type="person",
        aliases=[AtlasAlias(label="Trump")],
        identifiers=[AtlasIdentifier(value="Q22686", source="wikidata", identifier_type="wikidata-qid")],
        provenance=[AtlasProvenance(source="google-trends", retrieval_method="entity-resolution", confidence=0.93)],
    )

    assert entity.atlas_id == "atlas:donald-trump"
    assert entity.aliases[0].label == "Trump"
    assert entity.identifiers[0].value == "Q22686"
    assert entity.provenance[0].source == "google-trends"


@pytest.mark.anyio
async def test_resolve_entity_passes_context_to_classifier(monkeypatch):
    captured = {}

    async def fake_classifier(subject, resolution, context):
        captured["context"] = context
        return TypeClassification(canonical_type="Person", provenance=None, needs_curation=False)

    def fake_trends(subject):
        return {
            "canonical_label": subject,
            "normalized": subject,
            "mid": None,
            "type": "Person",
            "source": "resolver",
            "resolved_at": "2026-04-03T00:00:00Z",
            "candidates": [],
            "raw": subject,
        }

    writes = []

    async def fake_write(entity):
        writes.append(entity)
        return {"status": "ok"}

    monkeypatch.setattr("app.atlas.classify_entity_type", fake_classifier)
    monkeypatch.setattr("app.atlas.resolve_entity_via_trends", fake_trends)
    monkeypatch.setattr(atlas_module._storage, "write_entity", fake_write)

    entity = await resolve_entity("Sample", context="news paragraph")

    assert captured["context"] == "news paragraph"
    assert entity.entity_type == "Person"
    assert writes and writes[0].canonical_label == "Sample"


@pytest.mark.anyio
async def test_resolve_entity_persists_cached_hits(monkeypatch):
    cached_entity = AtlasEntity(atlas_id="atlas:x", canonical_label="Cached Entity")
    monkeypatch.setattr("app.atlas._entity_cache.get", lambda token: cached_entity)

    writes = []

    async def fake_write(entity):
        writes.append(entity)
        return {"status": "ok"}

    monkeypatch.setattr(atlas_module._storage, "write_entity", fake_write)

    entity = await resolve_entity("Cached Entity")

    assert entity is cached_entity
    assert writes and writes[0] is cached_entity


@pytest.mark.anyio
async def test_resolve_entity_marks_needs_curation(monkeypatch):
    async def fake_classifier(subject, resolution, context):
        return TypeClassification(canonical_type=None, provenance=None, needs_curation=True)

    def fake_trends(subject):
        return {
            "canonical_label": subject,
            "normalized": subject,
            "mid": None,
            "type": "Unknown",
            "source": "resolver",
            "resolved_at": "2026-04-03T00:00:00Z",
            "candidates": [],
            "raw": subject,
        }

    monkeypatch.setattr("app.atlas.classify_entity_type", fake_classifier)
    monkeypatch.setattr("app.atlas.resolve_entity_via_trends", fake_trends)

    entity = await resolve_entity("Mysterious")

    assert entity.needs_curation is True