370 lines
12 KiB
Python
370 lines
12 KiB
Python
|
import asyncio
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from gfa_core.mas.agents.planning_v2 import PlanningAgentV2, StageStatus
|
||
|
|
from gfa_core.mas.artifact_resolver import ArtifactResolver
|
||
|
|
from gfa_core.mas.context_index import build_project_context
|
||
|
|
from gfa_core.mas.stage_capabilities import get_stage_capability
|
||
|
|
import gfa_core.wgs_upstream_manifest as wgs_manifest
|
||
|
|
|
||
|
|
|
||
|
|
WGS_SAMPLE_SET_ID = "wgs-samples:" + "a" * 64
|
||
|
|
WGS_REFERENCE_ID = "wgs-reference:" + "b" * 64
|
||
|
|
WGS_QC_ANALYSIS_ID = "wgs-analysis:" + "c" * 64
|
||
|
|
WGS_INDEX_ANALYSIS_ID = "wgs-analysis:" + "d" * 64
|
||
|
|
WGS_RUN_ID = "1" * 32
|
||
|
|
|
||
|
|
|
||
|
|
def _wgs_artifact(
|
||
|
|
workspace: Path,
|
||
|
|
bundle: Path,
|
||
|
|
name: str,
|
||
|
|
role: str,
|
||
|
|
*,
|
||
|
|
sample_id: str = "",
|
||
|
|
component: str = "",
|
||
|
|
) -> dict:
|
||
|
|
path = bundle / name
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
path.write_bytes(f"{role}:{sample_id or component or 'reference'}".encode())
|
||
|
|
artifact = {
|
||
|
|
"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:
|
||
|
|
artifact["sample_id"] = sample_id
|
||
|
|
if component:
|
||
|
|
artifact["component"] = component
|
||
|
|
return artifact
|
||
|
|
|
||
|
|
|
||
|
|
def _write_wgs_align_authorities(
|
||
|
|
workspace: Path, prefix: str = "Cohort"
|
||
|
|
) -> dict[str, Path]:
|
||
|
|
final_dir = workspace / "final"
|
||
|
|
qc_bundle = final_dir / ".wgs_qc_runs" / prefix / WGS_RUN_ID
|
||
|
|
index_bundle = final_dir / ".wgs_index_runs" / prefix / WGS_RUN_ID
|
||
|
|
|
||
|
|
qc_artifacts = {}
|
||
|
|
for sample_id in ("S1", "S2"):
|
||
|
|
qc_artifacts[f"clean_reads_r1.{sample_id}"] = _wgs_artifact(
|
||
|
|
workspace,
|
||
|
|
qc_bundle,
|
||
|
|
f"{prefix}.wgs_qc.{sample_id}.R1.fastq.gz",
|
||
|
|
"clean_reads_r1",
|
||
|
|
sample_id=sample_id,
|
||
|
|
)
|
||
|
|
qc_artifacts[f"clean_reads_r2.{sample_id}"] = _wgs_artifact(
|
||
|
|
workspace,
|
||
|
|
qc_bundle,
|
||
|
|
f"{prefix}.wgs_qc.{sample_id}.R2.fastq.gz",
|
||
|
|
"clean_reads_r2",
|
||
|
|
sample_id=sample_id,
|
||
|
|
)
|
||
|
|
qc_manifest = {
|
||
|
|
"manifest_version": "1.0",
|
||
|
|
"stage": "wgs_qc",
|
||
|
|
"status": "success",
|
||
|
|
"prefix": prefix,
|
||
|
|
"job_id": "job-1",
|
||
|
|
"run_id": WGS_RUN_ID,
|
||
|
|
"workspace": str(workspace.resolve()),
|
||
|
|
"wgs_sample_set_id": WGS_SAMPLE_SET_ID,
|
||
|
|
"wgs_analysis_id": WGS_QC_ANALYSIS_ID,
|
||
|
|
"artifacts": qc_artifacts,
|
||
|
|
}
|
||
|
|
qc_path = final_dir / f"{prefix}.wgs_qc.manifest.json"
|
||
|
|
qc_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
qc_path.write_text(json.dumps(qc_manifest), encoding="utf-8")
|
||
|
|
|
||
|
|
index_artifacts = {
|
||
|
|
"reference_fasta": _wgs_artifact(
|
||
|
|
workspace, index_bundle, "reference.fasta", "reference_fasta"
|
||
|
|
),
|
||
|
|
"reference_fai": _wgs_artifact(
|
||
|
|
workspace, index_bundle, "reference.fasta.fai", "reference_fai"
|
||
|
|
),
|
||
|
|
"reference_dict": _wgs_artifact(
|
||
|
|
workspace, index_bundle, "reference.dict", "reference_dict"
|
||
|
|
),
|
||
|
|
}
|
||
|
|
for component in (".0123", ".amb", ".ann", ".bwt.2bit.64", ".pac"):
|
||
|
|
index_artifacts[f"bwa_mem2_index_component.{component}"] = _wgs_artifact(
|
||
|
|
workspace,
|
||
|
|
index_bundle,
|
||
|
|
f"reference.fasta{component}",
|
||
|
|
"bwa_mem2_index_component",
|
||
|
|
component=component,
|
||
|
|
)
|
||
|
|
index_manifest = {
|
||
|
|
"manifest_version": "1.0",
|
||
|
|
"stage": "wgs_index",
|
||
|
|
"status": "success",
|
||
|
|
"prefix": prefix,
|
||
|
|
"job_id": "job-1",
|
||
|
|
"run_id": WGS_RUN_ID,
|
||
|
|
"workspace": str(workspace.resolve()),
|
||
|
|
"wgs_reference_id": WGS_REFERENCE_ID,
|
||
|
|
"wgs_analysis_id": WGS_INDEX_ANALYSIS_ID,
|
||
|
|
"artifacts": index_artifacts,
|
||
|
|
}
|
||
|
|
index_path = final_dir / f"{prefix}.wgs_index.manifest.json"
|
||
|
|
index_path.write_text(json.dumps(index_manifest), encoding="utf-8")
|
||
|
|
return {
|
||
|
|
"qc_manifest": qc_path,
|
||
|
|
"index_manifest": index_path,
|
||
|
|
"missing_component": Path(
|
||
|
|
index_artifacts["bwa_mem2_index_component..ann"]["full_path"]
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def test_wgs_qc_resolves_only_canonical_sample_yaml(tmp_path: Path) -> None:
|
||
|
|
input_dir = tmp_path / "input" / "job-1"
|
||
|
|
input_dir.mkdir(parents=True)
|
||
|
|
canonical = input_dir / "wgs_samples.yaml"
|
||
|
|
canonical.write_text("samples: {}\n", encoding="utf-8")
|
||
|
|
(input_dir / "samples.yaml").write_text("samples: {}\n", encoding="utf-8")
|
||
|
|
|
||
|
|
context = build_project_context(
|
||
|
|
job_id="job-1",
|
||
|
|
prefix="At",
|
||
|
|
project_root=tmp_path,
|
||
|
|
)
|
||
|
|
capability = get_stage_capability("wgs_qc")
|
||
|
|
assert capability is not None
|
||
|
|
decision = ArtifactResolver(context).resolve_need(
|
||
|
|
"wgs_qc", capability.consumes[0]
|
||
|
|
)
|
||
|
|
|
||
|
|
assert decision.resolved
|
||
|
|
assert decision.selected is not None
|
||
|
|
assert Path(decision.selected.path) == canonical
|
||
|
|
|
||
|
|
|
||
|
|
def test_wgs_index_preflight_materializes_resolved_reference(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.web.vue3 import main as web_main
|
||
|
|
|
||
|
|
input_dir = tmp_path / "input" / "job-1"
|
||
|
|
input_dir.mkdir(parents=True)
|
||
|
|
reference = input_dir / "reference.9.1m.fasta"
|
||
|
|
reference.write_text(">chr1\nACGT\n", encoding="utf-8")
|
||
|
|
monkeypatch.setattr(web_main, "PROJECT_ROOT", tmp_path)
|
||
|
|
|
||
|
|
result = web_main._api_semantic_preflight(
|
||
|
|
"job-1",
|
||
|
|
"wgs_index",
|
||
|
|
"Cohort",
|
||
|
|
{"wgsReferenceFasta": ""},
|
||
|
|
)
|
||
|
|
|
||
|
|
assert result["valid"]
|
||
|
|
assert result["materialized_params"]["wgsReferenceFasta"] == str(reference)
|
||
|
|
|
||
|
|
|
||
|
|
def test_wgs_index_preflight_does_not_override_explicit_reference(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.web.vue3 import main as web_main
|
||
|
|
|
||
|
|
input_dir = tmp_path / "input" / "job-1"
|
||
|
|
input_dir.mkdir(parents=True)
|
||
|
|
(input_dir / "auto-reference.fasta").write_text(
|
||
|
|
">chr1\nACGT\n", encoding="utf-8"
|
||
|
|
)
|
||
|
|
explicit = tmp_path / "explicit-reference.fasta"
|
||
|
|
explicit.write_text(">chr1\nTGCA\n", encoding="utf-8")
|
||
|
|
monkeypatch.setattr(web_main, "PROJECT_ROOT", tmp_path)
|
||
|
|
|
||
|
|
result = web_main._api_semantic_preflight(
|
||
|
|
"job-1",
|
||
|
|
"wgs_index",
|
||
|
|
"Cohort",
|
||
|
|
{"wgsReferenceFasta": str(explicit)},
|
||
|
|
)
|
||
|
|
|
||
|
|
assert result["valid"]
|
||
|
|
assert "wgsReferenceFasta" not in result.get("materialized_params", {})
|
||
|
|
|
||
|
|
|
||
|
|
def test_wgs_align_planning_accepts_complete_authoritative_inputs(
|
||
|
|
tmp_path: Path,
|
||
|
|
) -> None:
|
||
|
|
workspace = tmp_path / "workspace" / "job-1"
|
||
|
|
_write_wgs_align_authorities(workspace)
|
||
|
|
input_dir = tmp_path / "input" / "job-1"
|
||
|
|
input_dir.mkdir(parents=True)
|
||
|
|
context = build_project_context(
|
||
|
|
job_id="job-1", prefix="Cohort", project_root=tmp_path
|
||
|
|
)
|
||
|
|
|
||
|
|
issues = PlanningAgentV2(None, project_root=str(tmp_path))._check_input_files(
|
||
|
|
"wgs_align", input_dir, workspace, "Cohort", context_index=context
|
||
|
|
)
|
||
|
|
|
||
|
|
assert not [issue for issue in issues if issue["severity"] == "error"]
|
||
|
|
assert not [
|
||
|
|
issue
|
||
|
|
for issue in issues
|
||
|
|
if "缺少 wgs_align 所需输入" in issue["message"]
|
||
|
|
and any(
|
||
|
|
role in issue["message"]
|
||
|
|
for role in (
|
||
|
|
"clean_reads_r1",
|
||
|
|
"clean_reads_r2",
|
||
|
|
"reference_fasta",
|
||
|
|
"reference_fai",
|
||
|
|
"reference_dict",
|
||
|
|
"bwa_mem2_index_component",
|
||
|
|
)
|
||
|
|
)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def test_wgs_align_planning_reports_one_invalid_authority_without_fallback(
|
||
|
|
tmp_path: Path,
|
||
|
|
) -> None:
|
||
|
|
workspace = tmp_path / "workspace" / "job-1"
|
||
|
|
authorities = _write_wgs_align_authorities(workspace)
|
||
|
|
authorities["missing_component"].unlink()
|
||
|
|
input_dir = tmp_path / "input" / "job-1"
|
||
|
|
input_dir.mkdir(parents=True)
|
||
|
|
context = build_project_context(
|
||
|
|
job_id="job-1", prefix="Cohort", project_root=tmp_path
|
||
|
|
)
|
||
|
|
|
||
|
|
issues = PlanningAgentV2(None, project_root=str(tmp_path))._check_input_files(
|
||
|
|
"wgs_align", input_dir, workspace, "Cohort", context_index=context
|
||
|
|
)
|
||
|
|
|
||
|
|
errors = [issue for issue in issues if issue["severity"] == "error"]
|
||
|
|
assert errors == [
|
||
|
|
{
|
||
|
|
"severity": "error",
|
||
|
|
"category": "input",
|
||
|
|
"code": "invalid_wgs_manifest_identity",
|
||
|
|
"message": (
|
||
|
|
"artifact bwa_mem2_index_component..ann file is missing or not regular"
|
||
|
|
),
|
||
|
|
"suggestion": "修复 wgs_qc 或 wgs_index 权威 Manifest 后重试",
|
||
|
|
"auto_fixable": False,
|
||
|
|
}
|
||
|
|
]
|
||
|
|
assert not [
|
||
|
|
issue for issue in issues if "缺少 wgs_align 所需输入" in issue["message"]
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def test_wgs_align_planning_allows_planned_upstream_rebuild(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
agent = PlanningAgentV2(None, project_root=str(tmp_path))
|
||
|
|
|
||
|
|
def fail_if_read(*_args, **_kwargs):
|
||
|
|
raise wgs_manifest.WgsUpstreamManifestError(
|
||
|
|
"wgs_qc authority is invalid: upstream Manifest is unreadable"
|
||
|
|
)
|
||
|
|
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"gfa_core.mas.agents.planning_v2.materialize_wgs_identity_params",
|
||
|
|
fail_if_read,
|
||
|
|
)
|
||
|
|
|
||
|
|
issues = agent._check_input_files(
|
||
|
|
"wgs_align",
|
||
|
|
tmp_path / "input",
|
||
|
|
tmp_path / "workspace" / "job-1",
|
||
|
|
"Cohort",
|
||
|
|
planned_stages=["wgs_qc", "wgs_index", "wgs_align"],
|
||
|
|
)
|
||
|
|
|
||
|
|
assert [issue["severity"] for issue in issues] == ["warning"]
|
||
|
|
assert issues[0]["code"] == "wgs_upstream_will_be_rebuilt"
|
||
|
|
|
||
|
|
|
||
|
|
def test_unique_wgs_prefix_ignores_stale_request_hint(tmp_path: Path) -> None:
|
||
|
|
workspace = tmp_path / "workspace" / "job-1"
|
||
|
|
_write_wgs_align_authorities(workspace, "Os")
|
||
|
|
|
||
|
|
assert wgs_manifest.resolve_unique_wgs_prefix(
|
||
|
|
stage="wgs_align",
|
||
|
|
job_id="job-1",
|
||
|
|
params={},
|
||
|
|
workspace=workspace,
|
||
|
|
) == "Os"
|
||
|
|
|
||
|
|
|
||
|
|
def test_unique_wgs_prefix_rejects_multiple_complete_prefixes(tmp_path: Path) -> None:
|
||
|
|
workspace = tmp_path / "workspace" / "job-1"
|
||
|
|
_write_wgs_align_authorities(workspace, "At")
|
||
|
|
_write_wgs_align_authorities(workspace, "Os")
|
||
|
|
|
||
|
|
with pytest.raises(
|
||
|
|
wgs_manifest.WgsUpstreamManifestError,
|
||
|
|
match="multiple complete WGS authority prefixes",
|
||
|
|
):
|
||
|
|
wgs_manifest.resolve_unique_wgs_prefix(
|
||
|
|
stage="wgs_align",
|
||
|
|
job_id="job-1",
|
||
|
|
params={},
|
||
|
|
workspace=workspace,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_planning_uses_manifest_prefix_for_wgs_consumer(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
workspace = tmp_path / "workspace" / "job-1"
|
||
|
|
_write_wgs_align_authorities(workspace, "Os")
|
||
|
|
input_dir = tmp_path / "input" / "job-1"
|
||
|
|
input_dir.mkdir(parents=True)
|
||
|
|
(input_dir / "wgs_samples.yaml").write_text("samples: []\n", encoding="utf-8")
|
||
|
|
(input_dir / "reference.fa").write_text(">chr1\nACGT\n", encoding="utf-8")
|
||
|
|
agent = PlanningAgentV2(None, project_root=str(tmp_path))
|
||
|
|
checked_prefixes: list[str] = []
|
||
|
|
monkeypatch.setattr(
|
||
|
|
agent,
|
||
|
|
"_check_stage_status_preflight",
|
||
|
|
lambda stage, prefix, checked_workspace, stage_params=None: (
|
||
|
|
checked_prefixes.append(prefix)
|
||
|
|
or StageStatus(stage=stage, status="not_found")
|
||
|
|
),
|
||
|
|
)
|
||
|
|
intent = {"stages": ["wgs_align"], "prefix": "At", "params": {}}
|
||
|
|
|
||
|
|
result = asyncio.run(
|
||
|
|
agent._run_preflight_validation(intent, "job-1", "At", str(workspace))
|
||
|
|
)
|
||
|
|
|
||
|
|
assert result["valid"]
|
||
|
|
assert intent["prefix"] == "Os"
|
||
|
|
assert checked_prefixes and set(checked_prefixes) == {"Os"}
|
||
|
|
|
||
|
|
|
||
|
|
def test_execute_prefix_uses_same_manifest_resolution(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.web.vue3 import main as web_main
|
||
|
|
|
||
|
|
workspace_root = tmp_path / "workspace"
|
||
|
|
_write_wgs_align_authorities(workspace_root / "job-1", "Os")
|
||
|
|
monkeypatch.setattr(web_main, "WORKSPACE_ROOT", workspace_root)
|
||
|
|
|
||
|
|
assert web_main._resolve_wgs_execute_prefix(
|
||
|
|
job_id="job-1",
|
||
|
|
stage="wgs_align",
|
||
|
|
prefix="At",
|
||
|
|
params={},
|
||
|
|
) == "Os"
|