367 lines
12 KiB
Python
367 lines
12 KiB
Python
"""测试任务取消与优雅停止."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
import pyflowx as px
|
|
from pyflowx.cancellation import CancelToken, _is_cancelled
|
|
from pyflowx.executors import _mark_remaining_skipped
|
|
from pyflowx.report import RunReport
|
|
from pyflowx.task import TaskStatus
|
|
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# CancelToken 单元测试
|
|
# ---------------------------------------------------------------------- #
|
|
class TestCancelToken:
|
|
"""测试 CancelToken 基本行为."""
|
|
|
|
def test_initial_state_not_cancelled(self) -> None:
|
|
"""新建的 token 不应处于取消状态."""
|
|
token = CancelToken()
|
|
assert not token.is_cancelled
|
|
assert not token.is_set()
|
|
assert not bool(token)
|
|
assert token.reason is None
|
|
|
|
def test_cancel_sets_state(self) -> None:
|
|
"""cancel() 后应处于取消状态."""
|
|
token = CancelToken()
|
|
token.cancel()
|
|
assert token.is_cancelled
|
|
assert token.is_set()
|
|
assert bool(token)
|
|
|
|
def test_cancel_with_reason(self) -> None:
|
|
"""cancel() 可附带原因描述."""
|
|
token = CancelToken()
|
|
token.cancel("用户中断")
|
|
assert token.is_cancelled
|
|
assert token.reason == "用户中断"
|
|
|
|
def test_cancel_without_reason_keeps_none(self) -> None:
|
|
"""cancel() 不传原因时 reason 保持 None."""
|
|
token = CancelToken()
|
|
token.cancel()
|
|
assert token.reason is None
|
|
|
|
def test_wait_returns_true_when_cancelled(self) -> None:
|
|
"""wait() 在已取消时应立即返回 True."""
|
|
token = CancelToken()
|
|
token.cancel()
|
|
assert token.wait(timeout=0.1)
|
|
|
|
def test_wait_timeout_returns_false(self) -> None:
|
|
"""wait() 在超时未取消时应返回 False."""
|
|
token = CancelToken()
|
|
assert not token.wait(timeout=0.05)
|
|
|
|
def test_wait_blocks_until_cancelled(self) -> None:
|
|
"""wait() 应阻塞直到另一线程调用 cancel()."""
|
|
|
|
token = CancelToken()
|
|
result: list[bool] = []
|
|
|
|
def _waiter() -> None:
|
|
result.append(token.wait(timeout=2.0))
|
|
|
|
t = threading.Thread(target=_waiter)
|
|
t.start()
|
|
time.sleep(0.05) # 等待 waiter 开始等待
|
|
token.cancel()
|
|
t.join(timeout=1.0)
|
|
assert result == [True]
|
|
|
|
def test_repr_active(self) -> None:
|
|
"""repr 在活跃状态应显示 active."""
|
|
token = CancelToken()
|
|
assert "active" in repr(token)
|
|
|
|
def test_repr_cancelled_with_reason(self) -> None:
|
|
"""repr 在取消状态应显示 cancelled 和 reason."""
|
|
token = CancelToken()
|
|
token.cancel("测试原因")
|
|
r = repr(token)
|
|
assert "cancelled" in r
|
|
assert "测试原因" in r
|
|
|
|
def test_link_threading_event(self) -> None:
|
|
"""link() 应在 threading.Event 设置时联动取消 token."""
|
|
token = CancelToken()
|
|
event = threading.Event()
|
|
token.link(event)
|
|
event.set()
|
|
time.sleep(0.1) # 等待 watcher 线程响应
|
|
assert token.is_cancelled
|
|
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# _is_cancelled 辅助函数测试
|
|
# ---------------------------------------------------------------------- #
|
|
class TestIsCancelled:
|
|
"""测试 _is_cancelled 统一检查函数."""
|
|
|
|
def test_none_returns_false(self) -> None:
|
|
"""None 应返回 False."""
|
|
assert not _is_cancelled(None)
|
|
|
|
def test_cancel_token_not_cancelled(self) -> None:
|
|
"""未取消的 CancelToken 应返回 False."""
|
|
token = CancelToken()
|
|
assert not _is_cancelled(token)
|
|
|
|
def test_cancel_token_cancelled(self) -> None:
|
|
"""已取消的 CancelToken 应返回 True."""
|
|
token = CancelToken()
|
|
token.cancel()
|
|
assert _is_cancelled(token)
|
|
|
|
def test_threading_event_not_set(self) -> None:
|
|
"""未设置的 threading.Event 应返回 False."""
|
|
event = threading.Event()
|
|
assert not _is_cancelled(event)
|
|
|
|
def test_threading_event_set(self) -> None:
|
|
"""已设置的 threading.Event 应返回 True."""
|
|
event = threading.Event()
|
|
event.set()
|
|
assert _is_cancelled(event)
|
|
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# _mark_remaining_skipped 测试
|
|
# ---------------------------------------------------------------------- #
|
|
class TestMarkRemainingSkipped:
|
|
"""测试 _mark_remaining_skipped 函数."""
|
|
|
|
def test_marks_all_unexecuted_tasks(self) -> None:
|
|
"""未执行的任务应被标记为 SKIPPED."""
|
|
graph = px.Graph.from_specs(
|
|
[
|
|
px.TaskSpec("a", lambda: 1),
|
|
px.TaskSpec("b", lambda: 2, depends_on=("a",)),
|
|
px.TaskSpec("c", lambda: 3, depends_on=("b",)),
|
|
]
|
|
)
|
|
report = RunReport()
|
|
# 模拟 a 已完成
|
|
from datetime import datetime
|
|
|
|
from pyflowx.task import TaskResult
|
|
|
|
report.results["a"] = TaskResult(
|
|
spec=graph.spec("a"),
|
|
status=TaskStatus.SUCCESS,
|
|
value=1,
|
|
finished_at=datetime.now(),
|
|
)
|
|
|
|
_mark_remaining_skipped(graph, report, None)
|
|
|
|
assert "a" in report.results
|
|
assert report.results["a"].status == TaskStatus.SUCCESS
|
|
assert report.results["b"].status == TaskStatus.SKIPPED
|
|
assert report.results["b"].reason == "执行被取消"
|
|
assert report.results["c"].status == TaskStatus.SKIPPED
|
|
|
|
def test_marks_all_when_none_executed(self) -> None:
|
|
"""无已完成任务时，所有任务应被标记为 SKIPPED."""
|
|
graph = px.Graph.from_specs(
|
|
[
|
|
px.TaskSpec("a", lambda: 1),
|
|
px.TaskSpec("b", lambda: 2, depends_on=("a",)),
|
|
]
|
|
)
|
|
report = RunReport()
|
|
_mark_remaining_skipped(graph, report, None)
|
|
assert len(report.results) == 2
|
|
assert report.results["a"].status == TaskStatus.SKIPPED
|
|
assert report.results["b"].status == TaskStatus.SKIPPED
|
|
|
|
def test_no_op_when_all_executed(self) -> None:
|
|
"""所有任务已完成时不应添加任何结果."""
|
|
from datetime import datetime
|
|
|
|
from pyflowx.task import TaskResult
|
|
|
|
graph = px.Graph.from_specs([px.TaskSpec("a", lambda: 1)])
|
|
report = RunReport()
|
|
report.results["a"] = TaskResult(
|
|
spec=graph.spec("a"),
|
|
status=TaskStatus.SUCCESS,
|
|
value=1,
|
|
finished_at=datetime.now(),
|
|
)
|
|
_mark_remaining_skipped(graph, report, None)
|
|
assert len(report.results) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# 外部取消集成测试
|
|
# ---------------------------------------------------------------------- #
|
|
def _make_cancellable_graph() -> px.Graph:
|
|
"""构建一个可被取消的图：a → b → c → d，每步有延迟."""
|
|
|
|
def a() -> int:
|
|
time.sleep(0.05)
|
|
return 1
|
|
|
|
def b(a: int) -> int:
|
|
time.sleep(0.05)
|
|
return a + 1
|
|
|
|
def c(b: int) -> int:
|
|
time.sleep(0.05)
|
|
return b + 1
|
|
|
|
def d(c: int) -> int:
|
|
time.sleep(0.05)
|
|
return c + 1
|
|
|
|
return px.Graph.from_specs(
|
|
[
|
|
px.TaskSpec("a", a),
|
|
px.TaskSpec("b", b, depends_on=("a",)),
|
|
px.TaskSpec("c", c, depends_on=("b",)),
|
|
px.TaskSpec("d", d, depends_on=("c",)),
|
|
]
|
|
)
|
|
|
|
|
|
class TestExternalCancel:
|
|
"""测试通过 CancelToken 外部取消 run()."""
|
|
|
|
@pytest.mark.parametrize("strategy", ["sequential", "thread", "async", "dependency"])
|
|
def test_cancel_before_run(self, strategy: str) -> None:
|
|
"""运行前取消：所有任务应标记为 SKIPPED."""
|
|
graph = _make_cancellable_graph()
|
|
token = CancelToken()
|
|
token.cancel()
|
|
|
|
report = px.run(graph, strategy=strategy, cancel_event=token) # type: ignore[arg-type]
|
|
|
|
assert not report.success
|
|
assert len(report.results) == 4
|
|
for name in ("a", "b", "c", "d"):
|
|
assert report.result_of(name).status == TaskStatus.SKIPPED
|
|
|
|
@pytest.mark.parametrize("strategy", ["sequential", "thread", "async", "dependency"])
|
|
def test_cancel_during_run(self, strategy: str) -> None:
|
|
"""运行中取消：已完成任务保持 SUCCESS，剩余任务标记为 SKIPPED."""
|
|
graph = _make_cancellable_graph()
|
|
token = CancelToken()
|
|
events: list[str] = []
|
|
|
|
def _on_event(event: px.TaskEvent) -> None:
|
|
events.append(event.task)
|
|
# a 完成后触发取消
|
|
if event.task == "a" and event.status == TaskStatus.SUCCESS:
|
|
token.cancel()
|
|
|
|
report = px.run(graph, strategy=strategy, cancel_event=token, on_event=_on_event) # type: ignore[arg-type]
|
|
|
|
assert not report.success
|
|
# a 应该成功
|
|
assert report.result_of("a").status == TaskStatus.SUCCESS
|
|
# 取消后未运行的任务应为 SKIPPED（d 一定未运行）
|
|
assert report.result_of("d").status == TaskStatus.SKIPPED
|
|
|
|
def test_cancel_with_threading_event(self) -> None:
|
|
"""threading.Event 也可作为取消信号."""
|
|
graph = _make_cancellable_graph()
|
|
event = threading.Event()
|
|
|
|
def _cancel_after_a() -> None:
|
|
time.sleep(0.08)
|
|
event.set()
|
|
|
|
t = threading.Thread(target=_cancel_after_a)
|
|
t.start()
|
|
report = px.run(graph, strategy="sequential", cancel_event=event) # type: ignore[arg-type]
|
|
t.join()
|
|
|
|
assert not report.success
|
|
assert report.result_of("a").status == TaskStatus.SUCCESS
|
|
|
|
def test_cancel_preserves_completed_results(self) -> None:
|
|
"""取消后已完成任务的结果应保留在 report 中."""
|
|
graph = _make_cancellable_graph()
|
|
token = CancelToken()
|
|
|
|
def _on_event(event: px.TaskEvent) -> None:
|
|
if event.task == "b" and event.status == TaskStatus.SUCCESS:
|
|
token.cancel()
|
|
|
|
report = px.run(graph, strategy="sequential", cancel_event=token, on_event=_on_event) # type: ignore[arg-type]
|
|
|
|
assert report.result_of("a").status == TaskStatus.SUCCESS
|
|
assert report.result_of("a").value == 1
|
|
assert report.result_of("b").status == TaskStatus.SUCCESS
|
|
assert report.result_of("b").value == 2
|
|
|
|
def test_no_cancel_returns_normal_report(self) -> None:
|
|
"""不传 cancel_event 时应正常完成."""
|
|
graph = _make_cancellable_graph()
|
|
report = px.run(graph, strategy="sequential")
|
|
assert report.success
|
|
assert report.result_of("d").value == 4
|
|
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# KeyboardInterrupt 测试
|
|
# ---------------------------------------------------------------------- #
|
|
class TestKeyboardInterrupt:
|
|
"""测试 KeyboardInterrupt 优雅处理."""
|
|
|
|
def test_keyboard_interrupt_marks_remaining_skipped(self) -> None:
|
|
"""KeyboardInterrupt 时未完成任务应标记为 SKIPPED."""
|
|
from unittest.mock import patch
|
|
|
|
from pyflowx import executors
|
|
|
|
graph = _make_cancellable_graph()
|
|
|
|
def _interrupt_after_first_layer(
|
|
graph: Any, layers: Any, context: Any, report: Any, backend: Any, on_event: Any, cancel_event: Any = None
|
|
) -> None:
|
|
# 只运行第一层，然后抛 KeyboardInterrupt
|
|
executors.SequentialLayerRunner.execute(layers[0], graph, context, report, backend, 1, on_event)
|
|
raise KeyboardInterrupt
|
|
|
|
with patch.object(executors, "_drive_sequential", _interrupt_after_first_layer):
|
|
report = px.run(graph, strategy="sequential")
|
|
|
|
assert not report.success
|
|
assert report.result_of("a").status == TaskStatus.SUCCESS
|
|
# b/c/d 应被标记为 SKIPPED
|
|
assert report.result_of("d").status == TaskStatus.SKIPPED
|
|
|
|
def test_keyboard_interrupt_with_dependency_strategy(self) -> None:
|
|
"""dependency 策略下 KeyboardInterrupt 也应优雅处理."""
|
|
from unittest.mock import patch
|
|
|
|
from pyflowx import executors
|
|
|
|
def fn() -> int:
|
|
return 1
|
|
|
|
graph = px.Graph.from_specs(
|
|
[
|
|
px.TaskSpec("a", fn),
|
|
px.TaskSpec("b", fn, depends_on=("a",)),
|
|
]
|
|
)
|
|
|
|
async def _interrupt(*args: Any, **kwargs: Any) -> None:
|
|
raise KeyboardInterrupt
|
|
|
|
with patch.object(executors.DependencyRunner, "execute", _interrupt):
|
|
report = px.run(graph, strategy="dependency")
|
|
|
|
assert not report.success
|