Fix MplexStream.read

This commit is contained in:
mhchia
2019-09-06 17:26:40 +08:00
parent 95926b7376
commit 649a230776
8 changed files with 116 additions and 143 deletions

View File

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

View File

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

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.host.basic_host import BasicHost
from libp2p.network.notifee_interface import INotifee
from tests.utils import (
cleanup,
echo_stream_handler,
perform_two_host_set_up_custom_handler,
)
from tests.utils import cleanup, perform_two_host_set_up
from tests.constants import MAX_READ_LEN
ACK = "ack:"
class MyNotifee(INotifee):
@ -67,38 +66,9 @@ class InvalidNotifee:
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
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
events = []
@ -113,11 +83,12 @@ async def test_one_notifier():
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()
@ -126,6 +97,7 @@ async def test_one_notifier():
@pytest.mark.asyncio
async def test_one_notifier_on_two_nodes():
events_b = []
messages = ["hello", "hello"]
async def my_stream_handler(stream):
# 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],
["opened_streamb", stream],
]
while True:
read_string = (await stream.read()).decode()
for message in messages:
read_string = (await stream.read(len(message))).decode()
resp = "ack:" + read_string
resp = ACK + read_string
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
events_a = []
@ -157,13 +129,13 @@ async def test_one_notifier_on_two_nodes():
# node_a
assert events_a == [["connecteda", stream.mplex_conn], ["opened_streama", stream]]
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()
@ -172,6 +144,7 @@ async def test_one_notifier_on_two_nodes():
@pytest.mark.asyncio
async def test_one_notifier_on_two_nodes_with_listen():
events_b = []
messages = ["hello", "hello"]
node_a_key_pair = create_new_key_pair()
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],
["opened_streamb", stream],
]
while True:
read_string = (await stream.read()).decode()
resp = "ack:" + read_string
for message in messages:
read_string = (await stream.read(len(message))).decode()
resp = ACK + read_string
await stream.write(resp.encode())
# Add notifee for node_a
@ -222,13 +194,13 @@ async def test_one_notifier_on_two_nodes_with_listen():
# node_a
assert events_a == [["connecteda", stream.mplex_conn], ["opened_streama", stream]]
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()
@ -236,7 +208,7 @@ async def test_one_notifier_on_two_nodes_with_listen():
@pytest.mark.asyncio
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
events0 = []
@ -255,11 +227,12 @@ async def test_two_notifiers():
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()
@ -269,7 +242,7 @@ async def test_two_notifiers():
async def test_ten_notifiers():
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
events_lst = []
@ -290,11 +263,12 @@ async def test_ten_notifiers():
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()
@ -315,12 +289,12 @@ async def test_ten_notifiers_on_two_nodes():
["opened_streamb" + str(i), stream],
]
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())
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
events_lst_a = []
@ -343,11 +317,12 @@ async def test_ten_notifiers_on_two_nodes():
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()
@ -357,7 +332,7 @@ async def test_ten_notifiers_on_two_nodes():
async def test_invalid_notifee():
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
events_lst = []
@ -372,11 +347,12 @@ async def test_invalid_notifee():
# given that InvalidNotifee should not have been added as a notifee)
messages = ["hello", "hello"]
for message in messages:
expected_resp = ACK + message
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.
await cleanup()

View File

@ -1,7 +1,7 @@
import pytest
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
# 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"]]
(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:
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())
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)
messages = ["hello" + str(x) for x in range(10)]
for message in messages:
expected_resp = "ack:" + message
await stream.write(message.encode())
response = (await stream.read()).decode()
assert response == ("ack:" + message)
response = (await stream.read(len(expected_resp))).decode()
assert response == expected_resp
assert expected_selected_protocol == stream.get_protocol()

View File

@ -6,6 +6,8 @@ import multiaddr
from libp2p import new_node
from libp2p.peer.peerinfo import info_from_p2p_addr
from tests.constants import MAX_READ_LEN
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):
while True:
read_string = (await stream.read()).decode()
read_string = (await stream.read(MAX_READ_LEN)).decode()
resp = "ack:" + read_string
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"]]
(node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list)