StarMirror/py-libp2p
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix MplexStream.read

Browse Source
This commit is contained in:
mhchia 2019-09-06 17:26:40 +08:00
parent 95926b7376
commit 649a230776
No known key found for this signature in database
GPG Key ID: 389EFBEA1362589A
8 changed files with 116 additions and 143 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
examples/chat/chat.py
View File

@ -11,11 +11,12 @@ from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.typing import TProtocol from libp2p.typing import TProtocol
PROTOCOL_ID = TProtocol("/chat/1.0.0") PROTOCOL_ID = TProtocol("/chat/1.0.0")
MAX_READ_LEN = 2 ** 32 - 1
async def read_data(stream: INetStream) -> None: async def read_data(stream: INetStream) -> None:
while True: while True:
read_bytes = await stream.read() read_bytes = await stream.read(MAX_READ_LEN)
if read_bytes is not None: if read_bytes is not None:
read_string = read_bytes.decode() read_string = read_bytes.decode()
if read_string != "\n": if read_string != "\n":
@ -24,7 +25,6 @@ async def read_data(stream: INetStream) -> None:
print("\x1b[32m %s\x1b[0m " % read_string, end="") print("\x1b[32m %s\x1b[0m " % read_string, end="")
# FIXME(mhchia): Reconsider whether we should use a thread pool here.
async def write_data(stream: INetStream) -> None: async def write_data(stream: INetStream) -> None:
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
while True: while True:

5
examples/echo/echo.py
View File

@ -14,6 +14,7 @@ PROTOCOL_ID = TProtocol("/echo/1.0.0")
async def _echo_stream_handler(stream: INetStream) -> None: async def _echo_stream_handler(stream: INetStream) -> None:
# Wait until EOF
msg = await stream.read() msg = await stream.read()
await stream.write(msg) await stream.write(msg)
await stream.close() await stream.close()
@ -72,13 +73,13 @@ async def run(port: int, destination: str, localhost: bool, seed: int = None) ->
msg = b"hi, there!\n" msg = b"hi, there!\n"
await stream.write(msg) await stream.write(msg)
# Notify the other side about EOF
await stream.close()
response = await stream.read() response = await stream.read()
print(f"Sent: {msg}") print(f"Sent: {msg}")
print(f"Got: {response}") print(f"Got: {response}")
await stream.close()
def main() -> None: def main() -> None:
description = """ description = """

40
libp2p/stream_muxer/mplex/mplex_stream.py
View File

@ -55,7 +55,6 @@ class MplexStream(IMuxedStream):
return self.stream_id.is_initiator return self.stream_id.is_initiator
async def _wait_for_data(self) -> None: async def _wait_for_data(self) -> None:
print("!@# _wait_for_data: 0")
done, pending = await asyncio.wait( done, pending = await asyncio.wait(
[ [
self.event_reset.wait(), self.event_reset.wait(),
@ -64,16 +63,25 @@ class MplexStream(IMuxedStream):
], ],
return_when=asyncio.FIRST_COMPLETED, return_when=asyncio.FIRST_COMPLETED,
) )
print("!@# _wait_for_data: 1")
if self.event_reset.is_set(): if self.event_reset.is_set():
raise MplexStreamReset raise MplexStreamReset
if self.event_remote_closed.is_set(): if self.event_remote_closed.is_set():
while not self.incoming_data.empty():
self._buf.extend(await self.incoming_data.get())
raise MplexStreamEOF raise MplexStreamEOF
# TODO: Handle timeout when deadline is used.
data = tuple(done)[0].result() data = tuple(done)[0].result()
self._buf.extend(data) self._buf.extend(data)
async def _read_until_eof(self) -> bytes:
while True:
try:
await self._wait_for_data()
except MplexStreamEOF:
break
payload = self._buf
self._buf = self._buf[len(payload) :]
return bytes(payload)
async def read(self, n: int = -1) -> bytes: async def read(self, n: int = -1) -> bytes:
""" """
Read up to n bytes. Read possibly returns fewer than `n` bytes, Read up to n bytes. Read possibly returns fewer than `n` bytes,
@ -87,21 +95,17 @@ class MplexStream(IMuxedStream):
raise ValueError( raise ValueError(
f"the number of bytes to read `n` must be positive or -1 to indicate read until EOF" f"the number of bytes to read `n` must be positive or -1 to indicate read until EOF"
) )
if self.event_reset.is_set():
# FIXME: If `n == -1`, we should blocking read until EOF, instead of returning when raise MplexStreamReset
# no message is available.
# If `n >= 0`, read up to `n` bytes.
# Else, read until no message is available.
while len(self._buf) < n or n == -1:
# new_bytes = await self.incoming_data.get()
try:
await self._wait_for_data()
except MplexStreamEOF:
break
payload: bytearray
if n == -1: if n == -1:
payload = self._buf return await self._read_until_eof()
else: if len(self._buf) == 0:
await self._wait_for_data()
# Read up to `n` bytes.
while len(self._buf) < n:
if self.incoming_data.empty() or self.event_remote_closed.is_set():
break
self._buf.extend(await self.incoming_data.get())
payload = self._buf[:n] payload = self._buf[:n]
self._buf = self._buf[len(payload) :] self._buf = self._buf[len(payload) :]
return bytes(payload) return bytes(payload)

28
tests/examples/test_chat.py
View File

@ -10,10 +10,13 @@ PROTOCOL_ID = "/chat/1.0.0"
async def hello_world(host_a, host_b): async def hello_world(host_a, host_b):
hello_world_from_host_a = b"hello world from host a"
hello_world_from_host_b = b"hello world from host b"
async def stream_handler(stream): async def stream_handler(stream):
read = await stream.read() read = await stream.read(len(hello_world_from_host_b))
assert read == b"hello world from host b" assert read == hello_world_from_host_b
await stream.write(b"hello world from host a") await stream.write(hello_world_from_host_a)
await stream.close() await stream.close()
host_a.set_stream_handler(PROTOCOL_ID, stream_handler) host_a.set_stream_handler(PROTOCOL_ID, stream_handler)
@ -21,9 +24,9 @@ async def hello_world(host_a, host_b):
# Start a stream with the destination. # Start a stream with the destination.
# Multiaddress of the destination peer is fetched from the peerstore using 'peerId'. # Multiaddress of the destination peer is fetched from the peerstore using 'peerId'.
stream = await host_b.new_stream(host_a.get_id(), [PROTOCOL_ID]) stream = await host_b.new_stream(host_a.get_id(), [PROTOCOL_ID])
await stream.write(b"hello world from host b") await stream.write(hello_world_from_host_b)
read = await stream.read() read = await stream.read()
assert read == b"hello world from host a" assert read == hello_world_from_host_a
await stream.close() await stream.close()
@ -32,11 +35,8 @@ async def connect_write(host_a, host_b):
received = [] received = []
async def stream_handler(stream): async def stream_handler(stream):
while True: for message in messages:
try: received.append((await stream.read(len(message))).decode())
received.append((await stream.read()).decode())
except Exception: # exception is raised when other side close the stream ?
break
host_a.set_stream_handler(PROTOCOL_ID, stream_handler) host_a.set_stream_handler(PROTOCOL_ID, stream_handler)
@ -67,12 +67,8 @@ async def connect_read(host_a, host_b):
# Multiaddress of the destination peer is fetched from the peerstore using 'peerId'. # Multiaddress of the destination peer is fetched from the peerstore using 'peerId'.
stream = await host_b.new_stream(host_a.get_id(), [PROTOCOL_ID]) stream = await host_b.new_stream(host_a.get_id(), [PROTOCOL_ID])
received = [] received = []
# while True: Seems the close stream event from the other host is not received for message in messages:
for _ in range(5): received.append(await stream.read(len(message)))
try:
received.append(await stream.read())
except Exception: # exception is raised when other side close the stream ?
break
await stream.close() await stream.close()
assert received == messages assert received == messages

51
tests/libp2p/test_libp2p.py
View File

@ -3,6 +3,7 @@ import pytest
from libp2p.peer.peerinfo import info_from_p2p_addr from libp2p.peer.peerinfo import info_from_p2p_addr
from tests.utils import cleanup, set_up_nodes_by_transport_opt from tests.utils import cleanup, set_up_nodes_by_transport_opt
from tests.constants import MAX_READ_LEN
@pytest.mark.asyncio @pytest.mark.asyncio
@ -12,7 +13,7 @@ async def test_simple_messages():
async def stream_handler(stream): async def stream_handler(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack:" + read_string response = "ack:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
@ -28,7 +29,7 @@ async def test_simple_messages():
for message in messages: for message in messages:
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(MAX_READ_LEN)).decode()
assert response == ("ack:" + message) assert response == ("ack:" + message)
@ -43,7 +44,7 @@ async def test_double_response():
async def stream_handler(stream): async def stream_handler(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack1:" + read_string response = "ack1:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
@ -61,10 +62,10 @@ async def test_double_response():
for message in messages: for message in messages:
await stream.write(message.encode()) await stream.write(message.encode())
response1 = (await stream.read()).decode() response1 = (await stream.read(MAX_READ_LEN)).decode()
assert response1 == ("ack1:" + message) assert response1 == ("ack1:" + message)
response2 = (await stream.read()).decode() response2 = (await stream.read(MAX_READ_LEN)).decode()
assert response2 == ("ack2:" + message) assert response2 == ("ack2:" + message)
# Success, terminate pending tasks. # Success, terminate pending tasks.
@ -80,14 +81,14 @@ async def test_multiple_streams():
async def stream_handler_a(stream): async def stream_handler_a(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_a:" + read_string response = "ack_a:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
async def stream_handler_b(stream): async def stream_handler_b(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_b:" + read_string response = "ack_b:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
@ -111,8 +112,8 @@ async def test_multiple_streams():
await stream_a.write(a_message.encode()) await stream_a.write(a_message.encode())
await stream_b.write(b_message.encode()) await stream_b.write(b_message.encode())
response_a = (await stream_a.read()).decode() response_a = (await stream_a.read(MAX_READ_LEN)).decode()
response_b = (await stream_b.read()).decode() response_b = (await stream_b.read(MAX_READ_LEN)).decode()
assert response_a == ("ack_b:" + a_message) and response_b == ( assert response_a == ("ack_b:" + a_message) and response_b == (
"ack_a:" + b_message "ack_a:" + b_message
@ -129,21 +130,21 @@ async def test_multiple_streams_same_initiator_different_protocols():
async def stream_handler_a1(stream): async def stream_handler_a1(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_a1:" + read_string response = "ack_a1:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
async def stream_handler_a2(stream): async def stream_handler_a2(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_a2:" + read_string response = "ack_a2:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
async def stream_handler_a3(stream): async def stream_handler_a3(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_a3:" + read_string response = "ack_a3:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
@ -171,9 +172,9 @@ async def test_multiple_streams_same_initiator_different_protocols():
await stream_a2.write(a2_message.encode()) await stream_a2.write(a2_message.encode())
await stream_a3.write(a3_message.encode()) await stream_a3.write(a3_message.encode())
response_a1 = (await stream_a1.read()).decode() response_a1 = (await stream_a1.read(MAX_READ_LEN)).decode()
response_a2 = (await stream_a2.read()).decode() response_a2 = (await stream_a2.read(MAX_READ_LEN)).decode()
response_a3 = (await stream_a3.read()).decode() response_a3 = (await stream_a3.read(MAX_READ_LEN)).decode()
assert ( assert (
response_a1 == ("ack_a1:" + a1_message) response_a1 == ("ack_a1:" + a1_message)
@ -192,28 +193,28 @@ async def test_multiple_streams_two_initiators():
async def stream_handler_a1(stream): async def stream_handler_a1(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_a1:" + read_string response = "ack_a1:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
async def stream_handler_a2(stream): async def stream_handler_a2(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_a2:" + read_string response = "ack_a2:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
async def stream_handler_b1(stream): async def stream_handler_b1(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_b1:" + read_string response = "ack_b1:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
async def stream_handler_b2(stream): async def stream_handler_b2(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack_b2:" + read_string response = "ack_b2:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
@ -249,11 +250,11 @@ async def test_multiple_streams_two_initiators():
await stream_b1.write(b1_message.encode()) await stream_b1.write(b1_message.encode())
await stream_b2.write(b2_message.encode()) await stream_b2.write(b2_message.encode())
response_a1 = (await stream_a1.read()).decode() response_a1 = (await stream_a1.read(MAX_READ_LEN)).decode()
response_a2 = (await stream_a2.read()).decode() response_a2 = (await stream_a2.read(MAX_READ_LEN)).decode()
response_b1 = (await stream_b1.read()).decode() response_b1 = (await stream_b1.read(MAX_READ_LEN)).decode()
response_b2 = (await stream_b2.read()).decode() response_b2 = (await stream_b2.read(MAX_READ_LEN)).decode()
assert ( assert (
response_a1 == ("ack_a1:" + a1_message) response_a1 == ("ack_a1:" + a1_message)
@ -277,7 +278,7 @@ async def test_triangle_nodes_connection():
async def stream_handler(stream): async def stream_handler(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
response = "ack:" + read_string response = "ack:" + read_string
await stream.write(response.encode()) await stream.write(response.encode())
@ -320,7 +321,7 @@ async def test_triangle_nodes_connection():
for stream in streams: for stream in streams:
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(MAX_READ_LEN)).decode()
assert response == ("ack:" + message) assert response == ("ack:" + message)

106
tests/libp2p/test_notify.py
View File

@ -16,11 +16,10 @@ from libp2p import initialize_default_swarm, new_node
from libp2p.crypto.rsa import create_new_key_pair from libp2p.crypto.rsa import create_new_key_pair
from libp2p.host.basic_host import BasicHost from libp2p.host.basic_host import BasicHost
from libp2p.network.notifee_interface import INotifee from libp2p.network.notifee_interface import INotifee
from tests.utils import ( from tests.utils import cleanup, perform_two_host_set_up
cleanup, from tests.constants import MAX_READ_LEN
echo_stream_handler,
perform_two_host_set_up_custom_handler, ACK = "ack:"
)
class MyNotifee(INotifee): class MyNotifee(INotifee):
@ -67,38 +66,9 @@ class InvalidNotifee:
assert False assert False
async def perform_two_host_simple_set_up():
node_a = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
node_b = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
async def my_stream_handler(stream):
while True:
read_string = (await stream.read()).decode()
resp = "ack:" + read_string
await stream.write(resp.encode())
node_b.set_stream_handler("/echo/1.0.0", my_stream_handler)
# Associate the peer with local ip address (see default parameters of Libp2p())
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
return node_a, node_b
async def perform_two_host_simple_set_up_custom_handler(handler):
node_a = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
node_b = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
node_b.set_stream_handler("/echo/1.0.0", handler)
# Associate the peer with local ip address (see default parameters of Libp2p())
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
return node_a, node_b
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_one_notifier(): async def test_one_notifier():
node_a, node_b = await perform_two_host_set_up_custom_handler(echo_stream_handler) node_a, node_b = await perform_two_host_set_up()
# Add notifee for node_a # Add notifee for node_a
events = [] events = []
@ -113,11 +83,12 @@ async def test_one_notifier():
messages = ["hello", "hello"] messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()
@ -126,6 +97,7 @@ async def test_one_notifier():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_one_notifier_on_two_nodes(): async def test_one_notifier_on_two_nodes():
events_b = [] events_b = []
messages = ["hello", "hello"]
async def my_stream_handler(stream): async def my_stream_handler(stream):
# Ensure the connected and opened_stream events were hit in Notifee obj # Ensure the connected and opened_stream events were hit in Notifee obj
@ -135,13 +107,13 @@ async def test_one_notifier_on_two_nodes():
["connectedb", stream.mplex_conn], ["connectedb", stream.mplex_conn],
["opened_streamb", stream], ["opened_streamb", stream],
] ]
while True: for message in messages:
read_string = (await stream.read()).decode() read_string = (await stream.read(len(message))).decode()
resp = "ack:" + read_string resp = ACK + read_string
await stream.write(resp.encode()) await stream.write(resp.encode())
node_a, node_b = await perform_two_host_set_up_custom_handler(my_stream_handler) node_a, node_b = await perform_two_host_set_up(my_stream_handler)
# Add notifee for node_a # Add notifee for node_a
events_a = [] events_a = []
@ -157,13 +129,13 @@ async def test_one_notifier_on_two_nodes():
# node_a # node_a
assert events_a == [["connecteda", stream.mplex_conn], ["opened_streama", stream]] assert events_a == [["connecteda", stream.mplex_conn], ["opened_streama", stream]]
messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()
@ -172,6 +144,7 @@ async def test_one_notifier_on_two_nodes():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_one_notifier_on_two_nodes_with_listen(): async def test_one_notifier_on_two_nodes_with_listen():
events_b = [] events_b = []
messages = ["hello", "hello"]
node_a_key_pair = create_new_key_pair() node_a_key_pair = create_new_key_pair()
node_a_transport_opt = ["/ip4/127.0.0.1/tcp/0"] node_a_transport_opt = ["/ip4/127.0.0.1/tcp/0"]
@ -196,10 +169,9 @@ async def test_one_notifier_on_two_nodes_with_listen():
["connectedb", stream.mplex_conn], ["connectedb", stream.mplex_conn],
["opened_streamb", stream], ["opened_streamb", stream],
] ]
while True: for message in messages:
read_string = (await stream.read()).decode() read_string = (await stream.read(len(message))).decode()
resp = ACK + read_string
resp = "ack:" + read_string
await stream.write(resp.encode()) await stream.write(resp.encode())
# Add notifee for node_a # Add notifee for node_a
@ -222,13 +194,13 @@ async def test_one_notifier_on_two_nodes_with_listen():
# node_a # node_a
assert events_a == [["connecteda", stream.mplex_conn], ["opened_streama", stream]] assert events_a == [["connecteda", stream.mplex_conn], ["opened_streama", stream]]
messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()
@ -236,7 +208,7 @@ async def test_one_notifier_on_two_nodes_with_listen():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_two_notifiers(): async def test_two_notifiers():
node_a, node_b = await perform_two_host_set_up_custom_handler(echo_stream_handler) node_a, node_b = await perform_two_host_set_up()
# Add notifee for node_a # Add notifee for node_a
events0 = [] events0 = []
@ -255,11 +227,12 @@ async def test_two_notifiers():
messages = ["hello", "hello"] messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()
@ -269,7 +242,7 @@ async def test_two_notifiers():
async def test_ten_notifiers(): async def test_ten_notifiers():
num_notifiers = 10 num_notifiers = 10
node_a, node_b = await perform_two_host_set_up_custom_handler(echo_stream_handler) node_a, node_b = await perform_two_host_set_up()
# Add notifee for node_a # Add notifee for node_a
events_lst = [] events_lst = []
@ -290,11 +263,12 @@ async def test_ten_notifiers():
messages = ["hello", "hello"] messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()
@ -315,12 +289,12 @@ async def test_ten_notifiers_on_two_nodes():
["opened_streamb" + str(i), stream], ["opened_streamb" + str(i), stream],
] ]
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
resp = "ack:" + read_string resp = ACK + read_string
await stream.write(resp.encode()) await stream.write(resp.encode())
node_a, node_b = await perform_two_host_set_up_custom_handler(my_stream_handler) node_a, node_b = await perform_two_host_set_up(my_stream_handler)
# Add notifee for node_a and node_b # Add notifee for node_a and node_b
events_lst_a = [] events_lst_a = []
@ -343,11 +317,12 @@ async def test_ten_notifiers_on_two_nodes():
messages = ["hello", "hello"] messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()
@ -357,7 +332,7 @@ async def test_ten_notifiers_on_two_nodes():
async def test_invalid_notifee(): async def test_invalid_notifee():
num_notifiers = 10 num_notifiers = 10
node_a, node_b = await perform_two_host_set_up_custom_handler(echo_stream_handler) node_a, node_b = await perform_two_host_set_up()
# Add notifee for node_a # Add notifee for node_a
events_lst = [] events_lst = []
@ -372,11 +347,12 @@ async def test_invalid_notifee():
# given that InvalidNotifee should not have been added as a notifee) # given that InvalidNotifee should not have been added as a notifee)
messages = ["hello", "hello"] messages = ["hello", "hello"]
for message in messages: for message in messages:
expected_resp = ACK + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read()).decode() response = (await stream.read(len(expected_resp))).decode()
assert response == ("ack:" + message) assert response == expected_resp
# Success, terminate pending tasks. # Success, terminate pending tasks.
await cleanup() await cleanup()

17
tests/protocol_muxer/test_protocol_muxer.py
View File

@ -1,7 +1,7 @@
import pytest import pytest
from libp2p.protocol_muxer.exceptions import MultiselectClientError from libp2p.protocol_muxer.exceptions import MultiselectClientError
from tests.utils import cleanup, set_up_nodes_by_transport_opt from tests.utils import cleanup, set_up_nodes_by_transport_opt, echo_stream_handler
# TODO: Add tests for multiple streams being opened on different # TODO: Add tests for multiple streams being opened on different
# protocols through the same connection # protocols through the same connection
@ -18,14 +18,8 @@ async def perform_simple_test(
transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]]
(node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list) (node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list)
async def stream_handler(stream):
while True:
read_string = (await stream.read()).decode()
response = "ack:" + read_string
await stream.write(response.encode())
for protocol in protocols_with_handlers: for protocol in protocols_with_handlers:
node_b.set_stream_handler(protocol, stream_handler) node_b.set_stream_handler(protocol, echo_stream_handler)
# Associate the peer with local ip address (see default parameters of Libp2p()) # Associate the peer with local ip address (see default parameters of Libp2p())
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10) node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
@ -33,11 +27,10 @@ async def perform_simple_test(
stream = await node_a.new_stream(node_b.get_id(), protocols_for_client) stream = await node_a.new_stream(node_b.get_id(), protocols_for_client)
messages = ["hello" + str(x) for x in range(10)] messages = ["hello" + str(x) for x in range(10)]
for message in messages: for message in messages:
expected_resp = "ack:" + message
await stream.write(message.encode()) await stream.write(message.encode())
response = (await stream.read(len(expected_resp))).decode()
response = (await stream.read()).decode() assert response == expected_resp
assert response == ("ack:" + message)
assert expected_selected_protocol == stream.get_protocol() assert expected_selected_protocol == stream.get_protocol()

6
tests/utils.py
View File

@ -6,6 +6,8 @@ import multiaddr
from libp2p import new_node from libp2p import new_node
from libp2p.peer.peerinfo import info_from_p2p_addr from libp2p.peer.peerinfo import info_from_p2p_addr
from tests.constants import MAX_READ_LEN
async def connect(node1, node2): async def connect(node1, node2):
""" """
@ -38,13 +40,13 @@ async def set_up_nodes_by_transport_opt(transport_opt_list):
async def echo_stream_handler(stream): async def echo_stream_handler(stream):
while True: while True:
read_string = (await stream.read()).decode() read_string = (await stream.read(MAX_READ_LEN)).decode()
resp = "ack:" + read_string resp = "ack:" + read_string
await stream.write(resp.encode()) await stream.write(resp.encode())
async def perform_two_host_set_up_custom_handler(handler): async def perform_two_host_set_up(handler=echo_stream_handler):
transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]]
(node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list) (node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list)