chakcy
f4de43c743
单任务受控执行（15 态状态机 + 审计轨迹）、DAG 管道、Controller 门面。
执行器三个：线程（默认，含 async def 任务）、内联（同步 / 测试）、进程
（spawn 子进程：硬性停止、绕开 GIL、崩溃隔离）。
调度：并发额度 + 每桶限流 + 具名资源信号量（按标签）+ 优先级排队。
持久层：TaskStore 契约 + 内存 / SQLite / Postgres 实现，契约测试随包发布。
跨进程与多节点：写穿 / 读穿、指令通道与租约、owned_buckets 划界、reconcile。
CLI、12 个可直接运行的示例、docs/ 下 5 份文档。
门禁：ruff / ruff format / mypy(strict) 干净；
pytest 563 passed（44 skipped 是需要真实 Postgres 的那批）；
带一台可连的 Postgres 时 607 passed。
440 lines
14 KiB
Python
440 lines
14 KiB
Python
"""``@task`` 装饰器与 ``TaskContext`` 的能力测试。
|
|
|
|
三块：
|
|
|
|
1. **``@task``**——元数据从装饰器来；不产生全局状态（导入无副作用）；
|
|
2. **进度**——进快照、进存储、**不碰版本号**；
|
|
3. **上游数据**——管道节点之间真正能传值，且失败的上游不传 ``None``。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from chongming_task.controller import Controller, TaskRegistry
|
|
from chongming_task.decorators import get_task_spec, task
|
|
from chongming_task.errors import ConfigurationError
|
|
from chongming_task.pipeline import Pipeline
|
|
from chongming_task.policies import RetryPolicy, Timeout
|
|
from chongming_task.states import TaskStatus
|
|
from chongming_task.stores.memory import InMemoryTaskStore
|
|
|
|
_WINDOW = 15.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. @task
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@task
|
|
def bare(ctx: Any) -> str:
|
|
"""裸装饰器。"""
|
|
del ctx
|
|
return "bare"
|
|
|
|
|
|
@task(tags={"io", "nightly"}, retry=RetryPolicy(max_attempts=4), timeout=Timeout(per_attempt=9.0))
|
|
def decorated(ctx: Any) -> str:
|
|
"""带元数据的任务。
|
|
|
|
第二行不该出现在描述里。
|
|
"""
|
|
del ctx
|
|
return "decorated"
|
|
|
|
|
|
@task(name="renamed")
|
|
def original_name(ctx: Any) -> str:
|
|
del ctx
|
|
return "renamed"
|
|
|
|
|
|
def undecorated(ctx: Any) -> str:
|
|
del ctx
|
|
return "nope"
|
|
|
|
|
|
def test_bare_decorator_works() -> None:
|
|
"""``@task`` 不带括号也太自然，不支持会让人反复踩空。"""
|
|
spec = get_task_spec(bare)
|
|
assert spec is not None
|
|
assert spec.name == "bare"
|
|
assert spec.description == "裸装饰器。"
|
|
|
|
|
|
def test_decorator_captures_metadata() -> None:
|
|
spec = get_task_spec(decorated)
|
|
assert spec is not None
|
|
assert spec.tags == frozenset({"io", "nightly"})
|
|
assert spec.retry is not None
|
|
assert spec.retry.max_attempts == 4
|
|
assert spec.timeout is not None
|
|
assert spec.timeout.per_attempt == 9.0
|
|
|
|
|
|
def test_description_defaults_to_the_first_docstring_line() -> None:
|
|
spec = get_task_spec(decorated)
|
|
assert spec is not None
|
|
assert spec.description == "带元数据的任务。", "只取第一行"
|
|
assert "第二行" not in spec.description
|
|
|
|
|
|
def test_name_override() -> None:
|
|
spec = get_task_spec(original_name)
|
|
assert spec is not None
|
|
assert spec.name == "renamed"
|
|
|
|
|
|
def test_undecorated_function_has_no_spec() -> None:
|
|
assert get_task_spec(undecorated) is None
|
|
|
|
|
|
def test_decorating_does_not_create_global_state() -> None:
|
|
"""规格只挂在函数对象上。
|
|
|
|
全局注册表意味着 **import 即注册**，而 import 可能发生在任何地方（包括测试收集
|
|
阶段），那种隐式行为极难排查。
|
|
"""
|
|
registry = TaskRegistry()
|
|
assert registry.names() == ()
|
|
assert get_task_spec(bare) is not None, "规格依然取得到，但没有登记进任何注册表"
|
|
|
|
|
|
def test_registry_registers_decorated_functions() -> None:
|
|
registry = TaskRegistry()
|
|
registration = registry.register_decorated(decorated)
|
|
assert registration.name == "decorated"
|
|
assert registration.tags == frozenset({"io", "nightly"})
|
|
assert registration.retry.max_attempts == 4
|
|
|
|
|
|
def test_registry_rejects_an_undecorated_function() -> None:
|
|
registry = TaskRegistry()
|
|
with pytest.raises(ConfigurationError, match="没有用 @task 装饰过"):
|
|
registry.register_decorated(undecorated)
|
|
|
|
|
|
def test_register_module_picks_up_every_decorated_function() -> None:
|
|
import sys
|
|
|
|
module = sys.modules[__name__]
|
|
registry = TaskRegistry()
|
|
registered = registry.register_module(module)
|
|
|
|
names = {item.name for item in registered}
|
|
assert {"bare", "decorated", "renamed"} <= names
|
|
assert "undecorated" not in names
|
|
|
|
|
|
def test_register_module_rejects_a_module_with_nothing_decorated() -> None:
|
|
import json
|
|
|
|
registry = TaskRegistry()
|
|
with pytest.raises(ConfigurationError, match="没有用 @task 装饰过的函数"):
|
|
registry.register_module(json)
|
|
|
|
|
|
def test_controller_registers_a_module_end_to_end() -> None:
|
|
import sys
|
|
|
|
controller = Controller(max_parallel=2)
|
|
controller.register_module(sys.modules[__name__])
|
|
run = controller.submit("bare")
|
|
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
|
|
assert outcome.ok
|
|
assert outcome.value == "bare"
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
def test_decorated_retry_actually_applies() -> None:
|
|
"""装饰器上写的重试次数必须真的生效，而不只是躺在元数据里。"""
|
|
store = InMemoryTaskStore()
|
|
controller = Controller(store=store, max_parallel=2)
|
|
registry = TaskRegistry()
|
|
registry.register(
|
|
"flaky",
|
|
_flaky,
|
|
retry=RetryPolicy(max_attempts=3, backoff=_no_backoff()),
|
|
)
|
|
controller.register("flaky", _flaky, retry=RetryPolicy(max_attempts=3, backoff=_no_backoff()))
|
|
run = controller.submit("flaky")
|
|
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
|
|
assert outcome.ok, "第三次应成功"
|
|
snapshot = store.get(run.snapshot.task_id)
|
|
assert snapshot is not None
|
|
assert snapshot.attempt == 3
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. TaskContext：进度与日志
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_progress_reaches_the_snapshot_without_bumping_the_version() -> None:
|
|
"""进度若参与乐观并发，会与状态写入持续冲突。"""
|
|
controller = Controller(max_parallel=2)
|
|
controller.register("reporter", _reporter)
|
|
run = controller.submit("reporter")
|
|
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
|
|
assert outcome.ok
|
|
|
|
snapshot = outcome.snapshot
|
|
assert snapshot.progress is not None
|
|
assert snapshot.progress.done == 3
|
|
assert snapshot.progress.total == 3
|
|
assert snapshot.progress.unit == "rows"
|
|
# 状态迁移只有 4 次（queued/starting/running/succeeded），进度不该挤进去。
|
|
assert snapshot.status_version == 4
|
|
|
|
|
|
def test_progress_is_persisted_and_visible_in_queries() -> None:
|
|
"""进度要经 touch 旁路落到存储，否则看板的 PROGRESS 列永远是空的。"""
|
|
store = InMemoryTaskStore()
|
|
controller = Controller(store=store, max_parallel=2, command_poll_interval=0.05)
|
|
controller.register("slow_reporter", _slow_reporter)
|
|
run = controller.submit("slow_reporter")
|
|
|
|
assert _wait_until(
|
|
lambda: (store.get(run.snapshot.task_id) or _stub()).progress is not None,
|
|
timeout=_WINDOW,
|
|
)
|
|
persisted = store.get(run.snapshot.task_id)
|
|
assert persisted is not None
|
|
assert persisted.progress is not None
|
|
assert persisted.progress.done >= 1
|
|
|
|
controller.stop(run.snapshot.task_id)
|
|
controller.wait_all(timeout=_WINDOW)
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
def test_progress_on_an_unsupported_context_is_silently_dropped() -> None:
|
|
"""没有回写通道时静默丢弃，而不是抛异常。
|
|
|
|
直接构造上下文（不经过执行器）是测试与工具里的常见用法，不该因此报错。
|
|
"""
|
|
from chongming_task.runtime.cancellation import CancellationToken
|
|
from chongming_task.runtime.context import RuntimeContext
|
|
from chongming_task.runtime.heartbeat import Heartbeat
|
|
|
|
context = RuntimeContext(
|
|
task_id="t",
|
|
name="t",
|
|
attempt=1,
|
|
token=CancellationToken(),
|
|
heartbeat=Heartbeat(),
|
|
)
|
|
context.progress(1, total=2)
|
|
context.log("hello")
|
|
assert list(context.logs) == ["hello"]
|
|
|
|
|
|
def test_log_records_into_the_buffer_and_the_stdlib_logger(
|
|
caplog: pytest.LogCaptureFixture,
|
|
) -> None:
|
|
# caplog 默认只捕获 WARNING 及以上，而 INFO 才是我们最常用的级别。
|
|
caplog.set_level(logging.INFO, logger="chongming_task.task")
|
|
|
|
controller = Controller(max_parallel=2)
|
|
controller.register("talker", _talker)
|
|
run = controller.submit("talker")
|
|
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
|
|
assert outcome.ok
|
|
|
|
assert outcome.value == "done"
|
|
messages = [record.message for record in caplog.records]
|
|
assert "hello" in messages, "INFO 日志应当走到标准库 logging"
|
|
assert "careful" in messages
|
|
|
|
record = next(r for r in caplog.records if r.message == "hello")
|
|
assert record.task_id == run.snapshot.task_id
|
|
assert record.task_name == "talker"
|
|
assert record.attempt == 1
|
|
|
|
|
|
def test_task_context_is_the_runtime_context() -> None:
|
|
"""公开名字与内部名字是**同一个类**。
|
|
|
|
任务作者只面对一个 ``ctx``；两个类会立刻带来"我该用哪个"的困惑。
|
|
"""
|
|
from chongming_task import RuntimeContext, TaskContext
|
|
|
|
assert TaskContext is RuntimeContext
|
|
|
|
|
|
def test_checkpoint_still_works_alongside_progress() -> None:
|
|
controller = Controller(max_parallel=2)
|
|
controller.register("stopme", _stopme)
|
|
run = controller.submit("stopme")
|
|
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
|
|
|
|
controller.stop(run.snapshot.task_id, reason="收工")
|
|
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
|
|
assert outcome.status is TaskStatus.STOPPED
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. 上游数据（管道传值）
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_downstream_receives_upstream_values() -> None:
|
|
controller = Controller(max_parallel=4)
|
|
controller.register("produce", _produce)
|
|
controller.register("consume", _consume)
|
|
|
|
pipeline = Pipeline("data")
|
|
source = pipeline.node("produce", node_id="source")
|
|
pipeline.node("consume", node_id="sink", depends_on=[source])
|
|
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="d1")
|
|
outcome = run.wait(timeout=_WINDOW)
|
|
assert outcome.status.value == "succeeded"
|
|
assert outcome.nodes["sink"] is TaskStatus.SUCCEEDED
|
|
|
|
snapshot = controller.get("d1:sink")
|
|
assert snapshot is not None
|
|
assert snapshot.outcome == {"value": 42}, "下游应当拿到上游的返回值 42"
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
def test_a_failed_upstream_contributes_no_value() -> None:
|
|
"""失败的上游不塞 ``None``：那会让下游以为"上游返回了 None"。"""
|
|
controller = Controller(max_parallel=4)
|
|
controller.register("boom", _boom)
|
|
controller.register("inspect", _inspect)
|
|
|
|
pipeline = Pipeline("missing")
|
|
bad = pipeline.node("boom", node_id="bad")
|
|
pipeline.node("inspect", node_id="probe", depends_on=[bad], join=_all_settled())
|
|
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="d2")
|
|
run.wait(timeout=_WINDOW)
|
|
|
|
snapshot = controller.get("d2:probe")
|
|
assert snapshot is not None
|
|
assert snapshot.outcome == {"value": []}, "上游没有值，ctx.upstream 必须是空的"
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
def test_input_can_be_read_with_a_default() -> None:
|
|
controller = Controller()
|
|
controller.register("use_default", _use_default)
|
|
run = controller.submit("use_default")
|
|
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
|
|
assert outcome.value == "fallback"
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
def test_inputs_cannot_be_injected_after_the_task_started() -> None:
|
|
controller = Controller(max_parallel=2)
|
|
controller.register("stopme", _stopme)
|
|
run = controller.submit("stopme")
|
|
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
|
|
|
|
with pytest.raises(RuntimeError, match="已经开始执行"):
|
|
run.set_inputs({"x": 1})
|
|
|
|
controller.stop(run.snapshot.task_id)
|
|
controller.wait_all(timeout=_WINDOW)
|
|
controller.close(timeout=_WINDOW)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 任务体
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _stub() -> Any:
|
|
from chongming_task.models import TaskSnapshot
|
|
|
|
return TaskSnapshot(task_id="stub", name="stub")
|
|
|
|
|
|
def _no_backoff() -> Any:
|
|
from chongming_task.policies import FixedBackoff
|
|
|
|
return FixedBackoff(0.0)
|
|
|
|
|
|
def _flaky(ctx: Any) -> str:
|
|
state = getattr(_flaky, "_count", 0) + 1
|
|
_flaky._count = state # type: ignore[attr-defined]
|
|
if state < 3:
|
|
raise RuntimeError(f"第 {state} 次失败")
|
|
return "ok"
|
|
|
|
|
|
def _reporter(ctx: Any) -> str:
|
|
for index in range(1, 4):
|
|
ctx.checkpoint()
|
|
ctx.progress(index, total=3, unit="rows")
|
|
return "reported"
|
|
|
|
|
|
def _slow_reporter(ctx: Any) -> str:
|
|
for index in range(1, 200):
|
|
ctx.checkpoint()
|
|
ctx.progress(index, total=200, unit="rows")
|
|
time.sleep(0.01)
|
|
return "done"
|
|
|
|
|
|
def _talker(ctx: Any) -> str:
|
|
ctx.log("hello")
|
|
ctx.warning("careful")
|
|
return "done"
|
|
|
|
|
|
def _stopme(ctx: Any) -> str:
|
|
for _ in range(1_000_000):
|
|
ctx.checkpoint()
|
|
time.sleep(0.005)
|
|
return "never"
|
|
|
|
|
|
def _produce(ctx: Any) -> int:
|
|
del ctx
|
|
return 42
|
|
|
|
|
|
def _consume(ctx: Any) -> int:
|
|
value = ctx.input("source")
|
|
assert value == 42, f"预期 42，收到 {value!r}"
|
|
return int(value)
|
|
|
|
|
|
def _inspect(ctx: Any) -> list[str]:
|
|
return sorted(ctx.upstream)
|
|
|
|
|
|
def _use_default(ctx: Any) -> str:
|
|
return str(ctx.input("nobody", "fallback"))
|
|
|
|
|
|
def _boom(ctx: Any) -> str:
|
|
del ctx
|
|
raise RuntimeError("故意失败")
|
|
|
|
|
|
def _all_settled() -> Any:
|
|
from chongming_task.policies import JoinMode
|
|
|
|
return JoinMode.ALL_SETTLED
|
|
|
|
|
|
def _wait_until(predicate: Any, timeout: float = _WINDOW) -> bool:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(0.01)
|
|
return False
|