513 lines
20 KiB
Python
513 lines
20 KiB
Python
# 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)
|