PlantAI-dev/tests/unit/test_wgs_align_cli.py

from __future__ import annotations

import hashlib
import json
import os
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import pytest

from apps.cli.run_wgs_align import _authority_header, build_parser, main


QC_RUN = "1" * 32
INDEX_RUN = "2" * 32


def _record(workspace: Path, path: Path, role: str, *, sample_id: str = "", component: str = "") -> dict:
    value = {"artifact": role, "role": role, "path": str(path.relative_to(workspace)),
             "full_path": str(path.resolve()), "size": path.stat().st_size,
             "size_bytes": path.stat().st_size, "sha256": hashlib.sha256(path.read_bytes()).hexdigest()}
    if sample_id: value["sample_id"] = sample_id
    if component: value["component"] = component
    return value


def _upstream(workspace: Path) -> tuple[Path, Path]:
    final = workspace / "final"; final.mkdir(parents=True, exist_ok=True)
    qc_bundle = final / ".wgs_qc_runs" / "Cohort" / QC_RUN; qc_bundle.mkdir(parents=True)
    qc_artifacts = {}
    for sample in ("S1", "S2", "S3"):
        for role, mate in (("clean_reads_r1", "R1"), ("clean_reads_r2", "R2")):
            path = qc_bundle / f"Cohort.wgs_qc.{sample}.clean.{mate}.fastq.gz"
            path.write_bytes(f"@{sample}/{mate}\nACGT\n+\nIIII\n".encode())
            qc_artifacts[f"{role}.{sample}"] = _record(workspace, path, role, sample_id=sample)
    sample_set = "wgs-samples:" + "a" * 64
    qc = {"manifest_version": "1.0", "stage": "wgs_qc", "status": "success", "prefix": "Cohort",
          "job_id": "job-1", "run_id": QC_RUN, "workspace": str(workspace.resolve()),
          "wgs_sample_set_id": sample_set, "wgs_analysis_id": "wgs-analysis:" + "b" * 64,
          "artifacts": qc_artifacts}
    qc_path = final / "Cohort.wgs_qc.manifest.json"; qc_path.write_text(json.dumps(qc))

    index_bundle = final / ".wgs_index_runs" / "Cohort" / INDEX_RUN; index_bundle.mkdir(parents=True)
    reference = index_bundle / "Cohort.wgs_index.reference.fasta"; reference.write_text(">chr1\nACGTACGT\n")
    fai = index_bundle / f"{reference.name}.fai"; fai.write_text("chr1\t8\t6\t8\t9\n")
    dictionary = index_bundle / "Cohort.wgs_index.reference.dict"; dictionary.write_text("@HD\tVN:1.6\n@SQ\tSN:chr1\tLN:8\n")
    index_artifacts = {"reference_fasta": _record(workspace, reference, "reference_fasta"),
                       "reference_fai": _record(workspace, fai, "reference_fai"),
                       "reference_dict": _record(workspace, dictionary, "reference_dict")}
    for component in (".0123", ".amb", ".ann", ".bwt.2bit.64", ".pac"):
        path = Path(f"{reference}{component}"); path.write_bytes(component.encode())
        index_artifacts[f"bwa_mem2_index_component.{component}"] = _record(
            workspace, path, "bwa_mem2_index_component", component=component,
        )
    reference_id = "wgs-reference:" + "c" * 64
    index = {"manifest_version": "1.0", "stage": "wgs_index", "status": "success", "prefix": "Cohort",
             "job_id": "job-1", "run_id": INDEX_RUN, "workspace": str(workspace.resolve()),
             "wgs_reference_id": reference_id, "wgs_analysis_id": "wgs-analysis:" + "d" * 64,
             "artifacts": index_artifacts}
    index_path = final / "Cohort.wgs_index.manifest.json"; index_path.write_text(json.dumps(index))
    return qc_path, index_path


def _tools(tmp_path: Path) -> tuple[Path, Path, Path]:
    argv_log = tmp_path / "argv.jsonl"
    bwa = tmp_path / "bwa-mem2"
    bwa.write_text("""#!/usr/bin/env python3
import json, os, sys
with open(os.environ['WGS_ARGV_LOG'], 'a') as out: out.write(json.dumps(['bwa'] + sys.argv[1:]) + '\\n')
if sys.argv[1] == 'version': print('bwa-mem2 2.2.1'); raise SystemExit(0)
if os.environ.get('WGS_FAIL_BWA'): print('forced bwa failure', file=sys.stderr); raise SystemExit(7)
rg = sys.argv[sys.argv.index('-R') + 1]
sample = dict(field.split(':', 1) for field in rg.removeprefix('@RG\\t').split('\\t'))['SM']
print('@HD\\tVN:1.6\\tSO:coordinate')
print('@SQ\\tSN:chr1\\tLN:8')
print(rg)
print(f'{sample}.read\\t0\\tchr1\\t1\\t60\\t4M\\t*\\t0\\t0\\tACGT\\tIIII')
""")
    samtools = tmp_path / "samtools"
    samtools.write_text("""#!/usr/bin/env python3
import json, os, pathlib, sys
with open(os.environ['WGS_ARGV_LOG'], 'a') as out: out.write(json.dumps(['samtools'] + sys.argv[1:]) + '\\n')
cmd = sys.argv[1]
if cmd == '--version': print('samtools 1.20'); raise SystemExit(0)
if cmd == 'sort':
    data = sys.stdin.read(); output = pathlib.Path(sys.argv[sys.argv.index('-o') + 1]); output.write_text(data)
    raise SystemExit(int(os.environ.get('WGS_FAIL_SORT', '0')))
if cmd == 'index': pathlib.Path(sys.argv[3]).write_bytes(b'BAI'); raise SystemExit(0)
if cmd == 'quickcheck': raise SystemExit(8 if os.environ.get('WGS_FAIL_QUICKCHECK') else (0 if pathlib.Path(sys.argv[-1]).stat().st_size else 1))
if cmd == 'idxstats': print('chr1\\t8\\t1\\t0'); raise SystemExit(0)
if cmd == 'view' and '-H' in sys.argv:
    print('\\n'.join(line for line in pathlib.Path(sys.argv[-1]).read_text().splitlines() if line.startswith('@')))
    raise SystemExit(0)
raise SystemExit(2)
""")
    bwa.chmod(0o755); samtools.chmod(0o755)
    return bwa, samtools, argv_log


def _run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, fail_bwa: bool = False) -> tuple[int, Path, Path]:
    workspace = tmp_path / "workspace"; qc, index = _upstream(workspace)
    bwa, samtools, argv_log = _tools(tmp_path)
    monkeypatch.setenv("WGS_ARGV_LOG", str(argv_log))
    if fail_bwa: monkeypatch.setenv("WGS_FAIL_BWA", "1")
    rc = main(["--workspace", str(workspace), "--job-id", "job-1", "--prefix", "Cohort",
               "--qc-manifest", str(qc), "--index-manifest", str(index), "--bwa-mem2", str(bwa),
               "--samtools", str(samtools), "--threads", "2", "--max-jobs", "2"])
    return rc, workspace, argv_log


def test_wgs_align_three_samples_have_exact_rg_and_validated_artifacts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    rc, workspace, argv_log = _run(tmp_path, monkeypatch)
    assert rc == 0
    manifest = json.loads((workspace / "final" / "Cohort.wgs_align.manifest.json").read_text())
    assert manifest["status"] == "success"
    assert manifest["metrics"]["input_sample_ids"] == manifest["metrics"]["output_sample_ids"] == ["S1", "S2", "S3"]
    assert manifest["wgs_sample_set_id"].startswith("wgs-samples:")
    assert manifest["wgs_reference_id"].startswith("wgs-reference:")
    assert manifest["wgs_analysis_id"].startswith("wgs-analysis:")
    assert len(manifest["artifacts"]) == 9
    assert all("/.wgs_align_runs/Cohort/" in item["full_path"] for item in manifest["artifacts"].values())
    calls = [json.loads(line) for line in argv_log.read_text().splitlines()]
    mem_calls = [call for call in calls if call[:2] == ["bwa", "mem"]]
    assert len(mem_calls) == 3
    for sample in ("S1", "S2", "S3"):
        call = next(call for call in mem_calls if f"SM:{sample}" in call[call.index("-R") + 1])
        assert call[call.index("-R") + 1] == f"@RG\tID:{sample}\tSM:{sample}\tPL:illumina"
    assert len([call for call in calls if call[:2] == ["samtools", "quickcheck"]]) == 3
    assert len([call for call in calls if call[:2] == ["samtools", "idxstats"]]) == 3


def test_failing_bwa_is_not_hidden_by_successful_sort_and_old_authority_survives(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    workspace = tmp_path / "workspace"; final = workspace / "final"; final.mkdir(parents=True)
    authority = final / "Cohort.wgs_align.manifest.json"
    old = {"stage": "wgs_align", "status": "success", "prefix": "Cohort", "run_id": "old", "artifacts": {"old": {"sha256": "f" * 64}}}
    authority.write_text(json.dumps(old)); old_hash = hashlib.sha256(authority.read_bytes()).hexdigest()
    qc, index = _upstream_existing(workspace)
    bwa, samtools, argv_log = _tools(tmp_path); monkeypatch.setenv("WGS_ARGV_LOG", str(argv_log)); monkeypatch.setenv("WGS_FAIL_BWA", "1")
    rc = main(["--workspace", str(workspace), "--job-id", "job-1", "--prefix", "Cohort", "--qc-manifest", str(qc),
               "--index-manifest", str(index), "--bwa-mem2", str(bwa), "--samtools", str(samtools)])
    assert rc == 1 and hashlib.sha256(authority.read_bytes()).hexdigest() == old_hash
    failures = list(final.glob("Cohort.wgs_align.failed.*.manifest.json"))
    assert len(failures) == 1 and json.loads(failures[0].read_text())["artifacts"] == {}
    calls = [json.loads(line) for line in argv_log.read_text().splitlines()]
    assert any(call[:2] == ["samtools", "sort"] for call in calls)


def _upstream_existing(workspace: Path) -> tuple[Path, Path]:
    # The authority may already contain an old Stage03 success; Stage01/02 are independent.
    return _upstream(workspace)


def test_parser_is_public_and_platform_defaults_to_illumina() -> None:
    parser = build_parser()
    args = parser.parse_args(["--workspace", "w", "--job-id", "j", "--prefix", "p", "--qc-manifest", "q", "--index-manifest", "i"])
    assert args.platform == "illumina"


def test_authority_header_accepts_same_job_through_mount_alias(tmp_path: Path) -> None:
    workspace = tmp_path / "current-project" / "workspace" / "job-1"
    manifest_path = workspace / "final" / "Cohort.wgs_qc.manifest.json"
    manifest_path.parent.mkdir(parents=True)
    manifest_path.write_text(json.dumps({
        "stage": "wgs_qc",
        "status": "success",
        "prefix": "Cohort",
        "job_id": "job-1",
        "workspace": "/legacy-mount/project/workspace/job-1",
    }))

    manifest = _authority_header(
        manifest_path,
        stage="wgs_qc",
        workspace=workspace,
        prefix="Cohort",
        job_id="job-1",
    )

    assert manifest["workspace"] == "/legacy-mount/project/workspace/job-1"


@pytest.mark.parametrize(("environment", "expected"), [("WGS_FAIL_SORT", "samtools sort failed"),
                                                          ("WGS_FAIL_QUICKCHECK", "quickcheck failed")])
def test_sort_and_postflight_failures_publish_no_artifacts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
                                                           environment: str, expected: str, capsys) -> None:
    monkeypatch.setenv(environment, "9")
    rc, workspace, _ = _run(tmp_path, monkeypatch)
    assert rc == 1 and expected in capsys.readouterr().err
    manifest = json.loads((workspace / "final" / "Cohort.wgs_align.manifest.json").read_text())
    assert manifest["status"] == "failed" and manifest["artifacts"] == {}


def test_concurrent_same_prefix_attempts_leave_one_authoritative_run_identity(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    workspace = tmp_path / "workspace"; qc, index = _upstream(workspace)
    bwa, samtools, argv_log = _tools(tmp_path); monkeypatch.setenv("WGS_ARGV_LOG", str(argv_log))
    argv = ["--workspace", str(workspace), "--job-id", "job-1", "--prefix", "Cohort", "--qc-manifest", str(qc),
            "--index-manifest", str(index), "--bwa-mem2", str(bwa), "--samtools", str(samtools), "--max-jobs", "2"]
    with ThreadPoolExecutor(max_workers=2) as pool:
        results = [future.result() for future in (pool.submit(main, argv), pool.submit(main, argv))]
    assert 0 in results
    manifest = json.loads((workspace / "final" / "Cohort.wgs_align.manifest.json").read_text())
    assert manifest["status"] == "success"
    run_roots = {Path(item["full_path"]).parent for item in manifest["artifacts"].values()}
    assert run_roots == {workspace / "final" / ".wgs_align_runs" / "Cohort" / manifest["run_id"]}