chongming-utils-task/tests/unit/test_cancellation.py

"""协作式取消令牌的行为测试。

这里必须用**真实线程**：令牌的价值恰恰在于跨线程正确握手（任务线程阻塞在
暂停点、调度线程请求停止并唤醒它）。用假的线程模型测不出竞态。

因此这些用例带小超时并断言线程确实退出——**同时也在验证没有死锁**。
"""

from __future__ import annotations

import asyncio
import contextlib
import threading

import pytest

from chongming_task.errors import TaskCancelled
from chongming_task.runtime.cancellation import CancellationToken

_JOIN_TIMEOUT = 5.0


def _run_in_thread(target: object) -> threading.Thread:
    thread = threading.Thread(target=target)  # type: ignore[arg-type]
    thread.daemon = True
    thread.start()
    return thread


# ---------------------------------------------------------------------------
# 停止
# ---------------------------------------------------------------------------


def test_checkpoint_is_noop_without_any_request() -> None:
    """没有请求时 checkpoint 必须立即返回，否则每轮循环都在付代价。"""
    CancellationToken().checkpoint()


def test_checkpoint_raises_when_stop_requested() -> None:
    token = CancellationToken()
    token.request_stop("发布窗口")
    with pytest.raises(TaskCancelled) as info:
        token.checkpoint()
    assert info.value.reason == "发布窗口"


def test_stop_is_idempotent_and_keeps_the_first_reason() -> None:
    """首个原因通常最接近真实起因，后续请求不应覆盖它。"""
    token = CancellationToken()
    token.request_stop("first")
    token.request_stop("second")
    with pytest.raises(TaskCancelled) as info:
        token.checkpoint()
    assert info.value.reason == "first"


def test_raise_if_stopped_ignores_pause() -> None:
    """用于清理路径：只响应停止，绝不阻塞。"""
    token = CancellationToken()
    token.request_pause()
    token.raise_if_stopped()  # 不应阻塞
    token.request_stop()
    with pytest.raises(TaskCancelled):
        token.raise_if_stopped()


def test_is_stop_requested_reflects_state() -> None:
    token = CancellationToken()
    assert not token.is_stop_requested
    token.request_stop()
    assert token.is_stop_requested


# ---------------------------------------------------------------------------
# 暂停
# ---------------------------------------------------------------------------


def test_pause_blocks_until_resume() -> None:
    token = CancellationToken()
    token.request_pause()  # 先请求，消除"任务先跑到检查点"的竞态
    observed: list[str] = []

    def worker() -> None:
        token.checkpoint()
        observed.append("resumed")

    thread = _run_in_thread(worker)
    assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True
    assert observed == [], "暂停请求生效前任务不应继续执行"

    token.resume()
    thread.join(_JOIN_TIMEOUT)
    assert not thread.is_alive(), "resume() 必须唤醒阻塞中的任务"
    assert observed == ["resumed"]


def test_stop_wakes_a_paused_task() -> None:
    """暂停中的任务必须能被停止，否则"暂停"就成了不可撤销的挂起。"""
    token = CancellationToken()
    token.request_pause()
    reasons: list[str | None] = []

    def worker() -> None:
        try:
            token.checkpoint()
        except TaskCancelled as exc:
            reasons.append(exc.reason)

    thread = _run_in_thread(worker)
    assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True

    token.request_stop("运维介入")
    thread.join(_JOIN_TIMEOUT)
    assert not thread.is_alive(), "停止请求必须能唤醒暂停点"
    assert reasons == ["运维介入"]


def test_repeated_pause_does_not_clear_acknowledgement() -> None:
    """重复请求暂停不应把"已确认"重置掉，否则 wait_paused 会误判为未生效。"""
    token = CancellationToken()
    token.request_pause()

    def worker() -> None:
        token.checkpoint()

    thread = _run_in_thread(worker)
    assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True
    token.request_pause()
    assert token.is_pause_acknowledged is True

    token.resume()
    thread.join(_JOIN_TIMEOUT)


def test_wait_paused_times_out_when_task_never_checkpoints() -> None:
    """调度器必须能区分"已暂停"与"暂停没生效"——后者要回滚为 RUNNING 并告警。"""
    token = CancellationToken()
    token.request_pause()
    assert token.wait_paused(timeout=0.05) is False
    assert token.is_pause_acknowledged is False


def test_wait_paused_reports_false_when_stopped_instead() -> None:
    """请求暂停后又被停止：这不是"暂停成功"，必须返回 False。"""
    token = CancellationToken()
    token.request_pause()
    token.request_stop()
    assert token.wait_paused(timeout=_JOIN_TIMEOUT) is False


def test_resume_without_pause_is_noop() -> None:
    token = CancellationToken()
    token.resume()
    assert not token.is_pause_requested
    token.checkpoint()  # 仍应立即返回


def _start_pausing_worker(token: CancellationToken) -> threading.Event:
    """起一个持续调用 ``checkpoint()`` 的工作线程，返回"完成过一轮"的事件。

    线程必须**持续**调用 checkpoint——只调一次就退出的话，之后的暂停请求自然没人
    确认，那是测试写错而不是缺陷。
    """

    def worker() -> None:
        # 收尾时的 stop 会让 checkpoint 抛 TaskCancelled：那正是预期行为，
        # 但线程里未捕获的异常会被 pytest 转成错误。
        with contextlib.suppress(TaskCancelled):
            while True:
                token.checkpoint()
                ran.set()

    ran = threading.Event()
    threading.Thread(target=worker, daemon=True).start()
    return ran


def test_pause_is_acknowledged_again_after_a_resume() -> None:
    """**回归用例**：暂停 → 恢复 → 再暂停里，第二次暂停曾经永久失效。

    ``checkpoint()`` 原先只在**进入**暂停分支时设置 ``_ack``。``resume()`` 把
    ``_pause`` 置回 False 并通知后，任务线程还没来得及重新取到锁，新的
    ``request_pause()`` 就把 ``_pause`` 又置成了 True——线程重新检查循环条件时发现
    "还要暂停"，于是继续等待，而它已经越过了设置 ``_ack`` 的那一行。

    结果 ``wait_paused()`` 永远等不到确认，暂停超时回滚。表现出来就是"暂停一次成功，
    之后再也暂停不了"。
    """
    token = CancellationToken()
    ran = _start_pausing_worker(token)
    assert ran.wait(_JOIN_TIMEOUT), "工作线程没有跑起来"

    for round_index in range(5):
        token.request_pause()
        assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True, (
            f"第 {round_index + 1} 轮暂停没有得到确认"
        )

        ran.clear()
        token.resume()
        assert ran.wait(_JOIN_TIMEOUT), f"第 {round_index + 1} 轮恢复后任务没有继续推进"

    token.request_stop()


def test_a_pause_right_after_a_resume_is_still_acknowledged() -> None:
    """恢复与再次暂停之间不留任何间隔——最容易踩中的时序。"""
    token = CancellationToken()
    ran = _start_pausing_worker(token)
    assert ran.wait(_JOIN_TIMEOUT)

    token.request_pause()
    assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True

    # 立刻恢复再立刻暂停：不给任务线程任何重新取锁的机会。
    token.resume()
    token.request_pause()
    assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True, "紧接着的第二次暂停丢了"

    token.resume()
    token.request_stop()


# ---------------------------------------------------------------------------
# 异步路径
# ---------------------------------------------------------------------------


def test_acheckpoint_fast_path_does_not_offload_to_thread() -> None:
    """紧密循环里的每次检查都丢线程会拖垮事件循环，快速路径必须留在本线程。"""

    async def main() -> str:
        token = CancellationToken()
        for _ in range(2000):
            await token.acheckpoint()
        return "done"

    assert asyncio.run(main()) == "done"


def test_acheckpoint_raises_when_stopped() -> None:
    token = CancellationToken()
    token.request_stop("已取消")

    async def main() -> None:
        await token.acheckpoint()

    with pytest.raises(TaskCancelled) as info:
        asyncio.run(main())
    assert info.value.reason == "已取消"


def test_acheckpoint_suspends_and_resumes_without_blocking_loop() -> None:
    """暂停期间事件循环必须仍能推进——这正是异步版本存在的理由。"""
    token = CancellationToken()
    token.request_pause()

    async def main() -> list[str]:
        events: list[str] = []

        async def heartbeat() -> None:
            while True:
                events.append("tick")
                await asyncio.sleep(0.01)

        ticker = asyncio.create_task(heartbeat())
        waiter = asyncio.create_task(token.acheckpoint())

        while not token.is_pause_acknowledged:
            await asyncio.sleep(0.01)
        ticks_while_paused = len(events)
        assert ticks_while_paused > 0, "暂停期间事件循环不应被阻塞"

        token.resume()
        await asyncio.wait_for(waiter, timeout=_JOIN_TIMEOUT)
        ticker.cancel()
        return events

    assert asyncio.run(main())