Files
PlantAI-dev/tests/unit/test_wgs_index_cli.py
2026-08-03 18:42:47 +08:00

309 lines
14 KiB
Python

import json
import os
import re
import shutil
import stat
import subprocess
import sys
from pathlib import Path
import pytest
from apps.cli.run_wgs_index import build_parser, main
COMPONENTS = (".0123", ".amb", ".ann", ".bwt.2bit.64", ".pac")
def _tool(path: Path, kind: str) -> Path:
path.write_text(
"#!/usr/bin/env python3\n"
"import os, sys\n"
f"KIND = {kind!r}\n"
"args = sys.argv[1:]\n"
"if '--version' in args or args == ['version']:\n"
" print(KIND + ' 1.0'); raise SystemExit(0)\n"
"if os.environ.get('WGS_INDEX_FAIL') == KIND:\n"
" print('forced failure', file=sys.stderr); raise SystemExit(9)\n"
"if KIND == 'samtools':\n"
" ref = args[args.index('faidx') + 1]\n"
" content = 'chr1\\t9\\t6\\t9\\t10\\n' if os.environ.get('WGS_INDEX_CORRUPT') == 'fai' else 'chr1\\t8\\t6\\t8\\t9\\n'\n"
" open(ref + '.fai', 'w').write(content)\n"
"elif KIND == 'gatk':\n"
" out = args[args.index('-O') + 1]\n"
" content = '@HD\\tVN:1.6\\n@SQ\\tSN:chr1\\tLN:9\\n' if os.environ.get('WGS_INDEX_CORRUPT') == 'dict' else '@HD\\tVN:1.6\\n@SQ\\tSN:chr1\\tLN:8\\n'\n"
" open(out, 'w').write(content)\n"
"else:\n"
" if '-t' in args: print('index: invalid option -- t', file=sys.stderr); raise SystemExit(11)\n"
" ref = args[-1]\n"
" for suffix in ('.0123', '.amb', '.ann', '.bwt.2bit.64', '.pac'):\n"
" open(ref + suffix, 'w').write(suffix + '\\n')\n",
encoding="utf-8",
)
path.chmod(path.stat().st_mode | stat.S_IXUSR)
return path
def _run(
workspace: Path, reference: Path, tools: dict[str, Path], *, prefix: str = "Cohort",
) -> subprocess.CompletedProcess[str]:
return subprocess.run(
[
sys.executable, "-m", "apps.cli.run_wgs_index", "--workspace", str(workspace),
"--job-id", "index-job", "--prefix", prefix, "--reference-fasta", str(reference),
"--samtools", str(tools["samtools"]), "--bwa-mem2", str(tools["bwa_mem2"]),
"--gatk", str(tools["gatk"]), "--threads", "2",
],
text=True,
capture_output=True,
env=os.environ.copy(),
)
def _fixture(tmp_path: Path) -> tuple[Path, dict[str, Path]]:
reference = tmp_path / "reference.fasta"
reference.write_text(">chr1\nACGTACGT\n", encoding="utf-8")
return reference, {
"samtools": _tool(tmp_path / "samtools", "samtools"),
"bwa_mem2": _tool(tmp_path / "bwa-mem2", "bwa-mem2"),
"gatk": _tool(tmp_path / "gatk", "gatk"),
}
def _manifest(workspace: Path) -> dict:
return json.loads((workspace / "final" / "Cohort.wgs_index.manifest.json").read_text(encoding="utf-8"))
def test_wgs_index_builds_all_required_hashed_roles_in_immutable_bundle(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
result = _run(workspace, reference, tools)
assert result.returncode == 0, result.stderr
manifest = _manifest(workspace)
assert manifest["status"] == "success"
assert re.fullmatch(r"[0-9a-f]{32}", manifest["run_id"])
assert manifest["wgs_reference_id"].startswith("wgs-reference:")
assert manifest["wgs_analysis_id"].startswith("wgs-analysis:")
records = list(manifest["artifacts"].values())
assert {item["role"] for item in records} >= {"reference_fasta", "reference_fai", "reference_dict", "bwa_mem2_index_component"}
assert {item.get("component") for item in records if item["role"] == "bwa_mem2_index_component"} == set(COMPONENTS)
assert all(item["sha256"] and item["size"] > 0 for item in records)
paths = [Path(item["full_path"]) for item in records]
assert all("/.wgs_index_runs/Cohort/" in str(path) for path in paths)
assert not any(path.name.endswith((".bwt", ".sa")) for path in paths)
def test_wgs_index_reuses_only_matching_complete_authority(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
first = _manifest(workspace)
assert _run(workspace, reference, tools).returncode == 0
assert _manifest(workspace)["run_id"] == first["run_id"]
reference.write_text(">chr1\nACGTACGA\n", encoding="utf-8")
assert _run(workspace, reference, tools).returncode == 0
changed = _manifest(workspace)
assert changed["run_id"] != first["run_id"]
missing = next(item for item in changed["artifacts"].values() if item.get("component") == ".ann")
Path(missing["full_path"]).unlink()
assert _run(workspace, reference, tools).returncode == 0
assert _manifest(workspace)["run_id"] != changed["run_id"]
def test_wgs_index_failed_rebuild_preserves_old_success_and_writes_sidecar(tmp_path: Path, monkeypatch) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
old = _manifest(workspace)
reference.write_text(">chr1\nACGTACGA\n", encoding="utf-8")
monkeypatch.setenv("WGS_INDEX_FAIL", "gatk")
result = _run(workspace, reference, tools)
assert result.returncode == 1
assert _manifest(workspace)["run_id"] == old["run_id"]
failures = list((workspace / "final").glob("Cohort.wgs_index.failed.*.manifest.json"))
assert len(failures) == 1
assert json.loads(failures[0].read_text(encoding="utf-8"))["artifacts"] == {}
def test_wgs_index_parser_is_public_and_requires_root_reference() -> None:
parser = build_parser()
assert parser.parse_args(["--workspace", "w", "--job-id", "j", "--prefix", "p", "--reference-fasta", "r"]).reference_fasta == "r"
assert main(["--workspace", "w", "--job-id", "j", "--prefix", "p", "--reference-fasta", "missing"]) == 1
def test_wgs_index_reuse_rejects_external_symlink_artifact(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
old = _manifest(workspace)
artifact = next(item for item in old["artifacts"].values() if item.get("component") == ".ann")
bundle_path = Path(artifact["full_path"])
external = tmp_path / "external.ann"
shutil.copy2(bundle_path, external)
bundle_path.unlink()
bundle_path.symlink_to(external)
assert _run(workspace, reference, tools).returncode == 0
assert _manifest(workspace)["run_id"] != old["run_id"]
def test_wgs_index_reuse_rejects_path_full_path_mismatch(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
old = _manifest(workspace)
artifact = next(iter(old["artifacts"].values()))
artifact["path"] = "final/not-the-declared-full-path.fasta"
(workspace / "final" / "Cohort.wgs_index.manifest.json").write_text(json.dumps(old), encoding="utf-8")
assert _run(workspace, reference, tools).returncode == 0
assert _manifest(workspace)["run_id"] != old["run_id"]
def test_wgs_index_reuse_rejects_artifact_from_wrong_run_bundle(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
old = _manifest(workspace)
artifact = next(item for item in old["artifacts"].values() if item.get("component") == ".pac")
source = Path(artifact["full_path"])
wrong = workspace / "final" / ".wgs_index_runs" / "Cohort" / "wrong-run" / source.name
wrong.parent.mkdir(parents=True)
shutil.copy2(source, wrong)
artifact["path"] = str(wrong.relative_to(workspace))
artifact["full_path"] = str(wrong.resolve())
(workspace / "final" / "Cohort.wgs_index.manifest.json").write_text(json.dumps(old), encoding="utf-8")
assert _run(workspace, reference, tools).returncode == 0
assert _manifest(workspace)["run_id"] != old["run_id"]
def test_wgs_index_reuse_requires_nonempty_run_id(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
old = _manifest(workspace)
old["run_id"] = ""
(workspace / "final" / "Cohort.wgs_index.manifest.json").write_text(json.dumps(old), encoding="utf-8")
assert _run(workspace, reference, tools).returncode == 0
assert _manifest(workspace)["run_id"]
def test_wgs_index_reuse_rejects_parent_run_id_pointing_at_another_prefix_bundle(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
authority = _manifest(workspace)
source_bundle = Path(next(iter(authority["artifacts"].values()))["full_path"]).parent
rogue_bundle = workspace / "final" / ".wgs_index_runs" / "OtherPrefix" / "other-run"
shutil.copytree(source_bundle, rogue_bundle)
for artifact in authority["artifacts"].values():
rogue_path = rogue_bundle / Path(artifact["full_path"]).name
artifact["path"] = str(rogue_path.relative_to(workspace))
artifact["full_path"] = str(rogue_path.resolve())
authority["run_id"] = ".."
authority_path = workspace / "final" / "Cohort.wgs_index.manifest.json"
authority_path.write_text(json.dumps(authority), encoding="utf-8")
assert _run(workspace, reference, tools).returncode == 0
rebuilt = _manifest(workspace)
assert re.fullmatch(r"[0-9a-f]{32}", rebuilt["run_id"])
assert rebuilt["run_id"] != authority["run_id"]
assert "reuse" not in rebuilt
def test_wgs_index_reuse_rejects_run_id_outside_generated_lowercase_hex_format(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
authority = _manifest(workspace)
forged_run_id = "A" * 32
source_bundle = Path(next(iter(authority["artifacts"].values()))["full_path"]).parent
forged_bundle = workspace / "final" / ".wgs_index_runs" / "Cohort" / forged_run_id
shutil.copytree(source_bundle, forged_bundle)
for artifact in authority["artifacts"].values():
forged_path = forged_bundle / Path(artifact["full_path"]).name
artifact["path"] = str(forged_path.relative_to(workspace))
artifact["full_path"] = str(forged_path.resolve())
authority["run_id"] = forged_run_id
(workspace / "final" / "Cohort.wgs_index.manifest.json").write_text(
json.dumps(authority), encoding="utf-8",
)
assert _run(workspace, reference, tools).returncode == 0
rebuilt = _manifest(workspace)
assert re.fullmatch(r"[0-9a-f]{32}", rebuilt["run_id"])
assert rebuilt["run_id"] != forged_run_id
assert "reuse" not in rebuilt
@pytest.mark.parametrize(
"prefix",
["", ".", "..", "../escape", "bad/name", "bad\\name", "bad\nname", "bad\x7fname"],
)
def test_wgs_index_rejects_unsafe_prefix_before_creating_any_workspace_path(tmp_path: Path, prefix: str) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "nested" / "workspace"
escaped_manifest = (workspace / "final" / f"{prefix}.wgs_index.manifest.json").resolve()
before = {path.relative_to(tmp_path) for path in tmp_path.rglob("*")}
result = _run(workspace, reference, tools, prefix=prefix)
assert result.returncode == 1
assert "prefix" in result.stderr.lower()
assert not workspace.exists()
assert not escaped_manifest.exists()
assert {path.relative_to(tmp_path) for path in tmp_path.rglob("*")} == before
def test_wgs_index_rejects_nonempty_corrupt_fai(tmp_path: Path, monkeypatch) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
monkeypatch.setenv("WGS_INDEX_CORRUPT", "fai")
result = _run(workspace, reference, tools)
assert result.returncode == 1
manifest = _manifest(workspace)
assert manifest["status"] == "failed" and manifest["artifacts"] == {}
assert "FAI" in manifest["errors"][0]
def test_wgs_index_rejects_nonempty_corrupt_sequence_dictionary(tmp_path: Path, monkeypatch) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
monkeypatch.setenv("WGS_INDEX_CORRUPT", "dict")
result = _run(workspace, reference, tools)
assert result.returncode == 1
manifest = _manifest(workspace)
assert manifest["status"] == "failed" and manifest["artifacts"] == {}
assert "dictionary" in manifest["errors"][0]
def test_older_publish_attempt_preserves_newer_authority_bundle_and_writes_stale_sidecar(tmp_path: Path) -> None:
reference, tools = _fixture(tmp_path)
workspace = tmp_path / "workspace"
assert _run(workspace, reference, tools).returncode == 0
authority_path = workspace / "final" / "Cohort.wgs_index.manifest.json"
authority = _manifest(workspace)
authority["start_time"] = "2999-01-01T00:00:00+00:00"
authority_path.write_text(json.dumps(authority, sort_keys=True), encoding="utf-8")
authority_bytes = authority_path.read_bytes()
bundle_hashes = {
item["full_path"]: item["sha256"] for item in authority["artifacts"].values()
}
reference.write_text(">chr1\nACGTACGA\n", encoding="utf-8")
result = _run(workspace, reference, tools)
assert result.returncode == 1
assert authority_path.read_bytes() == authority_bytes
assert all(Path(path).is_file() for path in bundle_hashes)
assert all(__import__("hashlib").sha256(Path(path).read_bytes()).hexdigest() == digest for path, digest in bundle_hashes.items())
failures = list((workspace / "final").glob("Cohort.wgs_index.failed.*.manifest.json"))
assert len(failures) == 1
failure = json.loads(failures[0].read_text(encoding="utf-8"))
assert failure["run_id"] != authority["run_id"]
assert failure["artifacts"] == {}
assert any("stale" in error.lower() for error in failure["errors"])