71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
|
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()
|