PlantAI-dev/tests/unit/test_ppi_snapshot_preflight_task3.py

from __future__ import annotations

from pathlib import Path

from openpyxl import Workbook

from apps.web import pipeline_backend as backend
from gfa_core.gene_family_scope import build_scope_id
from gfa_core.ortholog_materialize import (
    OrthologRunRecord,
    build_ortholog_run_identity,
    materialize_ortholog_final_outputs,
)


def _job(tmp_path: Path, *, code: str = "", name: str = "", taxon_id: str = "") -> backend.JobInfo:
    input_dir = tmp_path / "input"
    workspace = tmp_path / "workspace"
    input_dir.mkdir(parents=True)
    (workspace / "final").mkdir(parents=True)
    if code:
        workbook = Workbook()
        sheet = workbook.active
        sheet.title = "config"
        sheet.append(["Species_Code", "Species_Name", "TaxID"])
        sheet.append([code, name, taxon_id])
        workbook.save(input_dir / "config.xlsx")
    return backend.JobInfo(
        job_id="job_ppi_task3",
        input_dir=str(input_dir),
        workspace_dir=str(workspace),
        logs_dir=str(workspace / "logs"),
        created_at="",
    )


def _resolve_sv(job: backend.JobInfo, **kwargs):
    return backend.resolve_ppi_species_resolution(
        job=job,
        prefix="Sv.HKL",
        ppi_species="auto",
        target_species=["Sv"],
        family="HKL",
        requested_stages=["ppi"],
        **kwargs,
    )


def _materialize_trusted_mapping(
    job: backend.JobInfo, *, prefix: str, target: str, family: str,
    storage_key: str, taxon_id: str, query_id: str, subject_id: str,
) -> None:
    workspace = Path(job.workspace_dir)
    final_dir = workspace / "final"
    query = final_dir / f"{prefix}.final_members.pep.fa"
    reference = Path(job.input_dir) / storage_key / f"{storage_key}.pep.fasta"
    query.write_text(f">{query_id}\nMA\n", encoding="utf-8")
    pairs = workspace / "ortholog" / f"{prefix}.{target}_to_{storage_key}_pairs.txt"
    pairs.parent.mkdir(parents=True, exist_ok=True)
    pairs.write_text(f"{subject_id}\t{query_id}\t91\n", encoding="utf-8")
    scope_id = build_scope_id(family, [target], [], [], "target_only")
    identity = build_ortholog_run_identity(
        query_species=target, reference_taxon_id=taxon_id,
        reference_storage_key=storage_key, query_fasta=query,
        reference_pep=reference, scope_id=scope_id, member_set_id="",
        mapping_path=pairs,
    )
    materialize_ortholog_final_outputs(
        final_dir=final_dir, prefix=prefix,
        dependency_owner="ppi",
        run_records=[OrthologRunRecord(target, storage_key, query, reference, pairs)],
        job_id=job.job_id, workspace=workspace, reference_species=storage_key,
        reference_taxon_id=taxon_id, reference_storage_key=storage_key,
        ortholog_identity=identity,
        scope_fields={
            "scope_id": scope_id, "scope_mode": "target_only",
            "target_species": [target], "reference_species": [],
            "comparative_species": [], "selected_species": [target], "family": family,
        },
    )


def test_os_exact_string_is_ready_without_reference_or_ortholog(tmp_path: Path) -> None:
    result = backend.resolve_ppi_species_resolution(
        job=_job(tmp_path),
        prefix="Os.HKL",
        ppi_species="auto",
        target_species=["Os"],
        family="HKL",
        requested_stages=["ppi"],
    )

    assert result["status"] == "ready"
    assert result["decision"] == "ready_target_string"
    assert result["required_next_stage"] == ""
    assert result["selection_mode"] == "exact_string"
    assert result["selected_taxon_id"] == "39947"
    assert result["species_resolution"]["snapshot_id"]


def test_setaria_job_identity_returns_real_top_three_and_waits_for_confirmation(tmp_path: Path) -> None:
    result = _resolve_sv(
        _job(tmp_path, code="Sv", name="Setaria viridis", taxon_id="4556")
    )

    assert result["status"] == "confirmation_required"
    assert result["decision"] == "confirmation_required"
    assert result["required_next_stage"] == ""
    assert result["selection_mode"] == "tree_nearest"
    assert result["requires_user_confirmation"] is True
    assert [row["selected_string_taxon_id"] for row in result["candidates"]] == [
        "4555",
        "4540",
        "888268",
    ]
    assert result["species_resolution"]["candidate_source"] == "string_species_tree"


def test_unknown_conflict_and_missing_snapshot_are_structured_without_at_fallback(
    tmp_path: Path, monkeypatch
) -> None:
    unknown = backend.resolve_ppi_species_resolution(
        job=_job(tmp_path / "unknown"),
        prefix="Sv.HKL",
        target_species=["Sv"],
        requested_stages=["ppi"],
    )
    assert unknown["status"] == "unresolved"
    assert unknown["issue_code"] == "target_string_tree_unresolved"
    assert unknown["selected_taxon_id"] == ""
    assert unknown["required_next_stage"] == ""

    conflict = _resolve_sv(
        _job(tmp_path / "conflict", code="Sv", name="Setaria viridis", taxon_id="4556"),
        target_scientific_name="Oryza sativa",
    )
    assert conflict["status"] == "conflict"
    assert conflict["issue_code"] == "target_scientific_name_conflict"
    assert conflict["selected_taxon_id"] == ""
    assert conflict["required_next_stage"] == ""

    monkeypatch.setattr(
        backend,
        "STRING_SPECIES_SNAPSHOT_ROOT",
        tmp_path / "missing_snapshot",
        raising=False,
    )
    missing = backend.resolve_ppi_species_resolution(
        job=_job(tmp_path / "missing"),
        prefix="Os.HKL",
        target_species=["Os"],
        requested_stages=["ppi"],
    )
    assert missing["status"] == "resource_error"
    assert missing["issue_code"] == "string_registry_unavailable"
    assert missing["selected_taxon_id"] == ""
    assert missing["required_next_stage"] == ""


def test_confirmed_nearest_without_pep_requests_only_reference_data(tmp_path: Path) -> None:
    result = _resolve_sv(
        _job(tmp_path, code="Sv", name="Setaria viridis", taxon_id="4556"),
        ppi_reference_taxon_id="4555",
        ppi_reference_storage_key="string_taxon_4555",
        ppi_reference_selection_mode="tree_nearest",
        ppi_reference_user_confirmed=True,
        dependency_owner="ppi",
    )

    assert result["status"] == "needs_reference_data"
    assert result["issue_code"] == "ppi_reference_data_missing"
    assert result["reference_bundle_status"]["required_files"] == ["pep", "cds"]
    assert result["reference_bundle_status"]["missing_required"] == ["pep", "cds"]
    assert result["required_next_stage"] == ""
    assert result["selected_taxon_id"] == "4555"


def test_confirmed_nearest_with_pep_needs_ortholog(tmp_path: Path) -> None:
    job = _job(tmp_path, code="Sv", name="Setaria viridis", taxon_id="4556")
    reference_dir = Path(job.input_dir) / "string_taxon_4555"
    reference_dir.mkdir()
    (reference_dir / "string_taxon_4555.pep.fasta").write_text(">Seita.1\nMA\n", encoding="utf-8")
    (reference_dir / "string_taxon_4555.cds.fasta").write_text(">Seita.1\nATGGCT\n", encoding="utf-8")

    result = _resolve_sv(
        job,
        ppi_reference_taxon_id="4555",
        ppi_reference_storage_key="string_taxon_4555",
        ppi_reference_selection_mode="tree_nearest",
        ppi_reference_user_confirmed=True,
        dependency_owner="ppi",
    )

    assert result["status"] == "needs_ortholog"
    assert result["decision"] == "needs_reference_ortholog"
    assert result["required_next_stage"] == "ortholog"
    assert result["reference_bundle_status"]["missing_required"] == []


def test_confirmed_nearest_with_matching_mapping_is_ready(tmp_path: Path) -> None:
    job = _job(tmp_path, code="Sv", name="Setaria viridis", taxon_id="4556")
    reference_dir = Path(job.input_dir) / "string_taxon_4555"
    reference_dir.mkdir()
    (reference_dir / "string_taxon_4555.pep.fasta").write_text(">Seita.1\nMA\n", encoding="utf-8")
    (reference_dir / "string_taxon_4555.cds.fasta").write_text(">Seita.1\nATGGCT\n", encoding="utf-8")
    _materialize_trusted_mapping(
        job, prefix="Sv.HKL", target="Sv", family="HKL",
        storage_key="string_taxon_4555", taxon_id="4555",
        query_id="SvGene1", subject_id="Seita.1",
    )

    result = _resolve_sv(
        job,
        ppi_reference_taxon_id="4555",
        ppi_reference_storage_key="string_taxon_4555",
        ppi_reference_selection_mode="tree_nearest",
        ppi_reference_user_confirmed=True,
        dependency_owner="ppi",
    )

    assert result["status"] == "ready"
    assert result["decision"] == "ready_reference_mapping"
    assert result["required_next_stage"] == ""
    assert result["mapping_reference_species"] == "string_taxon_4555"
    assert result["selected_taxon_id"] == "4555"


def test_kaks_missing_cds_does_not_override_separate_direct_ppi(tmp_path: Path) -> None:
    job = _job(tmp_path)
    reference_dir = Path(job.input_dir) / "soe"
    reference_dir.mkdir()
    (reference_dir / "soe.pep.fasta").write_text(">SoeGene1\nMA\n", encoding="utf-8")
    ppi = backend.resolve_ppi_species_resolution(
        job=job,
        prefix="Os.HKL",
        target_species=["Os"],
        requested_stages=["ppi"],
    )
    kaks = backend.resolve_ppi_species_resolution(
        job=job,
        prefix="Os.HKL",
        target_species=["Os"],
        reference_species=["Spinacia oleracea"],
        requested_stages=["kaks"],
    )

    assert ppi["status"] == "ready"
    assert ppi["required_next_stage"] == ""
    assert kaks["status"] == "needs_reference_data"
    assert kaks["issue_code"] == "kaks_reference_cds_missing"
    assert kaks["reference_bundle_status"]["missing_required"] == ["cds"]


def test_api_request_exposes_ppi_owner_confirmation_fields() -> None:
    from apps.web.vue3.main import PpiSpeciesResolutionRequest

    request = PpiSpeciesResolutionRequest(
        job_id="job_ppi_task3",
        prefix="Sv.HKL",
        ppiReferenceTaxonId="4555",
        ppiReferenceStorageKey="string_taxon_4555",
        ppiReferenceSelectionMode="tree_nearest",
        ppiReferenceUserConfirmed=True,
        dependencyOwner="ppi",
    )

    assert request.ppiReferenceTaxonId == "4555"
    assert request.ppiReferenceStorageKey == "string_taxon_4555"
    assert request.ppiReferenceSelectionMode == "tree_nearest"
    assert request.ppiReferenceUserConfirmed is True
    assert request.dependencyOwner == "ppi"