256 lines
7.4 KiB
Python
256 lines
7.4 KiB
Python
"""run_iter() 流式执行测试。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import pytest
|
||
|
||
import pyflowx as px
|
||
from pyflowx.task import TaskStatus
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 基础流式执行
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_yields_all_results() -> None:
|
||
"""run_iter 逐个 yield 所有任务结果。"""
|
||
|
||
@px.task
|
||
def a() -> int:
|
||
return 1
|
||
|
||
@px.task(depends_on=("a",))
|
||
def b(a: int) -> int:
|
||
return a + 1
|
||
|
||
graph = px.Graph.from_specs([a, b])
|
||
results = dict(px.run_iter(graph))
|
||
assert len(results) == 2
|
||
assert results["a"].status == TaskStatus.SUCCESS
|
||
assert results["a"].value == 1
|
||
assert results["b"].status == TaskStatus.SUCCESS
|
||
assert results["b"].value == 2
|
||
|
||
|
||
def test_run_iter_returns_iterator() -> None:
|
||
"""run_iter 返回迭代器。"""
|
||
|
||
@px.task
|
||
def t() -> int:
|
||
return 42
|
||
|
||
graph = px.Graph.from_specs([t])
|
||
it = px.run_iter(graph)
|
||
assert hasattr(it, "__next__")
|
||
name, result = next(it)
|
||
assert name == "t"
|
||
assert result.value == 42
|
||
with pytest.raises(StopIteration):
|
||
next(it)
|
||
|
||
|
||
def test_run_iter_single_task() -> None:
|
||
"""单任务图流式执行。"""
|
||
|
||
@px.task
|
||
def solo() -> str:
|
||
return "hello"
|
||
|
||
graph = px.Graph.from_specs([solo])
|
||
results = list(px.run_iter(graph))
|
||
assert len(results) == 1
|
||
assert results[0][0] == "solo"
|
||
assert results[0][1].value == "hello"
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 四种策略
|
||
# ---------------------------------------------------------------------- #
|
||
@pytest.mark.parametrize("strategy", ["sequential", "thread", "async", "dependency"])
|
||
def test_run_iter_all_strategies(strategy: str) -> None:
|
||
"""run_iter 支持所有执行策略。"""
|
||
|
||
@px.task
|
||
def x() -> int:
|
||
return 10
|
||
|
||
@px.task(depends_on=("x",))
|
||
def y(x: int) -> int:
|
||
return x * 2
|
||
|
||
graph = px.Graph.from_specs([x, y])
|
||
results = dict(px.run_iter(graph, strategy=strategy)) # type: ignore[arg-type]
|
||
assert results["x"].value == 10
|
||
assert results["y"].value == 20
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 任务完成顺序
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_yields_in_completion_order_dependency() -> None:
|
||
"""dependency 策略下，x 在 y 之前完成。"""
|
||
|
||
@px.task
|
||
def x() -> int:
|
||
return 1
|
||
|
||
@px.task(depends_on=("x",))
|
||
def y(x: int) -> int:
|
||
return x + 1
|
||
|
||
graph = px.Graph.from_specs([x, y])
|
||
names = [name for name, _ in px.run_iter(graph, strategy="dependency")]
|
||
assert names.index("x") < names.index("y")
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 失败传播
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_yields_failed_then_reraises() -> None:
|
||
"""失败任务先 yield（FAILED），然后 re-raise TaskFailedError。"""
|
||
|
||
@px.task
|
||
def boom() -> None:
|
||
raise RuntimeError("explode")
|
||
|
||
graph = px.Graph.from_specs([boom])
|
||
results = []
|
||
with pytest.raises(px.TaskFailedError):
|
||
for name, result in px.run_iter(graph):
|
||
results.append((name, result))
|
||
assert len(results) == 1
|
||
assert results[0][0] == "boom"
|
||
assert results[0][1].status == TaskStatus.FAILED
|
||
|
||
|
||
def test_run_iter_failure_with_continue_on_error() -> None:
|
||
"""continue_on_error=True 时，失败任务 yield 后继续执行其他任务。"""
|
||
|
||
def fail() -> None:
|
||
raise ValueError("err")
|
||
|
||
graph = px.Graph.from_specs(
|
||
[
|
||
px.TaskSpec("fail", fn=fail, continue_on_error=True),
|
||
px.TaskSpec("ok", fn=lambda: 1),
|
||
]
|
||
)
|
||
results = dict(px.run_iter(graph))
|
||
assert results["fail"].status == TaskStatus.FAILED
|
||
assert results["ok"].status == TaskStatus.SUCCESS
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 跳过与缓存
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_yields_skipped_tasks() -> None:
|
||
"""条件跳过的任务以 SKIPPED 状态 yield。"""
|
||
never_true = lambda _ctx: False # noqa: E731
|
||
|
||
graph = px.Graph.from_specs(
|
||
[
|
||
px.TaskSpec("normal", fn=lambda: 1),
|
||
px.TaskSpec("skipped", fn=lambda: 2, conditions=(never_true,)),
|
||
]
|
||
)
|
||
results = dict(px.run_iter(graph))
|
||
assert results["normal"].status == TaskStatus.SUCCESS
|
||
assert results["skipped"].status == TaskStatus.SKIPPED
|
||
|
||
|
||
def test_run_iter_yields_cached_tasks() -> None:
|
||
"""缓存命中的任务以 SKIPPED 状态 yield。"""
|
||
|
||
@px.task
|
||
def cached() -> int:
|
||
return 42
|
||
|
||
graph = px.Graph.from_specs([cached])
|
||
backend = px.MemoryBackend()
|
||
backend.save("cached", 42)
|
||
|
||
results = dict(px.run_iter(graph, state=backend))
|
||
assert results["cached"].status == TaskStatus.SKIPPED
|
||
assert results["cached"].value == 42
|
||
assert results["cached"].reason is not None
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 早退（生成器未耗尽）
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_early_exit_does_not_block() -> None:
|
||
"""提前 break 不会阻塞生成器销毁。"""
|
||
|
||
@px.task
|
||
def a() -> int:
|
||
return 1
|
||
|
||
@px.task(depends_on=("a",))
|
||
def b(a: int) -> int:
|
||
return a + 1
|
||
|
||
@px.task(depends_on=("b",))
|
||
def c(b: int) -> int:
|
||
return b + 1
|
||
|
||
graph = px.Graph.from_specs([a, b, c])
|
||
count = 0
|
||
for _name, _result in px.run_iter(graph):
|
||
count += 1
|
||
if count >= 1:
|
||
break
|
||
assert count == 1
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 不可识别策略
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_unknown_strategy_raises() -> None:
|
||
"""不可识别的策略在生成器迭代时抛 ValueError。"""
|
||
|
||
@px.task
|
||
def t() -> int:
|
||
return 1
|
||
|
||
graph = px.Graph.from_specs([t])
|
||
with pytest.raises(ValueError, match="Unknown strategy"):
|
||
list(px.run_iter(graph, strategy="invalid")) # type: ignore[arg-type]
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 多任务图
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_diamond_graph() -> None:
|
||
"""菱形依赖图：a → b,c → d，所有结果均 yield。"""
|
||
|
||
@px.task
|
||
def a() -> int:
|
||
return 1
|
||
|
||
@px.task(depends_on=("a",))
|
||
def b(a: int) -> int:
|
||
return a + 10
|
||
|
||
@px.task(depends_on=("a",))
|
||
def c(a: int) -> int:
|
||
return a + 20
|
||
|
||
@px.task(depends_on=("b", "c"))
|
||
def d(b: int, c: int) -> int:
|
||
return b + c
|
||
|
||
graph = px.Graph.from_specs([a, b, c, d])
|
||
results = dict(px.run_iter(graph))
|
||
assert results["a"].value == 1
|
||
assert results["b"].value == 11
|
||
assert results["c"].value == 21
|
||
assert results["d"].value == 32
|
||
|
||
|
||
# ---------------------------------------------------------------------- #
|
||
# 公共 API 导出
|
||
# ---------------------------------------------------------------------- #
|
||
def test_run_iter_exported() -> None:
|
||
"""run_iter 从 pyflowx 顶层导出。"""
|
||
assert hasattr(px, "run_iter")
|
||
assert callable(px.run_iter)
|