Files
py-libp2p/tests/core/host/test_connected_peers.py
2025-03-27 14:21:45 -06:00

114 lines
3.7 KiB
Python

import pytest
from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
from tests.utils.factories import (
HostFactory,
)
async def connect_two(host_a, host_b, host_c):
# Initially all of the hosts are disconnected
assert (len(host_a.get_connected_peers())) == 0
assert (len(host_b.get_connected_peers())) == 0
assert (len(host_c.get_connected_peers())) == 0
# Connecting hostA with hostB
addr = host_b.get_addrs()[0]
info = info_from_p2p_addr(addr)
await host_a.connect(info)
# Since hostA and hostB are connected now
assert (len(host_a.get_connected_peers())) == 1
assert (len(host_b.get_connected_peers())) == 1
assert (len(host_c.get_connected_peers())) == 0
assert host_a.get_connected_peers()[0] == host_b.get_id()
assert host_b.get_connected_peers()[0] == host_a.get_id()
async def connect_three_cyclic(host_a, host_b, host_c):
# Connecting hostA with hostB
addr = host_b.get_addrs()[0]
infoB = info_from_p2p_addr(addr)
await host_a.connect(infoB)
# Connecting hostB with hostC
addr = host_c.get_addrs()[0]
infoC = info_from_p2p_addr(addr)
await host_b.connect(infoC)
# Connecting hostC with hostA
addr = host_a.get_addrs()[0]
infoA = info_from_p2p_addr(addr)
await host_c.connect(infoA)
# Performing checks
assert (len(host_a.get_connected_peers())) == 2
assert (len(host_b.get_connected_peers())) == 2
assert (len(host_c.get_connected_peers())) == 2
assert host_a.get_connected_peers() == [host_b.get_id(), host_c.get_id()]
assert host_b.get_connected_peers() == [host_a.get_id(), host_c.get_id()]
assert host_c.get_connected_peers() == [host_b.get_id(), host_a.get_id()]
async def connect_two_to_one(host_a, host_b, host_c):
# Connecting hostA and hostC to hostB
addr = host_b.get_addrs()[0]
infoB = info_from_p2p_addr(addr)
await host_a.connect(infoB)
await host_c.connect(infoB)
# Performing checks
assert (len(host_a.get_connected_peers())) == 1
assert (len(host_b.get_connected_peers())) == 2
assert (len(host_c.get_connected_peers())) == 1
assert host_a.get_connected_peers() == [host_b.get_id()]
assert host_b.get_connected_peers() == [host_a.get_id(), host_c.get_id()]
assert host_c.get_connected_peers() == [host_b.get_id()]
async def connect_and_disconnect(host_a, host_b, host_c):
# Connecting hostB to hostA and hostC
addr = host_a.get_addrs()[0]
infoA = info_from_p2p_addr(addr)
await host_b.connect(infoA)
addr = host_c.get_addrs()[0]
infoC = info_from_p2p_addr(addr)
await host_b.connect(infoC)
# Performing checks
assert (len(host_a.get_connected_peers())) == 1
assert (len(host_b.get_connected_peers())) == 2
assert (len(host_c.get_connected_peers())) == 1
assert host_a.get_connected_peers() == [host_b.get_id()]
assert host_b.get_connected_peers() == [host_a.get_id(), host_c.get_id()]
assert host_c.get_connected_peers() == [host_b.get_id()]
# Disconnecting hostB and hostA
await host_b.disconnect(host_a.get_id())
# Performing checks
assert (len(host_a.get_connected_peers())) == 0
assert (len(host_b.get_connected_peers())) == 1
assert (len(host_c.get_connected_peers())) == 1
assert host_b.get_connected_peers() == [host_c.get_id()]
assert host_c.get_connected_peers() == [host_b.get_id()]
@pytest.mark.parametrize(
"test",
[
(connect_two),
(connect_three_cyclic),
(connect_two_to_one),
(connect_and_disconnect),
],
)
@pytest.mark.trio
async def test_connected_peers(test, security_protocol):
async with HostFactory.create_batch_and_listen(
3, security_protocol=security_protocol
) as hosts:
await test(hosts[0], hosts[1], hosts[2])