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

137 lines
7.5 KiB
Python

from __future__ import annotations
import json
from pathlib import Path
import subprocess
from apps.web import pipeline_backend
from apps.web.vue3.main import ExecuteRequest
from gfa_core.mas.agents.planning_v2 import PlanningAgentV2
from gfa_core.mas.stage_capabilities import CAPABILITIES
from gfa_core.mas.stage_catalog import EXECUTABLE_STAGES, STAGE_DEPENDENCIES
from gfa_core.mas.tools.stage_tools import StageToolRegistry
SAMPLE_SET = "wgs-samples:" + "a" * 64
REFERENCE = "wgs-reference:" + "b" * 64
ANALYSIS = "wgs-analysis:" + "c" * 64
def _producer(path: Path, stage: str, *, sample: str = "", reference: str = "") -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps({"manifest_version": "1.0", "stage": stage, "status": "success",
"prefix": "Cohort", "job_id": "job-1", "workspace": str(path.parent.parent),
"wgs_sample_set_id": sample, "wgs_reference_id": reference, "artifacts": {}}))
def _align_manifest() -> dict:
artifacts = {}
for sample in ("S1", "S2", "S3"):
for role in ("sorted_bam", "sorted_bam_index", "align_log"):
artifacts[f"{role}.{sample}"] = {"role": role, "sample_id": sample, "sha256": "d" * 64}
return {"stage": "wgs_align", "status": "success", "prefix": "Cohort", "job_id": "job-1",
"wgs_sample_set_id": SAMPLE_SET, "wgs_reference_id": REFERENCE, "wgs_analysis_id": ANALYSIS,
"metrics": {"input_sample_ids": ["S1", "S2", "S3"], "output_sample_ids": ["S1", "S2", "S3"]},
"artifacts": artifacts}
def test_wgs_align_vertical_registration_is_exact() -> None:
assert STAGE_DEPENDENCIES["wgs_align"] == ["wgs_qc", "wgs_index"]
assert "wgs_align" in EXECUTABLE_STAGES
capability = CAPABILITIES["wgs_align"]
assert {item.role for item in capability.consumes} == {
"clean_reads_r1", "clean_reads_r2", "reference_fasta", "reference_fai", "reference_dict", "bwa_mem2_index_component",
}
assert {item.role for item in capability.provides} == {"sorted_bam", "sorted_bam_index", "align_log"}
definition = next(item for item in StageToolRegistry.STAGE_DEFINITIONS if item["id"] == "wgs_align")
assert definition["script"] == "run_wgs_align.py"
assert PlanningAgentV2.REQUIRED_INPUTS["wgs_align"] == []
def test_backend_uses_only_explicit_authority_manifests_and_fresh_complete_result(monkeypatch, tmp_path: Path) -> None:
workspace = tmp_path / "workspace"; final = workspace / "final"; final.mkdir(parents=True)
qc = final / "Cohort.wgs_qc.manifest.json"; index = final / "Cohort.wgs_index.manifest.json"
_producer(qc, "wgs_qc", sample=SAMPLE_SET); _producer(index, "wgs_index", reference=REFERENCE)
output = final / "Cohort.wgs_align.manifest.json"; output.write_text(json.dumps({"status": "success"}))
calls = []
monkeypatch.setattr(pipeline_backend, "resolve_wgs_artifacts", lambda *args, **kwargs: {role: (object(),) for role in args[2]})
def fake_run(command, *_args, **_kwargs):
calls.append(command); output.write_text(json.dumps(_align_manifest())); return 0
monkeypatch.setattr(pipeline_backend, "run_cmd_logged", fake_run)
message, returned = pipeline_backend.run_stage(
stage="wgs_align", job=pipeline_backend.JobInfo("job-1", str(tmp_path / "input"), str(workspace), str(workspace / "logs"), "now"),
prefix="Cohort", wgs_qc_manifest=str(qc), wgs_index_manifest=str(index),
wgs_sample_set_id=SAMPLE_SET, wgs_reference_id=REFERENCE,
wgs_samtools="sam", wgs_bwa_mem2="bwa", wgs_threads=3, wgs_max_jobs=2, wgs_platform="illumina",
)
assert "completed" in message.lower() and returned == output
command = calls[0]
assert command[:3] == [pipeline_backend.sys.executable, "-m", "apps.cli.run_wgs_align"]
assert command[command.index("--qc-manifest") + 1] == str(qc)
assert command[command.index("--index-manifest") + 1] == str(index)
assert pipeline_backend._valid_wgs_align_manifest(_align_manifest(), job_id="job-1", prefix="Cohort")
partial = _align_manifest(); partial["artifacts"].pop("sorted_bam.S1")
assert not pipeline_backend._valid_wgs_align_manifest(partial, job_id="job-1", prefix="Cohort")
def test_backend_never_scans_for_missing_align_manifests(tmp_path: Path) -> None:
workspace = tmp_path / "workspace"; final = workspace / "final"; final.mkdir(parents=True)
(final / "plausible.wgs_qc.manifest.json").write_text("{}")
message, returned = pipeline_backend.run_stage(
stage="wgs_align", job=pipeline_backend.JobInfo("job-1", str(tmp_path / "input"), str(workspace), str(workspace / "logs"), "now"), prefix="Cohort",
)
assert "explicit" in message.lower() and returned is None
def test_api_and_frontend_carry_both_authorities_parameters_and_identities() -> None:
request = ExecuteRequest(stage="wgs_align", wgsQcManifest="final/Cohort.wgs_qc.manifest.json",
wgsIndexManifest="final/Cohort.wgs_index.manifest.json", wgsSampleSetId=SAMPLE_SET,
wgsReferenceId=REFERENCE, wgsPlatform="illumina")
assert request.wgsQcManifest and request.wgsIndexManifest and request.wgsSampleSetId == SAMPLE_SET
text = (Path(__file__).resolve().parents[2] / "apps/web/vue3/index.html").read_text(encoding="utf-8")
for token in ("wgs_align", "wgsQcManifest", "wgsIndexManifest", "wgsSampleSetId", "wgsReferenceId", "wgsPlatform"):
assert token in text
worker_text = (Path(__file__).resolve().parents[2] / "gfa_core/mas/tools/stage_tools.py").read_text(encoding="utf-8")
for token in ("wgs_qc_manifest=", "wgs_index_manifest=", "wgs_sample_set_id=", "wgs_reference_id=", "wgs_platform="):
assert token in worker_text
def test_frontend_hands_completed_wgs_authorities_and_identities_to_pending_align() -> None:
text = (Path(__file__).resolve().parents[2] / "apps/web/vue3/index.html").read_text(encoding="utf-8")
start = text.index("const wgsAlignHandoffUpdates =")
end = text.index("const handoffWgsAlignProducerAuthority =", start)
helper = text[start:end]
payload = {
"qc": {
"result": {
"manifest_path": "/workspace/job/final/Cohort.wgs_qc.manifest.json",
"stage_status": {"completed": True, "stage": "wgs_qc", "prefix": "Cohort",
"wgs_sample_set_id": SAMPLE_SET, "wgs_analysis_id": ANALYSIS},
}
},
"index": {
"result": {
"manifest_path": "/workspace/job/final/Cohort.wgs_index.manifest.json",
"stage_status": {"completed": True, "stage": "wgs_index", "prefix": "Cohort",
"wgs_reference_id": REFERENCE, "wgs_analysis_id": "wgs-analysis:" + "e" * 64},
}
},
}
script = helper + "\n" + f"""
const payload = {json.dumps(payload)};
const qc = wgsAlignHandoffUpdates('wgs_qc', payload.qc, 'Cohort', {{}});
const index = wgsAlignHandoffUpdates('wgs_index', payload.index, 'Cohort', qc);
console.log(JSON.stringify({{qc, index}}));
"""
completed = subprocess.run(["/usr/bin/node", "-e", script], check=True, capture_output=True, text=True)
updates = json.loads(completed.stdout)
assert updates["qc"] == {
"wgsQcManifest": "/workspace/job/final/Cohort.wgs_qc.manifest.json",
"wgsSampleSetId": SAMPLE_SET,
}
assert updates["index"] == {
"wgsIndexManifest": "/workspace/job/final/Cohort.wgs_index.manifest.json",
"wgsReferenceId": REFERENCE,
}
assert "handoffWgsAlignProducerAuthority(stageAction, stageResult);" in text