PlantAI-dev/tests/unit/test_id_resolution.py

from pathlib import Path

import pytest

from gfa_core.id_resolution import (
    IdResolutionError,
    ReferenceAliasIndex,
)


LEGACY_COLUMNS = (
    "species_code",
    "protein_id",
    "pep_id",
    "cds_id",
    "gene_id",
    "cds_header",
)


def write_alias_map(
    tmp_path: Path,
    rows: list[dict[str, str]],
    *,
    include_ncbi_gene_id: bool = False,
) -> Path:
    path = tmp_path / "reference.pep_cds_idmap.tsv"
    columns = [*LEGACY_COLUMNS]
    if include_ncbi_gene_id:
        columns.append("ncbi_gene_id")
    lines = ["\t".join(columns)]
    for row in rows:
        lines.append("\t".join(row.get(column, "") for column in columns))
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return path


def legacy_row(
    *,
    species: str = "soe",
    accession: str = "XP_021867507.1",
    gene_symbol: str = "LOC110806180",
    gene_id: str = "110806180",
) -> dict[str, str]:
    return {
        "species_code": species,
        "protein_id": accession,
        "pep_id": accession,
        "cds_id": accession,
        "gene_id": gene_symbol,
        "cds_header": (
            f"{accession} [gene={gene_symbol}] "
            f"[db_xref=GeneID:{gene_id}]"
        ),
    }


def test_legacy_header_returns_versioned_accession_geneid_and_loc_for_versionless_lookup(
    tmp_path: Path,
):
    path = write_alias_map(tmp_path, [legacy_row()])

    index = ReferenceAliasIndex.from_tsv(path, species="soe")
    candidates = index.candidates("XP_021867507")

    assert [candidate.value for candidate in candidates] == [
        "XP_021867507.1",
        "110806180",
        "LOC110806180",
    ]
    assert all(candidate.namespace for candidate in candidates)
    assert all(candidate.method for candidate in candidates)
    assert all(candidate.evidence_source == str(path) for candidate in candidates)


def test_declared_species_mismatch_invalidates_file(tmp_path: Path):
    path = write_alias_map(tmp_path, [legacy_row(species="ath")])

    with pytest.raises(IdResolutionError, match="species") as exc_info:
        ReferenceAliasIndex.from_tsv(path, species="soe")

    assert exc_info.value.code == "species_mismatch"
    assert exc_info.value.details["declared_species"] == "soe"
    assert exc_info.value.details["observed_species"] == "ath"


def test_versionless_collision_with_conflicting_geneids_is_ambiguous(tmp_path: Path):
    path = write_alias_map(
        tmp_path,
        [
            legacy_row(
                accession="XP_021867507.1",
                gene_symbol="LOC110806180",
                gene_id="110806180",
            ),
            legacy_row(
                accession="XP_021867507.2",
                gene_symbol="LOC999999999",
                gene_id="999999999",
            ),
        ],
    )
    index = ReferenceAliasIndex.from_tsv(path, species="soe")

    assert index.candidates("XP_021867507.1")[0].value == "XP_021867507.1"
    assert index.candidates("XP_021867507.2")[0].value == "XP_021867507.2"
    with pytest.raises(IdResolutionError, match="ambiguous") as exc_info:
        index.candidates("XP_021867507")

    assert exc_info.value.code == "ambiguous_versionless"
    assert exc_info.value.details["raw_id"] == "XP_021867507"


def test_exact_unversioned_key_does_not_bypass_conflicting_versioned_rows(
    tmp_path: Path,
):
    path = write_alias_map(
        tmp_path,
        [
            legacy_row(
                accession="XP_021867507",
                gene_symbol="LOC110806180",
                gene_id="110806180",
            ),
            legacy_row(
                accession="XP_021867507.1",
                gene_symbol="LOC110806180",
                gene_id="110806180",
            ),
            legacy_row(
                accession="XP_021867507.2",
                gene_symbol="LOC999999999",
                gene_id="999999999",
            ),
        ],
    )
    index = ReferenceAliasIndex.from_tsv(path, species="soe")

    with pytest.raises(IdResolutionError, match="ambiguous") as exc_info:
        index.candidates("XP_021867507")

    assert exc_info.value.code == "ambiguous_versionless"
    assert exc_info.value.details["matching_accessions"] == [
        "XP_021867507",
        "XP_021867507.1",
        "XP_021867507.2",
    ]


def test_duplicate_accession_rows_with_same_geneid_deduplicate(tmp_path: Path):
    row = legacy_row()
    path = write_alias_map(tmp_path, [row, dict(row)])

    index = ReferenceAliasIndex.from_tsv(path, species="soe")

    assert [candidate.value for candidate in index.candidates("XP_021867507.1")] == [
        "XP_021867507.1",
        "110806180",
        "LOC110806180",
    ]
    assert [candidate.value for candidate in index.candidates("XP_021867507")] == [
        "XP_021867507.1",
        "110806180",
        "LOC110806180",
    ]


def test_explicit_ncbi_gene_id_column_is_consumed(tmp_path: Path):
    row = legacy_row()
    row["cds_header"] = "XP_021867507.1 [gene=LOC110806180]"
    row["ncbi_gene_id"] = "110806180"
    path = write_alias_map(tmp_path, [row], include_ncbi_gene_id=True)

    index = ReferenceAliasIndex.from_tsv(path, species="soe")

    assert [candidate.value for candidate in index.candidates("XP_021867507.1")] == [
        "XP_021867507.1",
        "110806180",
        "LOC110806180",
    ]