PlantAI-dev/tests/unit/test_identify_blast_direction.py

from pathlib import Path

from openpyxl import Workbook

from apps.cli import run_identify as identify


def _write_config(path: Path, rows: list[list[str]]) -> None:
    wb = Workbook()
    ws = wb.active
    ws.title = "config"
    ws.append([
        "Original_File_Name",
        "Species_Code",
        "Species_Name",
        "Gene_Family",
        "File_Type",
        "Ref",
        "Species_Role",
        "File_Scope",
        "Reference_Use",
    ])
    for row in rows:
        ws.append(row)
    wb.save(path)


def test_collect_reference_files_ignores_family_pep_filename_without_config(tmp_path):
    ref = tmp_path / "At.SPL.pep.fasta"
    ref.write_text(">AtSPL1\nMAAA\n", encoding="utf-8")

    references = identify.collect_reference_files(tmp_path, gene_family="SPL")

    assert references == []


def test_collect_reference_files_uses_config_family_pep_for_identify_blast(tmp_path):
    ref = tmp_path / "At.SPL.pep.fasta"
    ref.write_text(">AtSPL1\nMAAA\n", encoding="utf-8")
    _write_config(
        tmp_path / "config.xlsx",
        [[
            "At.SPL.pep.fasta",
            "At",
            "Arabidopsis thaliana",
            "SPL",
            "family_pep",
            "TRUE",
            "seed_reference",
            "known_family",
            "identify_blast",
        ]],
    )

    references = identify.collect_reference_files(tmp_path, gene_family="SPL")

    assert len(references) == 1
    assert references[0]["species"] == "At"
    assert references[0]["source"] == "config.xlsx"
    assert Path(references[0]["path"]) == ref.resolve()


def test_parse_blast_outfmt6_uses_sseqid_for_target_candidates(tmp_path):
    blast_out = tmp_path / "Pm.blast.At.out"
    blast_out.write_text(
        "AtNAC001\tPmNAC001\t45.0\t100\t0\t0\t1\t100\t5\t104\t1e-40\t180\t120\t130\n"
        "AtNAC002\tPmNAC002\t31.0\t90\t0\t0\t1\t90\t10\t99\t1e-06\t90\t110\t120\n"
        "AtNAC003\tPmNAC003\t30.0\t90\t0\t0\t1\t90\t10\t99\t1e-20\t100\t110\t120\n"
        "AtNAC004\tPmNAC004\t55.0\t90\t0\t0\t1\t90\t10\t99\t1e-03\t100\t110\t120\n",
        encoding="utf-8",
    )

    ids, evidence = identify.parse_blast_outfmt6_target_hits(blast_out, evalue=1e-5)

    assert ids == {"PmNAC001", "PmNAC002"}
    assert "AtNAC001" not in ids
    assert evidence["PmNAC001"]["query_ids"] == ["AtNAC001"]
    assert evidence["PmNAC001"]["best_evalue"] == 1e-40


def test_multi_reference_blast_unions_target_sseqids_and_query_support(monkeypatch, tmp_path):
    output_dir = tmp_path / "identify" / "Pm"
    analysis_dir = output_dir / "analysis"
    analysis_dir.mkdir(parents=True)
    pep = tmp_path / "Pm.pep.fasta"
    pep.write_text(">PmNAC001\nMAAA\n>PmNAC002\nMBBB\n", encoding="utf-8")
    at_ref = tmp_path / "At.NAC.pep.fasta"
    os_ref = tmp_path / "Os.NAC.pep.fasta"
    at_ref.write_text(">AtNAC001\nMAAA\n", encoding="utf-8")
    os_ref.write_text(">OsNAC009\nMBBB\n", encoding="utf-8")

    monkeypatch.setattr(identify, "cleanup_current_blast_intermediates", lambda analysis_dir, prefix: None)
    monkeypatch.setattr(
        identify,
        "build_target_proteome_blast_db",
        lambda analysis_dir, pep_file: {
            "status": "success",
            "target_database_input": str(pep_file.resolve()),
            "target_database_role": "target_proteome",
            "target_db_prefix": str((analysis_dir / "blast_db.target_proteome").resolve()),
            "target_database_files": [str((analysis_dir / "blast_db.target_proteome.pin").resolve())],
        },
    )

    def fake_run_blast_analysis(**kwargs):
        label = kwargs["reference_label"]
        if label == "At":
            return {
                "status": "success",
                "reference_label": "At",
                "query_file": str(at_ref.resolve()),
                "query_role": "reference_family_pep",
                "target_database_input": str(pep.resolve()),
                "target_database_role": "target_proteome",
                "target_db_prefix": str((analysis_dir / "blast_db.target_proteome").resolve()),
                "target_database_files": [],
                "candidate_id_column": "sseqid",
                "blast_output": str((analysis_dir / "Pm.blast.At.out").resolve()),
                "ids_file": str((analysis_dir / "Pm.blast.At.ids").resolve()),
                "members": 1,
                "ids": ["PmNAC001"],
                "evidence_by_gene": {"PmNAC001": {"query_ids": ["AtNAC001"], "best_evalue": 1e-40, "best_bitscore": 180.0}},
            }
        return {
            "status": "success",
            "reference_label": "Os",
            "query_file": str(os_ref.resolve()),
            "query_role": "reference_family_pep",
            "target_database_input": str(pep.resolve()),
            "target_database_role": "target_proteome",
            "target_db_prefix": str((analysis_dir / "blast_db.target_proteome").resolve()),
            "target_database_files": [],
            "candidate_id_column": "sseqid",
            "blast_output": str((analysis_dir / "Pm.blast.Os.out").resolve()),
            "ids_file": str((analysis_dir / "Pm.blast.Os.ids").resolve()),
            "members": 2,
            "ids": ["PmNAC001", "PmNAC002"],
            "evidence_by_gene": {
                "PmNAC001": {"query_ids": ["OsNAC009"], "best_evalue": 1e-20, "best_bitscore": 140.0},
                "PmNAC002": {"query_ids": ["OsNAC010"], "best_evalue": 1e-25, "best_bitscore": 150.0},
            },
        }

    monkeypatch.setattr(identify, "run_blast_analysis", fake_run_blast_analysis)

    result = identify.run_multi_reference_blast_analysis(
        input_dir=tmp_path,
        output_dir=output_dir,
        prefix="Pm",
        references=[{"label": "At", "path": str(at_ref)}, {"label": "Os", "path": str(os_ref)}],
        pep_file=pep,
        evalue=1e-5,
    )

    assert Path(result["ids_file"]).read_text(encoding="utf-8").splitlines() == ["PmNAC001", "PmNAC002"]
    matrix = Path(result["source_matrix"]).read_text(encoding="utf-8").splitlines()
    assert matrix[0].split("\t") == [
        "gene_id",
        "blast_reference_sources",
        "blast_reference_count",
        "blast_query_ids",
        "best_evalue",
        "best_bitscore",
    ]
    assert "PmNAC001\tAt;Os\t2\tAtNAC001;OsNAC009\t1e-40\t180.0" in matrix
    assert "PmNAC002\tOs\t1\tOsNAC010\t1e-25\t150.0" in matrix
    assert result["blast_database_role"] == "target_proteome"
    assert result["blast_query_role"] == "reference_family_pep"
    assert result["blast_candidate_id_column"] == "sseqid"