async def test_replay_ending_cancels_timeouts(event_loop, replay_deps):
mock_merger, mock_sender, _ = replay_deps
conf = MockReplayConfig(15, 100)
replay = Replay(*replay_deps, conf, 1)
exhaust_callbacks(event_loop)
replay.close()
mock_merger.stop_accepting_connections.assert_called()
mock_sender.stop_accepting_connections.assert_called()
mock_merger.stop_accepting_connections.reset_mock()
mock_sender.stop_accepting_connections.reset_mock()
# Replay expects these to end after calling close
mock_merger.wait_for_ended._lock.set()
mock_sender.wait_for_ended._lock.set()
await replay.wait_for_ended()
# Did replay stop waiting until merger says there aren't more connections?
await asyncio.sleep(1)
mock_merger.no_connections_for._lock.set()
# Did replay stop waiting for its own timeout?
await asyncio.sleep(19)
exhaust_callbacks(event_loop)
mock_merger.stop_accepting_connections.assert_not_called()
mock_sender.stop_accepting_connections.assert_not_called()
async def test_merger_active_connection_prevents_ending(
outside_source_stream, mock_merge_strategy, mock_stream_builder,
mock_connection_streams, mock_connections, event_loop):
connection = mock_connections()
connection_stream = mock_connection_streams()
mock_stream_builder.side_effect = [connection_stream]
canonical_stream = outside_source_stream
async def long_read():
await asyncio.sleep(0.05)
return b""
connection_stream.read.side_effect = long_read
connection_stream.ended.side_effect = lambda: connection_stream.read.called
merger = Merger(mock_stream_builder, 0.01, mock_merge_strategy,
canonical_stream)
f = asyncio.ensure_future(merger.wait_for_ended())
h = asyncio.ensure_future(merger.handle_connection(connection))
await asyncio.sleep(0.02)
exhaust_callbacks(event_loop)
assert not f.done()
await h
await asyncio.sleep(0.02)
exhaust_callbacks(event_loop)
assert f.done()
async def test_header_mixin_waits_on_ended(event_loop):
s = HeaderMixinStream()
f = asyncio.ensure_future(s.wait_for_header())
exhaust_callbacks(event_loop)
assert not f.done()
s._ended = True
s._signal_header_read_or_ended()
head = await f
assert head is None
async def test_data_mixin_waits_on_ended(event_loop):
s = DataMixinStream()
f = asyncio.ensure_future(s.wait_for_data())
exhaust_callbacks(event_loop)
assert not f.done()
s._ended = True
s._signal_new_data_or_ended()
data = await f
assert data == b""
async def test_data_mixin_wait_until_position(event_loop):
s = DataMixinStream()
f = asyncio.ensure_future(s.wait_for_data(3))
s._data += b"a"
s._signal_new_data_or_ended()
exhaust_callbacks(event_loop)
assert not f.done()
s._data += b"aab"
s._signal_new_data_or_ended()
data = await f
assert data == b"b"
async def test_outside_source_stream_read_header(event_loop):
stream = OutsideSourceReplayStream()
f = asyncio.ensure_future(stream.wait_for_header())
exhaust_callbacks(event_loop)
assert not f.done()
stream.set_header("header")
assert await f == "header"
assert stream.header == "header"
stream.finish()
await stream.wait_for_ended()
assert stream.data.bytes() == b""
async def test_outside_source_stream_wait_until_position(event_loop):
stream = OutsideSourceReplayStream()
f = asyncio.ensure_future(stream.wait_for_data(3))
stream.set_header("header")
stream.feed_data(b"a")
exhaust_callbacks(event_loop)
assert not f.done()
stream.feed_data(b"aa")
exhaust_callbacks(event_loop)
assert not f.done()
stream.feed_data(b"ccc")
assert await f == 3
assert stream.data[3:6] == b"ccc"
async def test_replay_closes_after_timeout(
event_loop, mock_merger, mock_sender, mock_bookkeeper):
timeout = 15
replay = Replay(mock_merger, mock_sender, mock_bookkeeper, timeout, 1)
mock_merger.close.assert_not_called()
mock_sender.close.assert_not_called()
await asyncio.sleep(20)
exhaust_callbacks(event_loop)
mock_merger.close.assert_called()
mock_sender.close.assert_called()
mock_merger._manual_end.set()
mock_sender._manual_end.set()
await replay.wait_for_ended()
async def test_outside_source_stream_immediate_end(event_loop):
stream = OutsideSourceReplayStream()
f1 = asyncio.ensure_future(stream.wait_for_header())
f2 = asyncio.ensure_future(stream.wait_for_data(0))
f3 = asyncio.ensure_future(stream.wait_for_ended())
exhaust_callbacks(event_loop)
assert not any(x.done() for x in [f1, f2, f3])
stream.finish()
assert await f1 is None
assert await f2 == 0
assert stream.header is None
assert stream.data.bytes() == b""
await f3
async def test_merger_no_connections_wait_empty(outside_source_stream,
mock_merge_strategy,
mock_reader_builder,
mock_delayed_stream_builder,
event_loop):
canonical_stream = outside_source_stream
merger = Merger(mock_reader_builder, mock_delayed_stream_builder,
mock_merge_strategy, canonical_stream)
f = asyncio.ensure_future(merger.no_connections_for(0.01))
exhaust_callbacks(event_loop)
assert not f.done()
await asyncio.sleep(0.02)
exhaust_callbacks(event_loop)
assert f.done()
await f
merger.stop_accepting_connections()
await merger.wait_for_ended()
async def test_merger_connection_extends_grace_period(
outside_source_stream, mock_merge_strategy, mock_stream_builder,
mock_connection_streams, mock_connections, event_loop):
connection = mock_connections()
connection_stream = mock_connection_streams()
mock_stream_builder.side_effect = [connection_stream]
canonical_stream = outside_source_stream
connection_stream.ended.return_value = True
merger = Merger(mock_stream_builder, 0.03, mock_merge_strategy,
canonical_stream)
f = asyncio.ensure_future(merger.wait_for_ended())
await asyncio.sleep(0.02)
await merger.handle_connection(connection)
await asyncio.sleep(0.02)
exhaust_callbacks(event_loop)
assert not f.done()
await f
async def test_replay_timeouts_while_ending_dont_explode(
event_loop, replay_deps):
mock_merger, mock_sender, _ = replay_deps
conf = MockReplayConfig(15, 100)
replay = Replay(*replay_deps, conf, 1)
exhaust_callbacks(event_loop)
# Won't finish until we tell mock merger & sender to end
f = asyncio.ensure_future(replay.wait_for_ended())
await asyncio.sleep(1)
mock_merger.no_connections_for._lock.set()
await asyncio.sleep(19)
exhaust_callbacks(event_loop)
mock_merger.wait_for_ended._lock.set()
mock_sender.wait_for_ended._lock.set()
await f
async def test_merger_no_connection_wait_active_connection(
outside_source_stream, mock_merge_strategy, mock_reader_builder,
mock_delayed_stream_builder, mock_readers, mock_connections,
event_loop):
connection = mock_connections()
reader = mock_readers()
mock_reader_builder.side_effect = [reader]
canonical_stream = outside_source_stream
async def long_read():
await asyncio.sleep(0.05)
return
reader.read.side_effect = long_read
merger = Merger(mock_reader_builder, mock_delayed_stream_builder,
mock_merge_strategy, canonical_stream)
f = asyncio.ensure_future(merger.no_connections_for(0.01))
h = asyncio.ensure_future(merger.handle_connection(connection))
await asyncio.sleep(0.02)
exhaust_callbacks(event_loop)
assert not f.done()
await h
exhaust_callbacks(event_loop)
assert not f.done()
await asyncio.sleep(0.02)
exhaust_callbacks(event_loop)
assert f.done()
await f
merger.stop_accepting_connections()
await merger.wait_for_ended()
async def test_replay_closes_after_timeout(event_loop, replay_deps,
mock_conn_plus_head):
mock_merger, mock_sender, _ = replay_deps
conf = MockReplayConfig(15, 100)
replay = Replay(*replay_deps, conf, 1)
mock_merger.stop_accepting_connections.assert_not_called()
mock_sender.stop_accepting_connections.assert_not_called()
writer = mock_conn_plus_head(ConnectionHeader.Type.WRITER, 1)
wf = asyncio.ensure_future(replay.handle_connection(*writer))
await asyncio.sleep(20)
exhaust_callbacks(event_loop)
mock_merger.stop_accepting_connections.assert_called()
mock_sender.stop_accepting_connections.assert_called()
writer[1].close.assert_called()
mock_merger.handle_connection._lock.set()
mock_merger.wait_for_ended._lock.set()
mock_sender.wait_for_ended._lock.set()
await wf
await replay.wait_for_ended()
async def test_replay_close_cancels_timeout(
event_loop, mock_merger, mock_sender, mock_bookkeeper):
timeout = 15
replay = Replay(mock_merger, mock_sender, mock_bookkeeper, timeout, 1)
exhaust_callbacks(event_loop)
replay.close()
mock_merger.close.assert_called()
mock_sender.close.assert_called()
mock_merger.close.reset_mock()
mock_sender.close.reset_mock()
# Replay expects these to end after calling close
mock_merger._manual_end.set()
mock_sender._manual_end.set()
await asyncio.sleep(20)
exhaust_callbacks(event_loop)
mock_merger.close.assert_not_called()
mock_sender.close.assert_not_called()
mock_merger._manual_end.set()
mock_sender._manual_end.set()
await replay.wait_for_ended()