rarelens/pipeline/tests/test_load_db.py

import json
import uuid
from pathlib import Path

import pytest
from load_db import (
    engine_for,
    load,
    parse_variant_id,
    read_vep_tab,
    sources,
    to_rows,
    vep_version,
)
from set_job_status import set_status
from sqlalchemy import text

HEADER = [
    "Uploaded_variation", "Location", "Allele", "Consequence", "IMPACT", "SYMBOL",
    "HGVSc", "HGVSp", "gnomADe_AF", "CLIN_SIG", "CADD_PHRED",
]
# A run with a VEP cache and plugins: the fixture header carries gnomADe_AF and CADD_PHRED.
CACHE_RUN = {"has_frequencies": True, "has_effect_scores": True}
# Longer than the old VARCHAR(120) column.
LONG_CLIN_SIG = (
    "conflicting_classifications_of_pathogenicity,uncertain_significance,"
    "likely_benign,benign,likely_pathogenic,pathogenic"
)


def vep_tab(tmp_path: Path, rows: list[list[str]]) -> Path:
    path = tmp_path / "x.vep.tsv"
    path.write_text("\n".join([
        "## ENSEMBL VARIANT EFFECT PREDICTOR v113.0",
        "## Output produced at 2026-09-11 12:00:00",
        "#" + "\t".join(HEADER),
        *("\t".join(r) for r in rows),
    ]) + "\n")
    return path


ROWS = [
    # SNV
    ["22_19710700_C_T", "22:19710700", "T", "missense_variant", "MODERATE", "TBX1",
     "ENST1:c.1C>T", "ENSP1:p.Arg1Trp", "0.0001", "pathogenic", "28.1"],
    # Deletion: VEP reports a trimmed allele ("-") and a shifted Location; the ID keeps the VCF truth.
    ["22_42126611_CT_C", "22:42126612", "-", "frameshift_variant", "HIGH", "CYP2D6",
     "-", "-", "-", LONG_CLIN_SIG, "-"],
]


@pytest.mark.parametrize(
    ("uid", "expected"),
    [
        ("22_19710700_C_T", ("22", 19710700, "C", "T")),
        ("chrUn_KI270742v1_100_A_AT", ("chrUn_KI270742v1", 100, "A", "AT")),
    ],
)
def test_parse_variant_id(uid: str, expected: tuple) -> None:
    assert parse_variant_id(uid) == expected


@pytest.mark.parametrize("uid", ["rs123", "12345", "22_x_A_G", "."])
def test_parse_variant_id_rejects_ids_not_set_by_normalise(uid: str) -> None:
    with pytest.raises(ValueError, match="CHROM_POS_REF_ALT"):
        parse_variant_id(uid)


def test_to_rows_takes_alleles_from_the_id_and_nulls_dashes(tmp_path: Path) -> None:
    rows = to_rows(read_vep_tab(vep_tab(tmp_path, ROWS)), job_id="j")
    snv, deletion = rows
    assert (snv["chrom"], snv["pos"], snv["ref"], snv["alt"]) == ("22", 19710700, "C", "T")
    assert (deletion["pos"], deletion["ref"], deletion["alt"]) == (42126611, "CT", "C")
    assert snv["gnomad_af"] == 0.0001 and deletion["gnomad_af"] is None
    assert deletion["hgvsc"] is None and deletion["gene"] == "CYP2D6"
    assert deletion["clinvar_sig"] == LONG_CLIN_SIG
    ann = json.loads(deletion["annotations"])
    assert "-" not in ann.values() and "CADD_PHRED" not in ann
    assert json.loads(snv["annotations"])["CADD_PHRED"] == "28.1"


def test_vep_version(tmp_path: Path) -> None:
    assert vep_version(vep_tab(tmp_path, ROWS)) == "113.0"


def test_engine_for_uses_psycopg_for_any_postgres_url() -> None:
    for url in ("postgresql+asyncpg://u:p@h/db", "postgresql://u:p@h/db"):
        assert engine_for(url).url.drivername == "postgresql+psycopg"


def new_job(engine, status: str = "running") -> str:
    job_id = str(uuid.uuid4())
    with engine.begin() as conn:
        conn.execute(text("INSERT INTO jobs (id, status) VALUES (:id, :s)"), {"id": job_id, "s": status})
    return job_id


def job_and_count(engine, job_id: str) -> tuple:
    with engine.connect() as conn:
        job = conn.execute(text("SELECT status, vep_version, log FROM jobs WHERE id=:id"),
                           {"id": job_id}).one()
        n = conn.execute(text("SELECT count(*) FROM variants WHERE job_id=:id"), {"id": job_id}).scalar()
    return (*job, n)


def test_load_is_idempotent_and_marks_job_succeeded(engine, tmp_path: Path) -> None:
    job_id = new_job(engine)
    rows = to_rows(read_vep_tab(vep_tab(tmp_path, ROWS)), job_id=job_id)
    load(engine, job_id, rows, vep="113.0", evidence=CACHE_RUN)
    # a retried task must not duplicate variants
    load(engine, job_id, rows, vep="113.0", evidence=CACHE_RUN)
    assert job_and_count(engine, job_id) == ("succeeded", "113.0", None, 2)


def test_load_with_no_variants_still_succeeds(engine, tmp_path: Path) -> None:
    job_id = new_job(engine)
    rows = to_rows(read_vep_tab(vep_tab(tmp_path, [])), job_id=job_id)
    load(engine, job_id, rows, vep="113.0", evidence=CACHE_RUN)
    assert job_and_count(engine, job_id) == ("succeeded", "113.0", None, 0)


def test_load_records_which_evidence_the_run_looked_up(engine, tmp_path: Path) -> None:
    """The API must be able to tell "absent from gnomAD" from "nobody asked gnomAD"."""
    job_id = new_job(engine)
    df = read_vep_tab(vep_tab(tmp_path, ROWS))
    load(engine, job_id, to_rows(df, job_id=job_id), vep="113.0", evidence=sources(df))
    with engine.begin() as conn:
        flags = conn.execute(
            text("SELECT has_frequencies, has_effect_scores FROM jobs WHERE id=:id"),
            {"id": job_id},
        ).one()
    assert flags == (True, True)  # this fixture has gnomADe_AF and CADD_PHRED columns


def test_a_database_mode_run_reports_no_frequencies(engine, tmp_path: Path) -> None:
    """VEP --database emits neither a frequency column nor plugin scores."""
    path = tmp_path / "db.vep.tsv"
    columns = [c for c in HEADER if c not in ("gnomADe_AF", "CADD_PHRED")]
    path.write_text("## ENSEMBL VARIANT EFFECT PREDICTOR v113.0\n#" + "\t".join(columns) + "\n")
    assert sources(read_vep_tab(path)) == {"has_frequencies": False, "has_effect_scores": False}


def test_set_status_failed_records_the_reason(engine) -> None:
    job_id = new_job(engine)
    set_status(engine, job_id, "failed", "workflow annotate-abc Failed")
    status, _, log, _ = job_and_count(engine, job_id)
    assert (status, log) == ("failed", "workflow annotate-abc Failed")


def test_set_status_never_overrides_a_succeeded_job(engine) -> None:
    job_id = new_job(engine, status="succeeded")
    set_status(engine, job_id, "failed", "late exit handler")
    assert job_and_count(engine, job_id)[0] == "succeeded"


@pytest.mark.parametrize(
    ("raw", "expected_query"),
    [
        # asyncpg's ssl= becomes libpq's sslmode= for psycopg.
        ("postgresql+asyncpg://u:p@h/db?ssl=require", {"sslmode": "require"}),
        ("postgresql://u:p@h/db?sslmode=require", {"sslmode": "require"}),
        ("postgresql+asyncpg://u:p@h/db", {}),
    ],
)
def test_engine_for_translates_ssl_options(raw: str, expected_query: dict) -> None:
    url = engine_for(raw).url
    assert url.drivername == "postgresql+psycopg"
    assert dict(url.query) == expected_query