moving to scripts
This commit is contained in:
eneller
32
asq-env/lib/python3.9/site-packages/trio/testing/__init__.py
Normal file
32
asq-env/lib/python3.9/site-packages/trio/testing/__init__.py
Normal file
@@ -0,0 +1,32 @@
|
||||
from .._core import wait_all_tasks_blocked, MockClock
|
||||
|
||||
from ._trio_test import trio_test
|
||||
|
||||
from ._checkpoints import assert_checkpoints, assert_no_checkpoints
|
||||
|
||||
from ._sequencer import Sequencer
|
||||
|
||||
from ._check_streams import (
|
||||
check_one_way_stream,
|
||||
check_two_way_stream,
|
||||
check_half_closeable_stream,
|
||||
)
|
||||
|
||||
from ._memory_streams import (
|
||||
MemorySendStream,
|
||||
MemoryReceiveStream,
|
||||
memory_stream_pump,
|
||||
memory_stream_one_way_pair,
|
||||
memory_stream_pair,
|
||||
lockstep_stream_one_way_pair,
|
||||
lockstep_stream_pair,
|
||||
)
|
||||
|
||||
from ._network import open_stream_to_socket_listener
|
||||
|
||||
################################################################
|
||||
|
||||
from .._util import fixup_module_metadata
|
||||
|
||||
fixup_module_metadata(__name__, globals())
|
||||
del fixup_module_metadata
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,512 @@
|
||||
# Generic stream tests
|
||||
|
||||
from contextlib import contextmanager
|
||||
import random
|
||||
|
||||
from .. import _core
|
||||
from .._highlevel_generic import aclose_forcefully
|
||||
from .._abc import SendStream, ReceiveStream, Stream, HalfCloseableStream
|
||||
from ._checkpoints import assert_checkpoints
|
||||
|
||||
|
||||
class _ForceCloseBoth:
|
||||
def __init__(self, both):
|
||||
self._both = list(both)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._both
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
try:
|
||||
await aclose_forcefully(self._both[0])
|
||||
finally:
|
||||
await aclose_forcefully(self._both[1])
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _assert_raises(exc):
|
||||
__tracebackhide__ = True
|
||||
try:
|
||||
yield
|
||||
except exc:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("expected exception: {}".format(exc))
|
||||
|
||||
|
||||
async def check_one_way_stream(stream_maker, clogged_stream_maker):
|
||||
"""Perform a number of generic tests on a custom one-way stream
|
||||
implementation.
|
||||
|
||||
Args:
|
||||
stream_maker: An async (!) function which returns a connected
|
||||
(:class:`~trio.abc.SendStream`, :class:`~trio.abc.ReceiveStream`)
|
||||
pair.
|
||||
clogged_stream_maker: Either None, or an async function similar to
|
||||
stream_maker, but with the extra property that the returned stream
|
||||
is in a state where ``send_all`` and
|
||||
``wait_send_all_might_not_block`` will block until ``receive_some``
|
||||
has been called. This allows for more thorough testing of some edge
|
||||
cases, especially around ``wait_send_all_might_not_block``.
|
||||
|
||||
Raises:
|
||||
AssertionError: if a test fails.
|
||||
|
||||
"""
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
assert isinstance(s, SendStream)
|
||||
assert isinstance(r, ReceiveStream)
|
||||
|
||||
async def do_send_all(data):
|
||||
with assert_checkpoints():
|
||||
assert await s.send_all(data) is None
|
||||
|
||||
async def do_receive_some(*args):
|
||||
with assert_checkpoints():
|
||||
return await r.receive_some(*args)
|
||||
|
||||
async def checked_receive_1(expected):
|
||||
assert await do_receive_some(1) == expected
|
||||
|
||||
async def do_aclose(resource):
|
||||
with assert_checkpoints():
|
||||
await resource.aclose()
|
||||
|
||||
# Simple sending/receiving
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_send_all, b"x")
|
||||
nursery.start_soon(checked_receive_1, b"x")
|
||||
|
||||
async def send_empty_then_y():
|
||||
# Streams should tolerate sending b"" without giving it any
|
||||
# special meaning.
|
||||
await do_send_all(b"")
|
||||
await do_send_all(b"y")
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(send_empty_then_y)
|
||||
nursery.start_soon(checked_receive_1, b"y")
|
||||
|
||||
# ---- Checking various argument types ----
|
||||
|
||||
# send_all accepts bytearray and memoryview
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_send_all, bytearray(b"1"))
|
||||
nursery.start_soon(checked_receive_1, b"1")
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_send_all, memoryview(b"2"))
|
||||
nursery.start_soon(checked_receive_1, b"2")
|
||||
|
||||
# max_bytes must be a positive integer
|
||||
with _assert_raises(ValueError):
|
||||
await r.receive_some(-1)
|
||||
with _assert_raises(ValueError):
|
||||
await r.receive_some(0)
|
||||
with _assert_raises(TypeError):
|
||||
await r.receive_some(1.5)
|
||||
# it can also be missing or None
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_send_all, b"x")
|
||||
assert await do_receive_some() == b"x"
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_send_all, b"x")
|
||||
assert await do_receive_some(None) == b"x"
|
||||
|
||||
with _assert_raises(_core.BusyResourceError):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_receive_some, 1)
|
||||
nursery.start_soon(do_receive_some, 1)
|
||||
|
||||
# Method always has to exist, and an empty stream with a blocked
|
||||
# receive_some should *always* allow send_all. (Technically it's legal
|
||||
# for send_all to wait until receive_some is called to run, though; a
|
||||
# stream doesn't *have* to have any internal buffering. That's why we
|
||||
# start a concurrent receive_some call, then cancel it.)
|
||||
async def simple_check_wait_send_all_might_not_block(scope):
|
||||
with assert_checkpoints():
|
||||
await s.wait_send_all_might_not_block()
|
||||
scope.cancel()
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(
|
||||
simple_check_wait_send_all_might_not_block, nursery.cancel_scope
|
||||
)
|
||||
nursery.start_soon(do_receive_some, 1)
|
||||
|
||||
# closing the r side leads to BrokenResourceError on the s side
|
||||
# (eventually)
|
||||
async def expect_broken_stream_on_send():
|
||||
with _assert_raises(_core.BrokenResourceError):
|
||||
while True:
|
||||
await do_send_all(b"x" * 100)
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(expect_broken_stream_on_send)
|
||||
nursery.start_soon(do_aclose, r)
|
||||
|
||||
# once detected, the stream stays broken
|
||||
with _assert_raises(_core.BrokenResourceError):
|
||||
await do_send_all(b"x" * 100)
|
||||
|
||||
# r closed -> ClosedResourceError on the receive side
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_receive_some(4096)
|
||||
|
||||
# we can close the same stream repeatedly, it's fine
|
||||
await do_aclose(r)
|
||||
await do_aclose(r)
|
||||
|
||||
# closing the sender side
|
||||
await do_aclose(s)
|
||||
|
||||
# now trying to send raises ClosedResourceError
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_send_all(b"x" * 100)
|
||||
|
||||
# even if it's an empty send
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_send_all(b"")
|
||||
|
||||
# ditto for wait_send_all_might_not_block
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
with assert_checkpoints():
|
||||
await s.wait_send_all_might_not_block()
|
||||
|
||||
# and again, repeated closing is fine
|
||||
await do_aclose(s)
|
||||
await do_aclose(s)
|
||||
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
# if send-then-graceful-close, receiver gets data then b""
|
||||
async def send_then_close():
|
||||
await do_send_all(b"y")
|
||||
await do_aclose(s)
|
||||
|
||||
async def receive_send_then_close():
|
||||
# We want to make sure that if the sender closes the stream before
|
||||
# we read anything, then we still get all the data. But some
|
||||
# streams might block on the do_send_all call. So we let the
|
||||
# sender get as far as it can, then we receive.
|
||||
await _core.wait_all_tasks_blocked()
|
||||
await checked_receive_1(b"y")
|
||||
await checked_receive_1(b"")
|
||||
await do_aclose(r)
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(send_then_close)
|
||||
nursery.start_soon(receive_send_then_close)
|
||||
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
await aclose_forcefully(r)
|
||||
|
||||
with _assert_raises(_core.BrokenResourceError):
|
||||
while True:
|
||||
await do_send_all(b"x" * 100)
|
||||
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_receive_some(4096)
|
||||
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
await aclose_forcefully(s)
|
||||
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_send_all(b"123")
|
||||
|
||||
# after the sender does a forceful close, the receiver might either
|
||||
# get BrokenResourceError or a clean b""; either is OK. Not OK would be
|
||||
# if it freezes, or returns data.
|
||||
try:
|
||||
await checked_receive_1(b"")
|
||||
except _core.BrokenResourceError:
|
||||
pass
|
||||
|
||||
# cancelled aclose still closes
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
with _core.CancelScope() as scope:
|
||||
scope.cancel()
|
||||
await r.aclose()
|
||||
|
||||
with _core.CancelScope() as scope:
|
||||
scope.cancel()
|
||||
await s.aclose()
|
||||
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_send_all(b"123")
|
||||
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await do_receive_some(4096)
|
||||
|
||||
# Check that we can still gracefully close a stream after an operation has
|
||||
# been cancelled. This can be challenging if cancellation can leave the
|
||||
# stream internals in an inconsistent state, e.g. for
|
||||
# SSLStream. Unfortunately this test isn't very thorough; the really
|
||||
# challenging case for something like SSLStream is it gets cancelled
|
||||
# *while* it's sending data on the underlying, not before. But testing
|
||||
# that requires some special-case handling of the particular stream setup;
|
||||
# we can't do it here. Maybe we could do a bit better with
|
||||
# https://github.com/python-trio/trio/issues/77
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
|
||||
async def expect_cancelled(afn, *args):
|
||||
with _assert_raises(_core.Cancelled):
|
||||
await afn(*args)
|
||||
|
||||
with _core.CancelScope() as scope:
|
||||
scope.cancel()
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(expect_cancelled, do_send_all, b"x")
|
||||
nursery.start_soon(expect_cancelled, do_receive_some, 1)
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(do_aclose, s)
|
||||
nursery.start_soon(do_aclose, r)
|
||||
|
||||
# Check that if a task is blocked in receive_some, then closing the
|
||||
# receive stream causes it to wake up.
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s, r):
|
||||
|
||||
async def receive_expecting_closed():
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await r.receive_some(10)
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(receive_expecting_closed)
|
||||
await _core.wait_all_tasks_blocked()
|
||||
await aclose_forcefully(r)
|
||||
|
||||
# check wait_send_all_might_not_block, if we can
|
||||
if clogged_stream_maker is not None:
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
record = []
|
||||
|
||||
async def waiter(cancel_scope):
|
||||
record.append("waiter sleeping")
|
||||
with assert_checkpoints():
|
||||
await s.wait_send_all_might_not_block()
|
||||
record.append("waiter wokeup")
|
||||
cancel_scope.cancel()
|
||||
|
||||
async def receiver():
|
||||
# give wait_send_all_might_not_block a chance to block
|
||||
await _core.wait_all_tasks_blocked()
|
||||
record.append("receiver starting")
|
||||
while True:
|
||||
await r.receive_some(16834)
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(waiter, nursery.cancel_scope)
|
||||
await _core.wait_all_tasks_blocked()
|
||||
nursery.start_soon(receiver)
|
||||
|
||||
assert record == [
|
||||
"waiter sleeping",
|
||||
"receiver starting",
|
||||
"waiter wokeup",
|
||||
]
|
||||
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
# simultaneous wait_send_all_might_not_block fails
|
||||
with _assert_raises(_core.BusyResourceError):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(s.wait_send_all_might_not_block)
|
||||
nursery.start_soon(s.wait_send_all_might_not_block)
|
||||
|
||||
# and simultaneous send_all and wait_send_all_might_not_block (NB
|
||||
# this test might destroy the stream b/c we end up cancelling
|
||||
# send_all and e.g. SSLStream can't handle that, so we have to
|
||||
# recreate afterwards)
|
||||
with _assert_raises(_core.BusyResourceError):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(s.wait_send_all_might_not_block)
|
||||
nursery.start_soon(s.send_all, b"123")
|
||||
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
# send_all and send_all blocked simultaneously should also raise
|
||||
# (but again this might destroy the stream)
|
||||
with _assert_raises(_core.BusyResourceError):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(s.send_all, b"123")
|
||||
nursery.start_soon(s.send_all, b"123")
|
||||
|
||||
# closing the receiver causes wait_send_all_might_not_block to return,
|
||||
# with or without an exception
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
|
||||
async def sender():
|
||||
try:
|
||||
with assert_checkpoints():
|
||||
await s.wait_send_all_might_not_block()
|
||||
except _core.BrokenResourceError: # pragma: no cover
|
||||
pass
|
||||
|
||||
async def receiver():
|
||||
await _core.wait_all_tasks_blocked()
|
||||
await aclose_forcefully(r)
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(sender)
|
||||
nursery.start_soon(receiver)
|
||||
|
||||
# and again with the call starting after the close
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
await aclose_forcefully(r)
|
||||
try:
|
||||
with assert_checkpoints():
|
||||
await s.wait_send_all_might_not_block()
|
||||
except _core.BrokenResourceError: # pragma: no cover
|
||||
pass
|
||||
|
||||
# Check that if a task is blocked in a send-side method, then closing
|
||||
# the send stream causes it to wake up.
|
||||
async def close_soon(s):
|
||||
await _core.wait_all_tasks_blocked()
|
||||
await aclose_forcefully(s)
|
||||
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(close_soon, s)
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await s.send_all(b"xyzzy")
|
||||
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s, r):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(close_soon, s)
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await s.wait_send_all_might_not_block()
|
||||
|
||||
|
||||
async def check_two_way_stream(stream_maker, clogged_stream_maker):
|
||||
"""Perform a number of generic tests on a custom two-way stream
|
||||
implementation.
|
||||
|
||||
This is similar to :func:`check_one_way_stream`, except that the maker
|
||||
functions are expected to return objects implementing the
|
||||
:class:`~trio.abc.Stream` interface.
|
||||
|
||||
This function tests a *superset* of what :func:`check_one_way_stream`
|
||||
checks – if you call this, then you don't need to also call
|
||||
:func:`check_one_way_stream`.
|
||||
|
||||
"""
|
||||
await check_one_way_stream(stream_maker, clogged_stream_maker)
|
||||
|
||||
async def flipped_stream_maker():
|
||||
return reversed(await stream_maker())
|
||||
|
||||
if clogged_stream_maker is not None:
|
||||
|
||||
async def flipped_clogged_stream_maker():
|
||||
return reversed(await clogged_stream_maker())
|
||||
|
||||
else:
|
||||
flipped_clogged_stream_maker = None
|
||||
await check_one_way_stream(flipped_stream_maker, flipped_clogged_stream_maker)
|
||||
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s1, s2):
|
||||
assert isinstance(s1, Stream)
|
||||
assert isinstance(s2, Stream)
|
||||
|
||||
# Duplex can be a bit tricky, might as well check it as well
|
||||
DUPLEX_TEST_SIZE = 2 ** 20
|
||||
CHUNK_SIZE_MAX = 2 ** 14
|
||||
|
||||
r = random.Random(0)
|
||||
i = r.getrandbits(8 * DUPLEX_TEST_SIZE)
|
||||
test_data = i.to_bytes(DUPLEX_TEST_SIZE, "little")
|
||||
|
||||
async def sender(s, data, seed):
|
||||
r = random.Random(seed)
|
||||
m = memoryview(data)
|
||||
while m:
|
||||
chunk_size = r.randint(1, CHUNK_SIZE_MAX)
|
||||
await s.send_all(m[:chunk_size])
|
||||
m = m[chunk_size:]
|
||||
|
||||
async def receiver(s, data, seed):
|
||||
r = random.Random(seed)
|
||||
got = bytearray()
|
||||
while len(got) < len(data):
|
||||
chunk = await s.receive_some(r.randint(1, CHUNK_SIZE_MAX))
|
||||
assert chunk
|
||||
got += chunk
|
||||
assert got == data
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(sender, s1, test_data, 0)
|
||||
nursery.start_soon(sender, s2, test_data[::-1], 1)
|
||||
nursery.start_soon(receiver, s1, test_data[::-1], 2)
|
||||
nursery.start_soon(receiver, s2, test_data, 3)
|
||||
|
||||
async def expect_receive_some_empty():
|
||||
assert await s2.receive_some(10) == b""
|
||||
await s2.aclose()
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(expect_receive_some_empty)
|
||||
nursery.start_soon(s1.aclose)
|
||||
|
||||
|
||||
async def check_half_closeable_stream(stream_maker, clogged_stream_maker):
|
||||
"""Perform a number of generic tests on a custom half-closeable stream
|
||||
implementation.
|
||||
|
||||
This is similar to :func:`check_two_way_stream`, except that the maker
|
||||
functions are expected to return objects that implement the
|
||||
:class:`~trio.abc.HalfCloseableStream` interface.
|
||||
|
||||
This function tests a *superset* of what :func:`check_two_way_stream`
|
||||
checks – if you call this, then you don't need to also call
|
||||
:func:`check_two_way_stream`.
|
||||
|
||||
"""
|
||||
await check_two_way_stream(stream_maker, clogged_stream_maker)
|
||||
|
||||
async with _ForceCloseBoth(await stream_maker()) as (s1, s2):
|
||||
assert isinstance(s1, HalfCloseableStream)
|
||||
assert isinstance(s2, HalfCloseableStream)
|
||||
|
||||
async def send_x_then_eof(s):
|
||||
await s.send_all(b"x")
|
||||
with assert_checkpoints():
|
||||
await s.send_eof()
|
||||
|
||||
async def expect_x_then_eof(r):
|
||||
await _core.wait_all_tasks_blocked()
|
||||
assert await r.receive_some(10) == b"x"
|
||||
assert await r.receive_some(10) == b""
|
||||
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(send_x_then_eof, s1)
|
||||
nursery.start_soon(expect_x_then_eof, s2)
|
||||
|
||||
# now sending is disallowed
|
||||
with _assert_raises(_core.ClosedResourceError):
|
||||
await s1.send_all(b"y")
|
||||
|
||||
# but we can do send_eof again
|
||||
with assert_checkpoints():
|
||||
await s1.send_eof()
|
||||
|
||||
# and we can still send stuff back the other way
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(send_x_then_eof, s2)
|
||||
nursery.start_soon(expect_x_then_eof, s1)
|
||||
|
||||
if clogged_stream_maker is not None:
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s1, s2):
|
||||
# send_all and send_eof simultaneously is not ok
|
||||
with _assert_raises(_core.BusyResourceError):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(s1.send_all, b"x")
|
||||
await _core.wait_all_tasks_blocked()
|
||||
nursery.start_soon(s1.send_eof)
|
||||
|
||||
async with _ForceCloseBoth(await clogged_stream_maker()) as (s1, s2):
|
||||
# wait_send_all_might_not_block and send_eof simultaneously is not
|
||||
# ok either
|
||||
with _assert_raises(_core.BusyResourceError):
|
||||
async with _core.open_nursery() as nursery:
|
||||
nursery.start_soon(s1.wait_send_all_might_not_block)
|
||||
await _core.wait_all_tasks_blocked()
|
||||
nursery.start_soon(s1.send_eof)
|
||||
@@ -0,0 +1,62 @@
|
||||
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)
|
||||
@@ -0,0 +1,591 @@
|
||||
import operator
|
||||
|
||||
from .. import _core
|
||||
from .._highlevel_generic import StapledStream
|
||||
from .. import _util
|
||||
from ..abc import SendStream, ReceiveStream
|
||||
|
||||
################################################################
|
||||
# In-memory streams - Unbounded buffer version
|
||||
################################################################
|
||||
|
||||
|
||||
class _UnboundedByteQueue:
|
||||
def __init__(self):
|
||||
self._data = bytearray()
|
||||
self._closed = False
|
||||
self._lot = _core.ParkingLot()
|
||||
self._fetch_lock = _util.ConflictDetector(
|
||||
"another task is already fetching data"
|
||||
)
|
||||
|
||||
# This object treats "close" as being like closing the send side of a
|
||||
# channel: so after close(), calling put() raises ClosedResourceError, and
|
||||
# calling the get() variants drains the buffer and then returns an empty
|
||||
# bytearray.
|
||||
def close(self):
|
||||
self._closed = True
|
||||
self._lot.unpark_all()
|
||||
|
||||
def close_and_wipe(self):
|
||||
self._data = bytearray()
|
||||
self.close()
|
||||
|
||||
def put(self, data):
|
||||
if self._closed:
|
||||
raise _core.ClosedResourceError("virtual connection closed")
|
||||
self._data += data
|
||||
self._lot.unpark_all()
|
||||
|
||||
def _check_max_bytes(self, max_bytes):
|
||||
if max_bytes is None:
|
||||
return
|
||||
max_bytes = operator.index(max_bytes)
|
||||
if max_bytes < 1:
|
||||
raise ValueError("max_bytes must be >= 1")
|
||||
|
||||
def _get_impl(self, max_bytes):
|
||||
assert self._closed or self._data
|
||||
if max_bytes is None:
|
||||
max_bytes = len(self._data)
|
||||
if self._data:
|
||||
chunk = self._data[:max_bytes]
|
||||
del self._data[:max_bytes]
|
||||
assert chunk
|
||||
return chunk
|
||||
else:
|
||||
return bytearray()
|
||||
|
||||
def get_nowait(self, max_bytes=None):
|
||||
with self._fetch_lock:
|
||||
self._check_max_bytes(max_bytes)
|
||||
if not self._closed and not self._data:
|
||||
raise _core.WouldBlock
|
||||
return self._get_impl(max_bytes)
|
||||
|
||||
async def get(self, max_bytes=None):
|
||||
with self._fetch_lock:
|
||||
self._check_max_bytes(max_bytes)
|
||||
if not self._closed and not self._data:
|
||||
await self._lot.park()
|
||||
else:
|
||||
await _core.checkpoint()
|
||||
return self._get_impl(max_bytes)
|
||||
|
||||
|
||||
class MemorySendStream(SendStream, metaclass=_util.Final):
|
||||
"""An in-memory :class:`~trio.abc.SendStream`.
|
||||
|
||||
Args:
|
||||
send_all_hook: An async function, or None. Called from
|
||||
:meth:`send_all`. Can do whatever you like.
|
||||
wait_send_all_might_not_block_hook: An async function, or None. Called
|
||||
from :meth:`wait_send_all_might_not_block`. Can do whatever you
|
||||
like.
|
||||
close_hook: A synchronous function, or None. Called from :meth:`close`
|
||||
and :meth:`aclose`. Can do whatever you like.
|
||||
|
||||
.. attribute:: send_all_hook
|
||||
wait_send_all_might_not_block_hook
|
||||
close_hook
|
||||
|
||||
All of these hooks are also exposed as attributes on the object, and
|
||||
you can change them at any time.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
send_all_hook=None,
|
||||
wait_send_all_might_not_block_hook=None,
|
||||
close_hook=None,
|
||||
):
|
||||
self._conflict_detector = _util.ConflictDetector(
|
||||
"another task is using this stream"
|
||||
)
|
||||
self._outgoing = _UnboundedByteQueue()
|
||||
self.send_all_hook = send_all_hook
|
||||
self.wait_send_all_might_not_block_hook = wait_send_all_might_not_block_hook
|
||||
self.close_hook = close_hook
|
||||
|
||||
async def send_all(self, data):
|
||||
"""Places the given data into the object's internal buffer, and then
|
||||
calls the :attr:`send_all_hook` (if any).
|
||||
|
||||
"""
|
||||
# Execute two checkpoints so we have more of a chance to detect
|
||||
# buggy user code that calls this twice at the same time.
|
||||
with self._conflict_detector:
|
||||
await _core.checkpoint()
|
||||
await _core.checkpoint()
|
||||
self._outgoing.put(data)
|
||||
if self.send_all_hook is not None:
|
||||
await self.send_all_hook()
|
||||
|
||||
async def wait_send_all_might_not_block(self):
|
||||
"""Calls the :attr:`wait_send_all_might_not_block_hook` (if any), and
|
||||
then returns immediately.
|
||||
|
||||
"""
|
||||
# Execute two checkpoints so we have more of a chance to detect
|
||||
# buggy user code that calls this twice at the same time.
|
||||
with self._conflict_detector:
|
||||
await _core.checkpoint()
|
||||
await _core.checkpoint()
|
||||
# check for being closed:
|
||||
self._outgoing.put(b"")
|
||||
if self.wait_send_all_might_not_block_hook is not None:
|
||||
await self.wait_send_all_might_not_block_hook()
|
||||
|
||||
def close(self):
|
||||
"""Marks this stream as closed, and then calls the :attr:`close_hook`
|
||||
(if any).
|
||||
|
||||
"""
|
||||
# XXX should this cancel any pending calls to the send_all_hook and
|
||||
# wait_send_all_might_not_block_hook? Those are the only places where
|
||||
# send_all and wait_send_all_might_not_block can be blocked.
|
||||
#
|
||||
# The way we set things up, send_all_hook is memory_stream_pump, and
|
||||
# wait_send_all_might_not_block_hook is unset. memory_stream_pump is
|
||||
# synchronous. So normally, send_all and wait_send_all_might_not_block
|
||||
# cannot block at all.
|
||||
self._outgoing.close()
|
||||
if self.close_hook is not None:
|
||||
self.close_hook()
|
||||
|
||||
async def aclose(self):
|
||||
"""Same as :meth:`close`, but async."""
|
||||
self.close()
|
||||
await _core.checkpoint()
|
||||
|
||||
async def get_data(self, max_bytes=None):
|
||||
"""Retrieves data from the internal buffer, blocking if necessary.
|
||||
|
||||
Args:
|
||||
max_bytes (int or None): The maximum amount of data to
|
||||
retrieve. None (the default) means to retrieve all the data
|
||||
that's present (but still blocks until at least one byte is
|
||||
available).
|
||||
|
||||
Returns:
|
||||
If this stream has been closed, an empty bytearray. Otherwise, the
|
||||
requested data.
|
||||
|
||||
"""
|
||||
return await self._outgoing.get(max_bytes)
|
||||
|
||||
def get_data_nowait(self, max_bytes=None):
|
||||
"""Retrieves data from the internal buffer, but doesn't block.
|
||||
|
||||
See :meth:`get_data` for details.
|
||||
|
||||
Raises:
|
||||
trio.WouldBlock: if no data is available to retrieve.
|
||||
|
||||
"""
|
||||
return self._outgoing.get_nowait(max_bytes)
|
||||
|
||||
|
||||
class MemoryReceiveStream(ReceiveStream, metaclass=_util.Final):
|
||||
"""An in-memory :class:`~trio.abc.ReceiveStream`.
|
||||
|
||||
Args:
|
||||
receive_some_hook: An async function, or None. Called from
|
||||
:meth:`receive_some`. Can do whatever you like.
|
||||
close_hook: A synchronous function, or None. Called from :meth:`close`
|
||||
and :meth:`aclose`. Can do whatever you like.
|
||||
|
||||
.. attribute:: receive_some_hook
|
||||
close_hook
|
||||
|
||||
Both hooks are also exposed as attributes on the object, and you can
|
||||
change them at any time.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, receive_some_hook=None, close_hook=None):
|
||||
self._conflict_detector = _util.ConflictDetector(
|
||||
"another task is using this stream"
|
||||
)
|
||||
self._incoming = _UnboundedByteQueue()
|
||||
self._closed = False
|
||||
self.receive_some_hook = receive_some_hook
|
||||
self.close_hook = close_hook
|
||||
|
||||
async def receive_some(self, max_bytes=None):
|
||||
"""Calls the :attr:`receive_some_hook` (if any), and then retrieves
|
||||
data from the internal buffer, blocking if necessary.
|
||||
|
||||
"""
|
||||
# Execute two checkpoints so we have more of a chance to detect
|
||||
# buggy user code that calls this twice at the same time.
|
||||
with self._conflict_detector:
|
||||
await _core.checkpoint()
|
||||
await _core.checkpoint()
|
||||
if self._closed:
|
||||
raise _core.ClosedResourceError
|
||||
if self.receive_some_hook is not None:
|
||||
await self.receive_some_hook()
|
||||
# self._incoming's closure state tracks whether we got an EOF.
|
||||
# self._closed tracks whether we, ourselves, are closed.
|
||||
# self.close() sends an EOF to wake us up and sets self._closed,
|
||||
# so after we wake up we have to check self._closed again.
|
||||
data = await self._incoming.get(max_bytes)
|
||||
if self._closed:
|
||||
raise _core.ClosedResourceError
|
||||
return data
|
||||
|
||||
def close(self):
|
||||
"""Discards any pending data from the internal buffer, and marks this
|
||||
stream as closed.
|
||||
|
||||
"""
|
||||
self._closed = True
|
||||
self._incoming.close_and_wipe()
|
||||
if self.close_hook is not None:
|
||||
self.close_hook()
|
||||
|
||||
async def aclose(self):
|
||||
"""Same as :meth:`close`, but async."""
|
||||
self.close()
|
||||
await _core.checkpoint()
|
||||
|
||||
def put_data(self, data):
|
||||
"""Appends the given data to the internal buffer."""
|
||||
self._incoming.put(data)
|
||||
|
||||
def put_eof(self):
|
||||
"""Adds an end-of-file marker to the internal buffer."""
|
||||
self._incoming.close()
|
||||
|
||||
|
||||
def memory_stream_pump(memory_send_stream, memory_receive_stream, *, max_bytes=None):
|
||||
"""Take data out of the given :class:`MemorySendStream`'s internal buffer,
|
||||
and put it into the given :class:`MemoryReceiveStream`'s internal buffer.
|
||||
|
||||
Args:
|
||||
memory_send_stream (MemorySendStream): The stream to get data from.
|
||||
memory_receive_stream (MemoryReceiveStream): The stream to put data into.
|
||||
max_bytes (int or None): The maximum amount of data to transfer in this
|
||||
call, or None to transfer all available data.
|
||||
|
||||
Returns:
|
||||
True if it successfully transferred some data, or False if there was no
|
||||
data to transfer.
|
||||
|
||||
This is used to implement :func:`memory_stream_one_way_pair` and
|
||||
:func:`memory_stream_pair`; see the latter's docstring for an example
|
||||
of how you might use it yourself.
|
||||
|
||||
"""
|
||||
try:
|
||||
data = memory_send_stream.get_data_nowait(max_bytes)
|
||||
except _core.WouldBlock:
|
||||
return False
|
||||
try:
|
||||
if not data:
|
||||
memory_receive_stream.put_eof()
|
||||
else:
|
||||
memory_receive_stream.put_data(data)
|
||||
except _core.ClosedResourceError:
|
||||
raise _core.BrokenResourceError("MemoryReceiveStream was closed")
|
||||
return True
|
||||
|
||||
|
||||
def memory_stream_one_way_pair():
|
||||
"""Create a connected, pure-Python, unidirectional stream with infinite
|
||||
buffering and flexible configuration options.
|
||||
|
||||
You can think of this as being a no-operating-system-involved
|
||||
Trio-streamsified version of :func:`os.pipe` (except that :func:`os.pipe`
|
||||
returns the streams in the wrong order – we follow the superior convention
|
||||
that data flows from left to right).
|
||||
|
||||
Returns:
|
||||
A tuple (:class:`MemorySendStream`, :class:`MemoryReceiveStream`), where
|
||||
the :class:`MemorySendStream` has its hooks set up so that it calls
|
||||
:func:`memory_stream_pump` from its
|
||||
:attr:`~MemorySendStream.send_all_hook` and
|
||||
:attr:`~MemorySendStream.close_hook`.
|
||||
|
||||
The end result is that data automatically flows from the
|
||||
:class:`MemorySendStream` to the :class:`MemoryReceiveStream`. But you're
|
||||
also free to rearrange things however you like. For example, you can
|
||||
temporarily set the :attr:`~MemorySendStream.send_all_hook` to None if you
|
||||
want to simulate a stall in data transmission. Or see
|
||||
:func:`memory_stream_pair` for a more elaborate example.
|
||||
|
||||
"""
|
||||
send_stream = MemorySendStream()
|
||||
recv_stream = MemoryReceiveStream()
|
||||
|
||||
def pump_from_send_stream_to_recv_stream():
|
||||
memory_stream_pump(send_stream, recv_stream)
|
||||
|
||||
async def async_pump_from_send_stream_to_recv_stream():
|
||||
pump_from_send_stream_to_recv_stream()
|
||||
|
||||
send_stream.send_all_hook = async_pump_from_send_stream_to_recv_stream
|
||||
send_stream.close_hook = pump_from_send_stream_to_recv_stream
|
||||
return send_stream, recv_stream
|
||||
|
||||
|
||||
def _make_stapled_pair(one_way_pair):
|
||||
pipe1_send, pipe1_recv = one_way_pair()
|
||||
pipe2_send, pipe2_recv = one_way_pair()
|
||||
stream1 = StapledStream(pipe1_send, pipe2_recv)
|
||||
stream2 = StapledStream(pipe2_send, pipe1_recv)
|
||||
return stream1, stream2
|
||||
|
||||
|
||||
def memory_stream_pair():
|
||||
"""Create a connected, pure-Python, bidirectional stream with infinite
|
||||
buffering and flexible configuration options.
|
||||
|
||||
This is a convenience function that creates two one-way streams using
|
||||
:func:`memory_stream_one_way_pair`, and then uses
|
||||
:class:`~trio.StapledStream` to combine them into a single bidirectional
|
||||
stream.
|
||||
|
||||
This is like a no-operating-system-involved, Trio-streamsified version of
|
||||
:func:`socket.socketpair`.
|
||||
|
||||
Returns:
|
||||
A pair of :class:`~trio.StapledStream` objects that are connected so
|
||||
that data automatically flows from one to the other in both directions.
|
||||
|
||||
After creating a stream pair, you can send data back and forth, which is
|
||||
enough for simple tests::
|
||||
|
||||
left, right = memory_stream_pair()
|
||||
await left.send_all(b"123")
|
||||
assert await right.receive_some() == b"123"
|
||||
await right.send_all(b"456")
|
||||
assert await left.receive_some() == b"456"
|
||||
|
||||
But if you read the docs for :class:`~trio.StapledStream` and
|
||||
:func:`memory_stream_one_way_pair`, you'll see that all the pieces
|
||||
involved in wiring this up are public APIs, so you can adjust to suit the
|
||||
requirements of your tests. For example, here's how to tweak a stream so
|
||||
that data flowing from left to right trickles in one byte at a time (but
|
||||
data flowing from right to left proceeds at full speed)::
|
||||
|
||||
left, right = memory_stream_pair()
|
||||
async def trickle():
|
||||
# left is a StapledStream, and left.send_stream is a MemorySendStream
|
||||
# right is a StapledStream, and right.recv_stream is a MemoryReceiveStream
|
||||
while memory_stream_pump(left.send_stream, right.recv_stream, max_bytes=1):
|
||||
# Pause between each byte
|
||||
await trio.sleep(1)
|
||||
# Normally this send_all_hook calls memory_stream_pump directly without
|
||||
# passing in a max_bytes. We replace it with our custom version:
|
||||
left.send_stream.send_all_hook = trickle
|
||||
|
||||
And here's a simple test using our modified stream objects::
|
||||
|
||||
async def sender():
|
||||
await left.send_all(b"12345")
|
||||
await left.send_eof()
|
||||
|
||||
async def receiver():
|
||||
async for data in right:
|
||||
print(data)
|
||||
|
||||
async with trio.open_nursery() as nursery:
|
||||
nursery.start_soon(sender)
|
||||
nursery.start_soon(receiver)
|
||||
|
||||
By default, this will print ``b"12345"`` and then immediately exit; with
|
||||
our trickle stream it instead sleeps 1 second, then prints ``b"1"``, then
|
||||
sleeps 1 second, then prints ``b"2"``, etc.
|
||||
|
||||
Pro-tip: you can insert sleep calls (like in our example above) to
|
||||
manipulate the flow of data across tasks... and then use
|
||||
:class:`MockClock` and its :attr:`~MockClock.autojump_threshold`
|
||||
functionality to keep your test suite running quickly.
|
||||
|
||||
If you want to stress test a protocol implementation, one nice trick is to
|
||||
use the :mod:`random` module (preferably with a fixed seed) to move random
|
||||
numbers of bytes at a time, and insert random sleeps in between them. You
|
||||
can also set up a custom :attr:`~MemoryReceiveStream.receive_some_hook` if
|
||||
you want to manipulate things on the receiving side, and not just the
|
||||
sending side.
|
||||
|
||||
"""
|
||||
return _make_stapled_pair(memory_stream_one_way_pair)
|
||||
|
||||
|
||||
################################################################
|
||||
# In-memory streams - Lockstep version
|
||||
################################################################
|
||||
|
||||
|
||||
class _LockstepByteQueue:
|
||||
def __init__(self):
|
||||
self._data = bytearray()
|
||||
self._sender_closed = False
|
||||
self._receiver_closed = False
|
||||
self._receiver_waiting = False
|
||||
self._waiters = _core.ParkingLot()
|
||||
self._send_conflict_detector = _util.ConflictDetector(
|
||||
"another task is already sending"
|
||||
)
|
||||
self._receive_conflict_detector = _util.ConflictDetector(
|
||||
"another task is already receiving"
|
||||
)
|
||||
|
||||
def _something_happened(self):
|
||||
self._waiters.unpark_all()
|
||||
|
||||
# Always wakes up when one side is closed, because everyone always reacts
|
||||
# to that.
|
||||
async def _wait_for(self, fn):
|
||||
while True:
|
||||
if fn():
|
||||
break
|
||||
if self._sender_closed or self._receiver_closed:
|
||||
break
|
||||
await self._waiters.park()
|
||||
await _core.checkpoint()
|
||||
|
||||
def close_sender(self):
|
||||
self._sender_closed = True
|
||||
self._something_happened()
|
||||
|
||||
def close_receiver(self):
|
||||
self._receiver_closed = True
|
||||
self._something_happened()
|
||||
|
||||
async def send_all(self, data):
|
||||
with self._send_conflict_detector:
|
||||
if self._sender_closed:
|
||||
raise _core.ClosedResourceError
|
||||
if self._receiver_closed:
|
||||
raise _core.BrokenResourceError
|
||||
assert not self._data
|
||||
self._data += data
|
||||
self._something_happened()
|
||||
await self._wait_for(lambda: not self._data)
|
||||
if self._sender_closed:
|
||||
raise _core.ClosedResourceError
|
||||
if self._data and self._receiver_closed:
|
||||
raise _core.BrokenResourceError
|
||||
|
||||
async def wait_send_all_might_not_block(self):
|
||||
with self._send_conflict_detector:
|
||||
if self._sender_closed:
|
||||
raise _core.ClosedResourceError
|
||||
if self._receiver_closed:
|
||||
await _core.checkpoint()
|
||||
return
|
||||
await self._wait_for(lambda: self._receiver_waiting)
|
||||
if self._sender_closed:
|
||||
raise _core.ClosedResourceError
|
||||
|
||||
async def receive_some(self, max_bytes=None):
|
||||
with self._receive_conflict_detector:
|
||||
# Argument validation
|
||||
if max_bytes is not None:
|
||||
max_bytes = operator.index(max_bytes)
|
||||
if max_bytes < 1:
|
||||
raise ValueError("max_bytes must be >= 1")
|
||||
# State validation
|
||||
if self._receiver_closed:
|
||||
raise _core.ClosedResourceError
|
||||
# Wake wait_send_all_might_not_block and wait for data
|
||||
self._receiver_waiting = True
|
||||
self._something_happened()
|
||||
try:
|
||||
await self._wait_for(lambda: self._data)
|
||||
finally:
|
||||
self._receiver_waiting = False
|
||||
if self._receiver_closed:
|
||||
raise _core.ClosedResourceError
|
||||
# Get data, possibly waking send_all
|
||||
if self._data:
|
||||
# Neat trick: if max_bytes is None, then obj[:max_bytes] is
|
||||
# the same as obj[:].
|
||||
got = self._data[:max_bytes]
|
||||
del self._data[:max_bytes]
|
||||
self._something_happened()
|
||||
return got
|
||||
else:
|
||||
assert self._sender_closed
|
||||
return b""
|
||||
|
||||
|
||||
class _LockstepSendStream(SendStream):
|
||||
def __init__(self, lbq):
|
||||
self._lbq = lbq
|
||||
|
||||
def close(self):
|
||||
self._lbq.close_sender()
|
||||
|
||||
async def aclose(self):
|
||||
self.close()
|
||||
await _core.checkpoint()
|
||||
|
||||
async def send_all(self, data):
|
||||
await self._lbq.send_all(data)
|
||||
|
||||
async def wait_send_all_might_not_block(self):
|
||||
await self._lbq.wait_send_all_might_not_block()
|
||||
|
||||
|
||||
class _LockstepReceiveStream(ReceiveStream):
|
||||
def __init__(self, lbq):
|
||||
self._lbq = lbq
|
||||
|
||||
def close(self):
|
||||
self._lbq.close_receiver()
|
||||
|
||||
async def aclose(self):
|
||||
self.close()
|
||||
await _core.checkpoint()
|
||||
|
||||
async def receive_some(self, max_bytes=None):
|
||||
return await self._lbq.receive_some(max_bytes)
|
||||
|
||||
|
||||
def lockstep_stream_one_way_pair():
|
||||
"""Create a connected, pure Python, unidirectional stream where data flows
|
||||
in lockstep.
|
||||
|
||||
Returns:
|
||||
A tuple
|
||||
(:class:`~trio.abc.SendStream`, :class:`~trio.abc.ReceiveStream`).
|
||||
|
||||
This stream has *absolutely no* buffering. Each call to
|
||||
:meth:`~trio.abc.SendStream.send_all` will block until all the given data
|
||||
has been returned by a call to
|
||||
:meth:`~trio.abc.ReceiveStream.receive_some`.
|
||||
|
||||
This can be useful for testing flow control mechanisms in an extreme case,
|
||||
or for setting up "clogged" streams to use with
|
||||
:func:`check_one_way_stream` and friends.
|
||||
|
||||
In addition to fulfilling the :class:`~trio.abc.SendStream` and
|
||||
:class:`~trio.abc.ReceiveStream` interfaces, the return objects
|
||||
also have a synchronous ``close`` method.
|
||||
|
||||
"""
|
||||
|
||||
lbq = _LockstepByteQueue()
|
||||
return _LockstepSendStream(lbq), _LockstepReceiveStream(lbq)
|
||||
|
||||
|
||||
def lockstep_stream_pair():
|
||||
"""Create a connected, pure-Python, bidirectional stream where data flows
|
||||
in lockstep.
|
||||
|
||||
Returns:
|
||||
A tuple (:class:`~trio.StapledStream`, :class:`~trio.StapledStream`).
|
||||
|
||||
This is a convenience function that creates two one-way streams using
|
||||
:func:`lockstep_stream_one_way_pair`, and then uses
|
||||
:class:`~trio.StapledStream` to combine them into a single bidirectional
|
||||
stream.
|
||||
|
||||
"""
|
||||
return _make_stapled_pair(lockstep_stream_one_way_pair)
|
||||
34
asq-env/lib/python3.9/site-packages/trio/testing/_network.py
Normal file
34
asq-env/lib/python3.9/site-packages/trio/testing/_network.py
Normal file
@@ -0,0 +1,34 @@
|
||||
from .. import socket as tsocket
|
||||
from .._highlevel_socket import SocketStream
|
||||
|
||||
|
||||
async def open_stream_to_socket_listener(socket_listener):
|
||||
"""Connect to the given :class:`~trio.SocketListener`.
|
||||
|
||||
This is particularly useful in tests when you want to let a server pick
|
||||
its own port, and then connect to it::
|
||||
|
||||
listeners = await trio.open_tcp_listeners(0)
|
||||
client = await trio.testing.open_stream_to_socket_listener(listeners[0])
|
||||
|
||||
Args:
|
||||
socket_listener (~trio.SocketListener): The
|
||||
:class:`~trio.SocketListener` to connect to.
|
||||
|
||||
Returns:
|
||||
SocketStream: a stream connected to the given listener.
|
||||
|
||||
"""
|
||||
family = socket_listener.socket.family
|
||||
sockaddr = socket_listener.socket.getsockname()
|
||||
if family in (tsocket.AF_INET, tsocket.AF_INET6):
|
||||
sockaddr = list(sockaddr)
|
||||
if sockaddr[0] == "0.0.0.0":
|
||||
sockaddr[0] = "127.0.0.1"
|
||||
if sockaddr[0] == "::":
|
||||
sockaddr[0] = "::1"
|
||||
sockaddr = tuple(sockaddr)
|
||||
|
||||
sock = tsocket.socket(family=family)
|
||||
await sock.connect(sockaddr)
|
||||
return SocketStream(sock)
|
||||
@@ -0,0 +1,82 @@
|
||||
from collections import defaultdict
|
||||
|
||||
import attr
|
||||
from async_generator import asynccontextmanager
|
||||
|
||||
from .. import _core
|
||||
from .. import _util
|
||||
from .. import Event
|
||||
|
||||
if False:
|
||||
from typing import DefaultDict, Set
|
||||
|
||||
|
||||
@attr.s(eq=False, hash=False)
|
||||
class Sequencer(metaclass=_util.Final):
|
||||
"""A convenience class for forcing code in different tasks to run in an
|
||||
explicit linear order.
|
||||
|
||||
Instances of this class implement a ``__call__`` method which returns an
|
||||
async context manager. The idea is that you pass a sequence number to
|
||||
``__call__`` to say where this block of code should go in the linear
|
||||
sequence. Block 0 starts immediately, and then block N doesn't start until
|
||||
block N-1 has finished.
|
||||
|
||||
Example:
|
||||
An extremely elaborate way to print the numbers 0-5, in order::
|
||||
|
||||
async def worker1(seq):
|
||||
async with seq(0):
|
||||
print(0)
|
||||
async with seq(4):
|
||||
print(4)
|
||||
|
||||
async def worker2(seq):
|
||||
async with seq(2):
|
||||
print(2)
|
||||
async with seq(5):
|
||||
print(5)
|
||||
|
||||
async def worker3(seq):
|
||||
async with seq(1):
|
||||
print(1)
|
||||
async with seq(3):
|
||||
print(3)
|
||||
|
||||
async def main():
|
||||
seq = trio.testing.Sequencer()
|
||||
async with trio.open_nursery() as nursery:
|
||||
nursery.start_soon(worker1, seq)
|
||||
nursery.start_soon(worker2, seq)
|
||||
nursery.start_soon(worker3, seq)
|
||||
|
||||
"""
|
||||
|
||||
_sequence_points = attr.ib(
|
||||
factory=lambda: defaultdict(Event), init=False
|
||||
) # type: DefaultDict[int, Event]
|
||||
_claimed = attr.ib(factory=set, init=False) # type: Set[int]
|
||||
_broken = attr.ib(default=False, init=False)
|
||||
|
||||
@asynccontextmanager
|
||||
async def __call__(self, position: int):
|
||||
if position in self._claimed:
|
||||
raise RuntimeError("Attempted to re-use sequence point {}".format(position))
|
||||
if self._broken:
|
||||
raise RuntimeError("sequence broken!")
|
||||
self._claimed.add(position)
|
||||
if position != 0:
|
||||
try:
|
||||
await self._sequence_points[position].wait()
|
||||
except _core.Cancelled:
|
||||
self._broken = True
|
||||
for event in self._sequence_points.values():
|
||||
event.set()
|
||||
raise RuntimeError("Sequencer wait cancelled -- sequence broken")
|
||||
else:
|
||||
if self._broken:
|
||||
raise RuntimeError("sequence broken!")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._sequence_points[position + 1].set()
|
||||
@@ -0,0 +1,29 @@
|
||||
from functools import wraps, partial
|
||||
|
||||
from .. import _core
|
||||
from ..abc import Clock, Instrument
|
||||
|
||||
|
||||
# Use:
|
||||
#
|
||||
# @trio_test
|
||||
# async def test_whatever():
|
||||
# await ...
|
||||
#
|
||||
# Also: if a pytest fixture is passed in that subclasses the Clock abc, then
|
||||
# that clock is passed to trio.run().
|
||||
def trio_test(fn):
|
||||
@wraps(fn)
|
||||
def wrapper(**kwargs):
|
||||
__tracebackhide__ = True
|
||||
clocks = [c for c in kwargs.values() if isinstance(c, Clock)]
|
||||
if not clocks:
|
||||
clock = None
|
||||
elif len(clocks) == 1:
|
||||
clock = clocks[0]
|
||||
else:
|
||||
raise ValueError("too many clocks spoil the broth!")
|
||||
instruments = [i for i in kwargs.values() if isinstance(i, Instrument)]
|
||||
return _core.run(partial(fn, **kwargs), clock=clock, instruments=instruments)
|
||||
|
||||
return wrapper
|
||||
Reference in New Issue
Block a user