mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
67ab6e27d8
* feat: implement AsyncContextManager for IMuxedStream to support async with * doc: add newsfragment
128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
import pytest
|
|
import trio
|
|
|
|
from libp2p.stream_muxer.exceptions import (
|
|
MuxedStreamClosed,
|
|
MuxedStreamError,
|
|
)
|
|
from libp2p.stream_muxer.mplex.datastructures import (
|
|
StreamID,
|
|
)
|
|
from libp2p.stream_muxer.mplex.mplex_stream import (
|
|
MplexStream,
|
|
)
|
|
from libp2p.stream_muxer.yamux.yamux import (
|
|
YamuxStream,
|
|
)
|
|
|
|
|
|
class DummySecuredConn:
|
|
async def write(self, data):
|
|
pass
|
|
|
|
|
|
class MockMuxedConn:
|
|
def __init__(self):
|
|
self.streams = {}
|
|
self.streams_lock = trio.Lock()
|
|
self.event_shutting_down = trio.Event()
|
|
self.event_closed = trio.Event()
|
|
self.event_started = trio.Event()
|
|
self.secured_conn = DummySecuredConn() # For YamuxStream
|
|
|
|
async def send_message(self, flag, data, stream_id):
|
|
pass
|
|
|
|
def get_remote_address(self):
|
|
return None
|
|
|
|
|
|
@pytest.mark.trio
|
|
async def test_mplex_stream_async_context_manager():
|
|
muxed_conn = MockMuxedConn()
|
|
stream_id = StreamID(1, True) # Use real StreamID
|
|
stream = MplexStream(
|
|
name="test_stream",
|
|
stream_id=stream_id,
|
|
muxed_conn=muxed_conn,
|
|
incoming_data_channel=trio.open_memory_channel(8)[1],
|
|
)
|
|
async with stream as s:
|
|
assert s is stream
|
|
assert not stream.event_local_closed.is_set()
|
|
assert not stream.event_remote_closed.is_set()
|
|
assert not stream.event_reset.is_set()
|
|
assert stream.event_local_closed.is_set()
|
|
|
|
|
|
@pytest.mark.trio
|
|
async def test_yamux_stream_async_context_manager():
|
|
muxed_conn = MockMuxedConn()
|
|
stream = YamuxStream(stream_id=1, conn=muxed_conn, is_initiator=True)
|
|
async with stream as s:
|
|
assert s is stream
|
|
assert not stream.closed
|
|
assert not stream.send_closed
|
|
assert not stream.recv_closed
|
|
assert stream.send_closed
|
|
|
|
|
|
@pytest.mark.trio
|
|
async def test_mplex_stream_async_context_manager_with_error():
|
|
muxed_conn = MockMuxedConn()
|
|
stream_id = StreamID(1, True)
|
|
stream = MplexStream(
|
|
name="test_stream",
|
|
stream_id=stream_id,
|
|
muxed_conn=muxed_conn,
|
|
incoming_data_channel=trio.open_memory_channel(8)[1],
|
|
)
|
|
with pytest.raises(ValueError):
|
|
async with stream as s:
|
|
assert s is stream
|
|
assert not stream.event_local_closed.is_set()
|
|
assert not stream.event_remote_closed.is_set()
|
|
assert not stream.event_reset.is_set()
|
|
raise ValueError("Test error")
|
|
assert stream.event_local_closed.is_set()
|
|
|
|
|
|
@pytest.mark.trio
|
|
async def test_yamux_stream_async_context_manager_with_error():
|
|
muxed_conn = MockMuxedConn()
|
|
stream = YamuxStream(stream_id=1, conn=muxed_conn, is_initiator=True)
|
|
with pytest.raises(ValueError):
|
|
async with stream as s:
|
|
assert s is stream
|
|
assert not stream.closed
|
|
assert not stream.send_closed
|
|
assert not stream.recv_closed
|
|
raise ValueError("Test error")
|
|
assert stream.send_closed
|
|
|
|
|
|
@pytest.mark.trio
|
|
async def test_mplex_stream_async_context_manager_write_after_close():
|
|
muxed_conn = MockMuxedConn()
|
|
stream_id = StreamID(1, True)
|
|
stream = MplexStream(
|
|
name="test_stream",
|
|
stream_id=stream_id,
|
|
muxed_conn=muxed_conn,
|
|
incoming_data_channel=trio.open_memory_channel(8)[1],
|
|
)
|
|
async with stream as s:
|
|
assert s is stream
|
|
with pytest.raises(MuxedStreamClosed):
|
|
await stream.write(b"test data")
|
|
|
|
|
|
@pytest.mark.trio
|
|
async def test_yamux_stream_async_context_manager_write_after_close():
|
|
muxed_conn = MockMuxedConn()
|
|
stream = YamuxStream(stream_id=1, conn=muxed_conn, is_initiator=True)
|
|
async with stream as s:
|
|
assert s is stream
|
|
with pytest.raises(MuxedStreamError):
|
|
await stream.write(b"test data")
|