144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from gfa_core.mas.agents.planning_v2 import PlanningAgentV2
|
|
|
|
|
|
GENOTYPE_SCOPE_ID = "genotype:" + "1" * 64
|
|
PHENOTYPE_SCOPE_ID = "phenotype:" + "2" * 64
|
|
COVARIATE_SCOPE_ID = "covariate:" + "3" * 64
|
|
|
|
|
|
def _write_manifest(final_dir: Path, stage: str, prefix: str, **scope_ids: str) -> Path:
|
|
path = final_dir / f"{prefix}.{stage}.manifest.json"
|
|
path.write_text(
|
|
json.dumps(
|
|
{
|
|
"stage": stage,
|
|
"prefix": prefix,
|
|
"status": "success",
|
|
**scope_ids,
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
return path
|
|
|
|
|
|
def _completed_preflight(workspace: Path, prefix: str = "Os") -> dict:
|
|
final_dir = workspace / "final"
|
|
final_dir.mkdir(parents=True)
|
|
manifests = {
|
|
"gwas_fmt": _write_manifest(
|
|
final_dir,
|
|
"gwas_fmt",
|
|
prefix,
|
|
genotype_scope_id=GENOTYPE_SCOPE_ID,
|
|
),
|
|
"gwas_kinship": _write_manifest(
|
|
final_dir,
|
|
"gwas_kinship",
|
|
prefix,
|
|
genotype_scope_id=GENOTYPE_SCOPE_ID,
|
|
),
|
|
"gwas_trait": _write_manifest(
|
|
final_dir,
|
|
"gwas_trait",
|
|
prefix,
|
|
phenotype_scope_id=PHENOTYPE_SCOPE_ID,
|
|
),
|
|
"gwas_cov": _write_manifest(
|
|
final_dir,
|
|
"gwas_cov",
|
|
prefix,
|
|
genotype_scope_id=GENOTYPE_SCOPE_ID,
|
|
covariate_scope_id=COVARIATE_SCOPE_ID,
|
|
),
|
|
}
|
|
stage_status = {
|
|
stage: {
|
|
"status": "success",
|
|
"manifest_path": str(path),
|
|
"output_files": [],
|
|
"missing_outputs": [],
|
|
}
|
|
for stage, path in manifests.items()
|
|
}
|
|
stage_status["gwas_mlm"] = {
|
|
"status": "not_found",
|
|
"manifest_path": None,
|
|
"output_files": [],
|
|
"missing_outputs": [],
|
|
}
|
|
return {"stage_status": stage_status}
|
|
|
|
|
|
def _plan_mlm(workspace: Path, preflight_result: dict) -> dict:
|
|
agent = PlanningAgentV2.__new__(PlanningAgentV2)
|
|
intent = {
|
|
"stages": ["gwas_mlm"],
|
|
"prefix": "Os",
|
|
"params": {},
|
|
"task_kind": "pipeline",
|
|
}
|
|
return asyncio.run(
|
|
agent._generate_plan_with_llm(
|
|
"使用 TASSEL 执行 GWAS MLM 分析",
|
|
intent,
|
|
completed_stages=[],
|
|
job_id="job_test",
|
|
workspace=str(workspace),
|
|
preflight_result=preflight_result,
|
|
)
|
|
)
|
|
|
|
|
|
def test_completed_gwas_dependencies_handoff_exact_manifest_scopes_to_mlm(tmp_path: Path) -> None:
|
|
workspace = tmp_path / "workspace" / "job_test"
|
|
plan = _plan_mlm(workspace, _completed_preflight(workspace))
|
|
|
|
assert [stage["id"] for stage in plan["stages"]] == ["gwas_mlm"]
|
|
params = plan["stages"][0]["params"]
|
|
assert params["genotypeScopeId"] == GENOTYPE_SCOPE_ID
|
|
assert params["phenotypeScopeId"] == PHENOTYPE_SCOPE_ID
|
|
assert params["covariateScopeId"] == COVARIATE_SCOPE_ID
|
|
|
|
|
|
def _rewrite_manifest(path: Path, **updates: str) -> None:
|
|
manifest = json.loads(path.read_text(encoding="utf-8"))
|
|
manifest.update(updates)
|
|
path.write_text(json.dumps(manifest), encoding="utf-8")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"invalid_case",
|
|
["missing_manifest", "wrong_prefix", "illegal_scope", "genotype_conflict"],
|
|
)
|
|
def test_completed_gwas_dependency_identity_is_all_or_nothing(
|
|
tmp_path: Path,
|
|
invalid_case: str,
|
|
) -> None:
|
|
workspace = tmp_path / "workspace" / "job_test"
|
|
preflight = _completed_preflight(workspace)
|
|
status = preflight["stage_status"]
|
|
if invalid_case == "missing_manifest":
|
|
status["gwas_trait"]["manifest_path"] = None
|
|
elif invalid_case == "wrong_prefix":
|
|
_rewrite_manifest(Path(status["gwas_trait"]["manifest_path"]), prefix="At")
|
|
elif invalid_case == "illegal_scope":
|
|
_rewrite_manifest(Path(status["gwas_trait"]["manifest_path"]), phenotype_scope_id="bad")
|
|
else:
|
|
_rewrite_manifest(
|
|
Path(status["gwas_cov"]["manifest_path"]),
|
|
genotype_scope_id="genotype:" + "4" * 64,
|
|
)
|
|
|
|
params = _plan_mlm(workspace, preflight)["stages"][0]["params"]
|
|
|
|
assert "genotypeScopeId" not in params
|
|
assert "phenotypeScopeId" not in params
|
|
assert "covariateScopeId" not in params
|