PlantAI-dev/tests/unit/test_gwas_mlm_frontend_handoff.py

import json
import subprocess
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[2]
FRONTEND_PATH = PROJECT_ROOT / "apps/web/vue3/index.html"
GENOTYPE_SCOPE_ID = "genotype:" + "1" * 64
PHENOTYPE_SCOPE_ID = "phenotype:" + "2" * 64
COVARIATE_SCOPE_ID = "covariate:" + "3" * 64


def _run_handoff(
    tmp_path: Path,
    producer_results: list[list],
    *,
    mlm_status: str = "pending",
    initial_params: dict | None = None,
) -> dict:
    tmp_path.mkdir(parents=True, exist_ok=True)
    frontend = FRONTEND_PATH.read_text(encoding="utf-8")
    handoff_start = frontend.index("const handoffGwasMlmProducerIdentity =")
    handoff_end = frontend.index("const handoffGwasEmmaxProducerIdentity =", handoff_start)
    stages = [{"id": "gwas_mlm", "params": initial_params or {}}]
    script = tmp_path / "mlm-handoff.cjs"
    script.write_text(
        "const batchExecuteStages={value:" + json.dumps(stages) + "};"
        "const stageExecutionStatus={value:{gwas_mlm:" + json.dumps(mlm_status) + "}};"
        "const getStageId=s=>s.id;const executionPrefix=()=> 'Os';"
        + frontend[handoff_start:handoff_end]
        + "const rows=" + json.dumps(producer_results) + ";"
        + "const accepted=rows.map(([stage,result])=>handoffGwasMlmProducerIdentity(stage,result));"
        + "process.stdout.write(JSON.stringify({accepted,params:batchExecuteStages.value[0].params}));",
        encoding="utf-8",
    )
    completed = subprocess.run(["node", str(script)], text=True, capture_output=True, check=False)
    assert completed.returncode == 0, completed.stderr
    return json.loads(completed.stdout)


def _producer_results() -> list[list]:
    return [
        ["gwas_fmt", {"result": {"stage_status": {
            "completed": True, "stage": "gwas_fmt", "prefix": "Os",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
        }}}],
        ["gwas_kinship", {"already_completed": True, "stage_status": {
            "completed": True, "stage": "gwas_kinship", "prefix": "Os",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
        }}],
        ["gwas_trait", {
            "stage_status": {"completed": True, "stage": "gwas_trait", "prefix": "Os"},
            "stage_identity": {
                "stage": "gwas_trait", "prefix": "Os",
                "phenotype_scope_id": PHENOTYPE_SCOPE_ID,
            },
        }],
        ["gwas_cov", {"status": {"identity": {
            "completed": True, "stage": "gwas_cov", "prefix": "Os",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
            "covariate_scope_id": COVARIATE_SCOPE_ID,
        }}}],
    ]


def test_frontend_handoffs_all_completed_producer_scopes_to_pending_mlm(tmp_path: Path) -> None:
    result = _run_handoff(tmp_path, _producer_results())

    assert result["accepted"] == [True, True, True, True]
    assert result["params"] == {
        "genotypeScopeId": GENOTYPE_SCOPE_ID,
        "phenotypeScopeId": PHENOTYPE_SCOPE_ID,
        "covariateScopeId": COVARIATE_SCOPE_ID,
    }


def test_frontend_rejects_invalid_or_irrelevant_mlm_handoffs(tmp_path: Path) -> None:
    cases = [
        ["gwas_fmt", {"stage_status": {
            "completed": False, "stage": "gwas_fmt", "prefix": "Os",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
        }}],
        ["gwas_fmt", {"stage_status": {
            "completed": True, "stage": "gwas_trait", "prefix": "Os",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
        }}],
        ["gwas_fmt", {"stage_status": {
            "completed": True, "stage": "gwas_fmt", "prefix": "At",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
        }}],
        ["gwas_fmt", {"stage_status": {
            "completed": True, "stage": "gwas_fmt", "prefix": "Os",
            "genotype_scope_id": "invalid",
        }}],
        ["gwas_emmax", {"stage_status": {
            "completed": True, "stage": "gwas_emmax", "prefix": "Os",
            "genotype_scope_id": GENOTYPE_SCOPE_ID,
        }}],
    ]

    for index, row in enumerate(cases):
        result = _run_handoff(tmp_path / str(index), [row])
        assert result == {"accepted": [False], "params": {}}

    completed_consumer = _run_handoff(
        tmp_path / "completed",
        [_producer_results()[0]],
        mlm_status="completed",
    )
    assert completed_consumer == {"accepted": [False], "params": {}}


def test_frontend_rejects_conflicting_cov_identity_atomically(tmp_path: Path) -> None:
    existing = "genotype:" + "4" * 64
    cov_result = _producer_results()[-1]

    result = _run_handoff(
        tmp_path,
        [cov_result],
        initial_params={"genotypeScopeId": existing},
    )

    assert result["accepted"] == [False]
    assert result["params"] == {"genotypeScopeId": existing}


def _apply_mlm_request(tmp_path: Path, stage_params: dict) -> dict:
    frontend = FRONTEND_PATH.read_text(encoding="utf-8")
    apply_start = frontend.index("const applyGwasExecutionParams =")
    apply_end = frontend.index("const executeSingleStageWithParams =", apply_start)
    script = tmp_path / "mlm-request.cjs"
    script.write_text(
        "const batchExecuteParams={value:{plan:{params:{}}}};const confirmParams={};"
        + frontend[apply_start:apply_end]
        + "const body={};applyGwasExecutionParams(body,'gwas_mlm',"
        + json.dumps(stage_params)
        + ");process.stdout.write(JSON.stringify(body));",
        encoding="utf-8",
    )
    completed = subprocess.run(["node", str(script)], text=True, capture_output=True, check=False)
    assert completed.returncode == 0, completed.stderr
    return json.loads(completed.stdout)


def test_mlm_request_binds_only_scopes_and_tassel_runtime_params(tmp_path: Path) -> None:
    body = _apply_mlm_request(tmp_path, {
        "genotypeScopeId": GENOTYPE_SCOPE_ID,
        "phenotypeScopeId": PHENOTYPE_SCOPE_ID,
        "covariateScopeId": COVARIATE_SCOPE_ID,
        "tasselPipeline": "/opt/tassel/run_pipeline.pl",
        "tasselXmx": "20g",
        "gwasSortedVcf": "/forbidden/sorted.vcf.gz",
        "gwasKinshipFile": "/forbidden/kinship.txt",
        "gwasTraitFile": "/forbidden/trait.txt",
        "gwasCovariateFile": "/forbidden/cov.txt",
        "gwasVcfManifest": "/forbidden/fmt.manifest.json",
    })

    assert body == {
        "genotypeScopeId": GENOTYPE_SCOPE_ID,
        "phenotypeScopeId": PHENOTYPE_SCOPE_ID,
        "covariateScopeId": COVARIATE_SCOPE_ID,
        "tasselPipeline": "/opt/tassel/run_pipeline.pl",
        "tasselXmx": "20g",
    }


def test_mlm_request_accepts_snake_case_scope_params(tmp_path: Path) -> None:
    body = _apply_mlm_request(tmp_path, {
        "genotype_scope_id": GENOTYPE_SCOPE_ID,
        "phenotype_scope_id": PHENOTYPE_SCOPE_ID,
        "covariate_scope_id": COVARIATE_SCOPE_ID,
        "tassel_pipeline": "/opt/tassel/run_pipeline.pl",
        "tassel_xmx": "12g",
    })

    assert body == {
        "genotypeScopeId": GENOTYPE_SCOPE_ID,
        "phenotypeScopeId": PHENOTYPE_SCOPE_ID,
        "covariateScopeId": COVARIATE_SCOPE_ID,
        "tasselPipeline": "/opt/tassel/run_pipeline.pl",
        "tasselXmx": "12g",
    }