291 lines
9.2 KiB
Python
291 lines
9.2 KiB
Python
"""持久化与恢复测试：检查点恢复 + 运行历史管理。
|
||
|
||
覆盖：
|
||
* run(resume_from=RunReport) 跳过已成功任务。
|
||
* run(resume_from=Path) 从 JSON 文件恢复。
|
||
* 恢复后 FAILED/SKIPPED 任务被重新执行。
|
||
* 恢复后下游任务能获取上游恢复结果。
|
||
* RunHistory: save/load/list/latest/delete。
|
||
* RunHistory: __len__ / __contains__。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
import pyflowx as px
|
||
from pyflowx import Graph, RunHistory, TaskSpec
|
||
|
||
|
||
def test_resume_skips_successful_tasks() -> None:
|
||
"""resume_from 中的 SUCCESS 任务被跳过，不重新执行。"""
|
||
call_count: dict[str, int] = {"a": 0, "b": 0}
|
||
|
||
def make_a():
|
||
def _a() -> int:
|
||
call_count["a"] += 1
|
||
return 1
|
||
|
||
return _a
|
||
|
||
def make_b():
|
||
def _b(a: int) -> int:
|
||
call_count["b"] += 1
|
||
return a + 1
|
||
|
||
return _b
|
||
|
||
graph = Graph.from_specs(
|
||
[
|
||
TaskSpec("a", fn=make_a()),
|
||
TaskSpec("b", fn=make_b(), depends_on=("a",)),
|
||
]
|
||
)
|
||
# 第一次运行
|
||
report1 = px.run(graph, strategy="sequential")
|
||
assert call_count == {"a": 1, "b": 1}
|
||
|
||
# 第二次运行，从 report1 恢复
|
||
report2 = px.run(graph, strategy="sequential", resume_from=report1)
|
||
assert call_count == {"a": 1, "b": 1} # 没有重新执行
|
||
assert report2["a"] == 1
|
||
assert report2["b"] == 2
|
||
|
||
|
||
def test_resume_from_json_path(tmp_path: Path) -> None:
|
||
"""resume_from 接受 JSON 文件路径。"""
|
||
graph = Graph.from_specs([TaskSpec("a", fn=lambda: 42)])
|
||
report1 = px.run(graph, strategy="sequential")
|
||
|
||
json_path = tmp_path / "report.json"
|
||
json_path.write_text(report1.to_json(), encoding="utf-8")
|
||
|
||
call_count = 0
|
||
|
||
def _a() -> int:
|
||
nonlocal call_count
|
||
call_count += 1
|
||
return 42
|
||
|
||
graph2 = Graph.from_specs([TaskSpec("a", fn=_a)])
|
||
report2 = px.run(graph2, strategy="sequential", resume_from=json_path)
|
||
assert call_count == 0 # 从检查点恢复，未重新执行
|
||
assert report2["a"] == 42
|
||
|
||
|
||
def test_resume_reruns_failed_tasks() -> None:
|
||
"""resume_from 中 FAILED 的任务被重新执行。"""
|
||
state: dict[str, bool] = {"should_fail": True}
|
||
|
||
def _a() -> str:
|
||
if state["should_fail"]:
|
||
raise RuntimeError("第一次失败")
|
||
return "success"
|
||
|
||
graph = Graph.from_specs([TaskSpec("a", fn=_a)])
|
||
# 第一次运行失败
|
||
with pytest.raises(px.TaskFailedError):
|
||
px.run(graph, strategy="sequential")
|
||
|
||
# 修复问题
|
||
state["should_fail"] = False
|
||
|
||
# 第二次运行：从头开始（无 resume_from）
|
||
report2 = px.run(graph, strategy="sequential")
|
||
assert report2["a"] == "success"
|
||
|
||
|
||
def test_resume_downstream_gets_restored_result() -> None:
|
||
"""恢复后下游任务能获取上游恢复结果。
|
||
|
||
a 第一次成功，b 第一次失败；第二次 a 被恢复，b 重新执行获取 a 的恢复值。
|
||
"""
|
||
state: dict[str, bool] = {"b_fail": True}
|
||
b_call_count = 0
|
||
|
||
def _a() -> int:
|
||
return 10
|
||
|
||
def _b(a: int) -> int:
|
||
nonlocal b_call_count
|
||
b_call_count += 1
|
||
if state["b_fail"]:
|
||
raise RuntimeError("b 第一次失败")
|
||
return a * 2
|
||
|
||
graph = Graph.from_specs(
|
||
[
|
||
TaskSpec("a", fn=_a),
|
||
TaskSpec("b", fn=_b, depends_on=("a",)),
|
||
]
|
||
)
|
||
# 第一次运行：a 成功，b 失败
|
||
with pytest.raises(px.TaskFailedError):
|
||
px.run(graph, strategy="sequential")
|
||
|
||
# 修复 b
|
||
state["b_fail"] = False
|
||
b_call_count = 0 # 重置计数
|
||
|
||
# 第二次运行：a 被恢复（值=10），b 重新执行获取 a=10
|
||
# 需要构造一个只含 a 成功结果的 prev_report
|
||
prev_report = px.run(
|
||
Graph.from_specs([TaskSpec("a", fn=_a)]),
|
||
strategy="sequential",
|
||
)
|
||
report2 = px.run(graph, strategy="sequential", resume_from=prev_report)
|
||
assert report2["a"] == 10 # a 被恢复
|
||
assert b_call_count == 1 # b 被重新执行一次
|
||
assert report2["b"] == 20 # b 获取 a=10
|
||
|
||
|
||
def test_resume_with_dependency_strategy() -> None:
|
||
"""dependency 策略下 resume_from 也正常工作。"""
|
||
graph = Graph.from_specs(
|
||
[
|
||
TaskSpec("a", fn=lambda: 1),
|
||
TaskSpec("b", fn=lambda a: a + 1, depends_on=("a",)),
|
||
]
|
||
)
|
||
report1 = px.run(graph, strategy="dependency")
|
||
assert report1["a"] == 1
|
||
assert report1["b"] == 2
|
||
|
||
call_count = 0
|
||
|
||
def _a() -> int:
|
||
nonlocal call_count
|
||
call_count += 1
|
||
return 1
|
||
|
||
graph2 = Graph.from_specs(
|
||
[
|
||
TaskSpec("a", fn=_a),
|
||
TaskSpec("b", fn=lambda a: a + 1, depends_on=("a",)),
|
||
]
|
||
)
|
||
report2 = px.run(graph2, strategy="dependency", resume_from=report1)
|
||
assert call_count == 0 # a 被恢复
|
||
assert report2["a"] == 1
|
||
assert report2["b"] == 2
|
||
|
||
|
||
def test_resume_ignores_tasks_not_in_graph() -> None:
|
||
"""resume_from 中不在当前图里的任务被忽略。"""
|
||
graph1 = Graph.from_specs([TaskSpec("a", fn=lambda: 1), TaskSpec("b", fn=lambda: 2)])
|
||
report1 = px.run(graph1, strategy="sequential")
|
||
|
||
# 新图只有 a，没有 b
|
||
graph2 = Graph.from_specs([TaskSpec("a", fn=lambda: 999)])
|
||
report2 = px.run(graph2, strategy="sequential", resume_from=report1)
|
||
assert "b" not in report2.results
|
||
assert report2["a"] == 1 # a 被恢复
|
||
|
||
|
||
class TestRunHistory:
|
||
"""RunHistory 测试。"""
|
||
|
||
def test_save_and_load(self, tmp_path: Path) -> None:
|
||
"""保存后能正确加载。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
graph = Graph.from_specs([TaskSpec("a", fn=lambda: 42)])
|
||
report = px.run(graph, strategy="sequential")
|
||
path = history.save(report)
|
||
assert path.exists()
|
||
|
||
loaded = history.load(report.run_id)
|
||
assert loaded["a"] == 42
|
||
assert loaded.run_id == report.run_id
|
||
|
||
def test_list_runs(self, tmp_path: Path) -> None:
|
||
"""list_runs 返回全部 run_id。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
for i in range(3):
|
||
graph = Graph.from_specs([TaskSpec(f"t{i}", fn=lambda i=i: i)])
|
||
report = px.run(graph, strategy="sequential")
|
||
history.save(report)
|
||
runs = history.list_runs()
|
||
assert len(runs) == 3
|
||
|
||
def test_latest(self, tmp_path: Path) -> None:
|
||
"""latest 返回最近保存的报告。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
assert history.latest() is None
|
||
|
||
graph1 = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
|
||
r1 = px.run(graph1, strategy="sequential")
|
||
history.save(r1)
|
||
|
||
graph2 = Graph.from_specs([TaskSpec("b", fn=lambda: 2)])
|
||
r2 = px.run(graph2, strategy="sequential")
|
||
history.save(r2)
|
||
|
||
latest = history.latest()
|
||
assert latest is not None
|
||
assert latest.run_id == r2.run_id
|
||
|
||
def test_delete(self, tmp_path: Path) -> None:
|
||
"""delete 删除指定报告。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
|
||
report = px.run(graph, strategy="sequential")
|
||
history.save(report)
|
||
assert report.run_id in history
|
||
assert history.delete(report.run_id)
|
||
assert report.run_id not in history
|
||
assert not history.delete(report.run_id) # 再删返回 False
|
||
|
||
def test_load_not_found(self, tmp_path: Path) -> None:
|
||
"""加载不存在的 run_id 抛 FileNotFoundError。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
with pytest.raises(FileNotFoundError, match="不存在"):
|
||
history.load("nonexistent")
|
||
|
||
def test_len(self, tmp_path: Path) -> None:
|
||
"""__len__ 返回存储数量。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
assert len(history) == 0
|
||
graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
|
||
report = px.run(graph, strategy="sequential")
|
||
history.save(report)
|
||
assert len(history) == 1
|
||
|
||
def test_contains(self, tmp_path: Path) -> None:
|
||
"""__contains__ 检查 run_id 是否存在。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
assert "abc" not in history
|
||
graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
|
||
report = px.run(graph, strategy="sequential")
|
||
history.save(report)
|
||
assert report.run_id in history
|
||
|
||
def test_creates_dir(self, tmp_path: Path) -> None:
|
||
"""目录不存在时自动创建。"""
|
||
dir = tmp_path / "nested" / "runs"
|
||
assert not dir.exists()
|
||
RunHistory(dir)
|
||
assert dir.exists()
|
||
|
||
def test_resume_from_history(self, tmp_path: Path) -> None:
|
||
"""从 RunHistory 加载报告用于 resume_from。"""
|
||
history = RunHistory(tmp_path / "runs")
|
||
graph = Graph.from_specs([TaskSpec("a", fn=lambda: 42)])
|
||
report1 = px.run(graph, strategy="sequential")
|
||
history.save(report1)
|
||
|
||
call_count = 0
|
||
|
||
def _a() -> int:
|
||
nonlocal call_count
|
||
call_count += 1
|
||
return 999
|
||
|
||
graph2 = Graph.from_specs([TaskSpec("a", fn=_a)])
|
||
prev = history.latest()
|
||
assert prev is not None
|
||
report2 = px.run(graph2, strategy="sequential", resume_from=prev)
|
||
assert call_count == 0
|
||
assert report2["a"] == 42
|