#!/bin/env python3 import asyncio import sys import click import multiaddr import pytest from libp2p.libp2p import * from peer.peerinfo import info_from_p2p_addr from protocol_muxer.multiselect_client import MultiselectClientError PROTOCOL_ID = '/chat/1.0.0' async def hello_world(host_a, 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') await stream.close() host_a.set_stream_handler(PROTOCOL_ID, stream_handler) # 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') read = await stream.read() assert read == b'hello world from host a' await stream.close() async def connect_write(host_a, host_b): messages = [b'data %d' % i for i in range(5)] async def stream_handler(stream): received = [] while True: try: received.append((await stream.read()).decode()) except Exception: # exception is raised when other side close the stream ? break await stream.close() assert received == messages host_a.set_stream_handler(PROTOCOL_ID, stream_handler) # 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]) for message in messages: await stream.write(message) await stream.close() async def connect_read(host_a, host_b): messages = [b'data %d' % i for i in range(5)] async def stream_handler(stream): for message in messages: await stream.write(message) await stream.close() host_a.set_stream_handler(PROTOCOL_ID, stream_handler) # 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]) 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 await stream.close() assert received == messages async def no_common_protocol(host_a, host_b): messages = [b'data %d' % i for i in range(5)] async def stream_handler(stream): for message in messages: await stream.write(message) await stream.close() host_a.set_stream_handler(PROTOCOL_ID, stream_handler) # try to creates a new new with a procotol not known by the other host with pytest.raises(MultiselectClientError): _ = await host_b.new_stream(host_a.get_id(), ['/fakeproto/0.0.1']) @pytest.mark.asyncio @pytest.mark.parametrize("test", [ (hello_world), (connect_write), (connect_read), (no_common_protocol), ]) async def test_chat(test): host_a = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"]) host_b = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"]) addr = host_a.get_addrs()[0] info = info_from_p2p_addr(addr) await host_b.connect(info) await test(host_a, host_b)