205 lines
11 KiB
Python
205 lines
11 KiB
Python
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"]}
|