diff --git a/tests/host/test_ping.py b/tests/host/test_ping.py new file mode 100644 index 00000000..d6f02449 --- /dev/null +++ b/tests/host/test_ping.py @@ -0,0 +1,70 @@ +import asyncio +import secrets + +import pytest + +from libp2p.host.ping import ID, PING_LENGTH, handle_ping +from libp2p.peer.peerinfo import info_from_p2p_addr +from tests.utils import set_up_nodes_by_transport_opt + + +def _add_ping_to(host): + host.set_stream_handler(ID, handle_ping) + + +@pytest.mark.asyncio +async def test_ping_once(): + transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] + (host_a, host_b) = await set_up_nodes_by_transport_opt(transport_opt_list) + + # NOTE: this will fail after we add ping as a default handler + # as a forced reminder to fix this test by removing the calls to + # `_add_ping_to` + assert host_a.get_mux().handlers == {} + assert host_b.get_mux().handlers == {} + _add_ping_to(host_a) + _add_ping_to(host_b) + + addr = host_a.get_addrs()[0] + info = info_from_p2p_addr(addr) + await host_b.connect(info) + + stream = await host_b.new_stream(host_a.get_id(), (ID,)) + some_ping = secrets.token_bytes(PING_LENGTH) + await stream.write(some_ping) + some_pong = await stream.read(PING_LENGTH) + assert some_ping == some_pong + await stream.close() + + +SOME_PING_COUNT = 3 + + +@pytest.mark.asyncio +async def test_ping_several(): + transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] + (host_a, host_b) = await set_up_nodes_by_transport_opt(transport_opt_list) + + # NOTE: this will fail after we add ping as a default handler + # as a forced reminder to fix this test by removing the calls to + # `_add_ping_to` + assert host_a.get_mux().handlers == {} + assert host_b.get_mux().handlers == {} + _add_ping_to(host_a) + _add_ping_to(host_b) + + addr = host_a.get_addrs()[0] + info = info_from_p2p_addr(addr) + await host_b.connect(info) + + stream = await host_b.new_stream(host_a.get_id(), (ID,)) + for _ in range(SOME_PING_COUNT): + some_ping = secrets.token_bytes(PING_LENGTH) + await stream.write(some_ping) + some_pong = await stream.read(PING_LENGTH) + assert some_ping == some_pong + # NOTE: simulate some time to sleep to mirror a real + # world usage where a peer sends pings on some periodic interval + # NOTE: this interval can be `0` for this test. + await asyncio.sleep(0) + await stream.close()