atlas2-mcp/tests/test_maintenance_helpers.py

import rdflib

from app.maintenance import (
    _extract_wikidata_qids_from_entity_dump,
    _infer_atlas_type_from_qids,
)


def test_extract_wikidata_qids_from_dump_dedupes():
    dump = {
        "entities": {
            "Q1": {
                "claims": {
                    "P31": [
                        {
                            "mainsnak": {
                                "datavalue": {"value": {"id": "Q5"}},
                            }
                        },
                        {
                            "mainsnak": {
                                "datavalue": {"value": {"id": "Q5"}},
                            }
                        },
                    ],
                    "P279": [
                        {
                            "mainsnak": {
                                "datavalue": {"value": {"id": "Q6"}},
                            }
                        }
                    ],
                },
            }
        }
    }

    qids = _extract_wikidata_qids_from_entity_dump(dump)
    assert qids == ["Q5", "Q6"]


def test_infer_atlas_type_from_qids_heuristic_person():
    g = rdflib.Graph()
    wikidata_uri = rdflib.URIRef("http://www.wikidata.org/entity/Q5")
    # object should contain a keyword like "Person" so the heuristic maps.
    g.add((wikidata_uri, rdflib.RDFS.subClassOf, rdflib.URIRef("http://dbpedia.org/ontology/Person")))

    atlas_type = _infer_atlas_type_from_qids(["Q5"], g)
    assert atlas_type == "atlas:Person"