1 22
from contextlib import contextmanager
2

3 22
from .. import _core
4

5

6 22
@contextmanager
7 9
def _assert_yields_or_not(expected):
8 22
    __tracebackhide__ = True
9 22
    task = _core.current_task()
10 22
    orig_cancel = task._cancel_points
11 22
    orig_schedule = task._schedule_points
12 22
    try:
13 22
        yield
14 22
        if expected and (
15
            task._cancel_points == orig_cancel or task._schedule_points == orig_schedule
16
        ):
17 22
            raise AssertionError("assert_checkpoints block did not yield!")
18
    finally:
19 22
        if not expected and (
20
            task._cancel_points != orig_cancel or task._schedule_points != orig_schedule
21
        ):
22 22
            raise AssertionError("assert_no_checkpoints block yielded!")
23

24

25 22
def assert_checkpoints():
26
    """Use as a context manager to check that the code inside the ``with``
27
    block either exits with an exception or executes at least one
28
    :ref:`checkpoint <checkpoints>`.
29

30
    Raises:
31
      AssertionError: if no checkpoint was executed.
32

33
    Example:
34
      Check that :func:`trio.sleep` is a checkpoint, even if it doesn't
35
      block::
36

37
         with trio.testing.assert_checkpoints():
38
             await trio.sleep(0)
39

40
    """
41 22
    __tracebackhide__ = True
42 22
    return _assert_yields_or_not(True)
43

44

45 22
def assert_no_checkpoints():
46
    """Use as a context manager to check that the code inside the ``with``
47
    block does not execute any :ref:`checkpoints <checkpoints>`.
48

49
    Raises:
50
      AssertionError: if a checkpoint was executed.
51

52
    Example:
53
      Synchronous code never contains any checkpoints, but we can double-check
54
      that::
55

56
         send_channel, receive_channel = trio.open_memory_channel(10)
57
         with trio.testing.assert_no_checkpoints():
58
             send_channel.send_nowait(None)
59

60
    """
61 22
    __tracebackhide__ = True
62 22
    return _assert_yields_or_not(False)

Read our documentation on viewing source code .

Loading