PlantAI-dev/tests/unit/test_wgs_qc_cli.py

import gzip
import json
import os
import stat
import subprocess
import sys
from pathlib import Path

import yaml
import pytest

from apps.cli.run_wgs_qc import _publish_transaction, build_parser, main


ROOT = Path(__file__).resolve().parents[2]


def _fastq(path: Path) -> None:
    with gzip.open(path, "wt", encoding="utf-8") as handle:
        handle.write("@read\nACGT\n+\nIIII\n")


def _fake_fastp(path: Path, *, fail_sample: str = "") -> Path:
    path.write_text(
        "#!/usr/bin/env python3\n"
        "import pathlib, shutil, sys\n"
        "if '--version' in sys.argv:\n"
        "    print('fastp 0.23.2'); raise SystemExit(0)\n"
        "a = sys.argv\n"
        "def value(flag): return pathlib.Path(a[a.index(flag) + 1])\n"
        f"if {fail_sample!r} and {fail_sample!r} in str(value('-i')):\n"
        "    raise SystemExit(17)\n"
        "shutil.copyfile(value('-i'), value('-o'))\n"
        "shutil.copyfile(value('-I'), value('-O'))\n"
        "value('-h').write_text('<html>ok</html>', encoding='utf-8')\n"
        "value('-j').write_text('{\\\"summary\\\": {}}', encoding='utf-8')\n",
        encoding="utf-8",
    )
    path.chmod(path.stat().st_mode | stat.S_IXUSR)
    return path


def _samples(tmp_path: Path, ids: tuple[str, ...] = ("S01", "S02", "S03")) -> Path:
    samples = {}
    for sample_id in ids:
        r1, r2 = tmp_path / f"{sample_id}.R1.fastq.gz", tmp_path / f"{sample_id}.R2.fastq.gz"
        _fastq(r1)
        _fastq(r2)
        samples[sample_id] = {"r1": r1.name, "r2": r2.name}
    config = tmp_path / "wgs_samples.yaml"
    config.write_text(yaml.safe_dump({"samples": samples}), encoding="utf-8")
    return config


def _run(workspace: Path, config: Path, fastp: Path) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [
            sys.executable, "-m", "apps.cli.run_wgs_qc", "--workspace", str(workspace),
            "--job-id", "qc-job", "--prefix", "Cohort", "--wgs-samples", str(config),
            "--fastp", str(fastp), "--threads", "2", "--max-jobs", "2",
        ],
        cwd=ROOT,
        text=True,
        capture_output=True,
        check=False,
    )


def test_wgs_qc_publishes_every_declared_sample_with_four_hashed_roles(tmp_path: Path) -> None:
    config = _samples(tmp_path)
    result = _run(tmp_path / "workspace", config, _fake_fastp(tmp_path / "fake_fastp"))

    assert result.returncode == 0, result.stderr
    manifest_path = tmp_path / "workspace" / "final" / "Cohort.wgs_qc.manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    assert manifest["status"] == "success"
    assert manifest["metrics"]["input_sample_ids"] == ["S01", "S02", "S03"]
    assert manifest["metrics"]["output_sample_ids"] == ["S01", "S02", "S03"]
    assert manifest["wgs_sample_set_id"].startswith("wgs-samples:")
    assert manifest["wgs_analysis_id"].startswith("wgs-analysis:")
    run_root = tmp_path / "workspace" / "final" / ".wgs_qc_runs" / "Cohort" / manifest["run_id"]
    assert run_root.is_dir()
    for sample_id in ("S01", "S02", "S03"):
        roles = [item for item in manifest["artifacts"].values() if item.get("sample_id") == sample_id]
        assert {item["role"] for item in roles} == {"clean_reads_r1", "clean_reads_r2", "fastp_html", "fastp_json"}
        for item in roles:
            output = tmp_path / "workspace" / item["path"]
            assert output.is_file() and output.stat().st_size > 0
            assert len(item["sha256"]) == 64
    assert manifest["artifacts"]["wgs_sample_config"]["role"] == "wgs_sample_config"


def test_wgs_qc_failure_writes_empty_fresh_manifest_without_publishing_old_final(tmp_path: Path) -> None:
    config = _samples(tmp_path)
    workspace = tmp_path / "workspace"
    old = workspace / "final" / "old.clean.R1.fastq.gz"
    old.parent.mkdir(parents=True)
    old.write_bytes(b"old-result")

    result = _run(workspace, config, _fake_fastp(tmp_path / "fake_fastp", fail_sample="S02"))

    assert result.returncode == 1
    manifest = json.loads((workspace / "final" / "Cohort.wgs_qc.manifest.json").read_text(encoding="utf-8"))
    assert manifest["status"] == "failed"
    assert manifest["artifacts"] == {}
    assert old.read_bytes() == b"old-result"
    assert not list((workspace / "final").glob("Cohort.wgs_qc.*.clean.*"))


def test_wgs_qc_publish_failure_keeps_old_authority_and_orphans_new_bundle(monkeypatch, tmp_path: Path) -> None:
    final = tmp_path / "final"; final.mkdir()
    manifest = final / "Cohort.wgs_qc.manifest.json"
    manifest.write_text(json.dumps({"status": "success", "start_time": "2000-01-01T00:00:00+00:00"}), encoding="utf-8")
    bundle_work = tmp_path / "bundle"; bundle_work.mkdir(); (bundle_work / "new.txt").write_text("new", encoding="utf-8")
    bundle_destination = final / ".wgs_qc_runs" / "Cohort" / "run"
    real_replace = os.replace
    calls = {"count": 0}
    def failing_replace(source, destination):
        calls["count"] += 1
        if calls["count"] == 2: raise OSError("injected manifest-switch failure")
        return real_replace(source, destination)
    monkeypatch.setattr(os, "replace", failing_replace)
    staged_manifest = tmp_path / "manifest.json"; staged_manifest.write_text('{"status":"success"}', encoding="utf-8")
    with pytest.raises(OSError):
        _publish_transaction(bundle_work, bundle_destination, staged_manifest, manifest, "run", "2026-01-01T00:00:00+00:00")
    assert json.loads(manifest.read_text(encoding="utf-8"))["status"] == "success"
    assert (bundle_destination / "new.txt").read_text(encoding="utf-8") == "new"


def test_wgs_qc_public_parser_and_main_argv(tmp_path: Path) -> None:
    parser = build_parser()
    assert parser.parse_args(["--workspace", "w", "--job-id", "j", "--prefix", "p", "--wgs-samples", "s"]).prefix == "p"
    assert main(["--workspace", str(tmp_path), "--job-id", "j", "--prefix", "p", "--wgs-samples", "missing.yaml"]) == 1