You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

63 lines
1.8 KiB
Python

from contextlib import contextmanager
from .. import _core
@contextmanager
def _assert_yields_or_not(expected):
__tracebackhide__ = True
task = _core.current_task()
orig_cancel = task._cancel_points
orig_schedule = task._schedule_points
try:
yield
if expected and (
task._cancel_points == orig_cancel or task._schedule_points == orig_schedule
):
raise AssertionError("assert_checkpoints block did not yield!")
finally:
if not expected and (
task._cancel_points != orig_cancel or task._schedule_points != orig_schedule
):
raise AssertionError("assert_no_checkpoints block yielded!")
def assert_checkpoints():
"""Use as a context manager to check that the code inside the ``with``
block either exits with an exception or executes at least one
:ref:`checkpoint <checkpoints>`.
Raises:
AssertionError: if no checkpoint was executed.
Example:
Check that :func:`trio.sleep` is a checkpoint, even if it doesn't
block::
with trio.testing.assert_checkpoints():
await trio.sleep(0)
"""
__tracebackhide__ = True
return _assert_yields_or_not(True)
def assert_no_checkpoints():
"""Use as a context manager to check that the code inside the ``with``
block does not execute any :ref:`checkpoints <checkpoints>`.
Raises:
AssertionError: if a checkpoint was executed.
Example:
Synchronous code never contains any checkpoints, but we can double-check
that::
send_channel, receive_channel = trio.open_memory_channel(10)
with trio.testing.assert_no_checkpoints():
send_channel.send_nowait(None)
"""
__tracebackhide__ = True
return _assert_yields_or_not(False)