Merge branch 'main' into feature/bootstrap

tests/core/identity/identify_push/test_identify_push.py

@@ -35,6 +35,8 @@ from tests.utils.factories import (
 )
 from tests.utils.utils import (
     create_mock_connections,
+    run_host_forever,
+    wait_until_listening,
 )
 
 logger = logging.getLogger("libp2p.identity.identify-push-test")
@@ -503,3 +505,91 @@ async def test_push_identify_to_peers_respects_concurrency_limit():
     assert state["max_observed"] <= CONCURRENCY_LIMIT, (
         f"Max concurrency observed: {state['max_observed']}"
     )
+
+
+@pytest.mark.trio
+async def test_all_peers_receive_identify_push_with_semaphore(security_protocol):
+    dummy_peers = []
+
+    async with host_pair_factory(security_protocol=security_protocol) as (host_a, _):
+        # Create dummy peers
+        for _ in range(50):
+            key_pair = create_new_key_pair()
+            dummy_host = new_host(key_pair=key_pair)
+            dummy_host.set_stream_handler(
+                ID_PUSH, identify_push_handler_for(dummy_host)
+            )
+            listen_addr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
+            dummy_peers.append((dummy_host, listen_addr))
+
+        async with trio.open_nursery() as nursery:
+            # Start all dummy hosts
+            for host, listen_addr in dummy_peers:
+                nursery.start_soon(run_host_forever, host, listen_addr)
+
+            # Wait for all hosts to finish setting up listeners
+            for host, _ in dummy_peers:
+                await wait_until_listening(host)
+
+            # Now connect host_a → dummy peers
+            for host, _ in dummy_peers:
+                await host_a.connect(info_from_p2p_addr(host.get_addrs()[0]))
+
+            await push_identify_to_peers(
+                host_a,
+            )
+
+            await trio.sleep(0.5)
+
+            peer_id_a = host_a.get_id()
+            for host, _ in dummy_peers:
+                dummy_peerstore = host.get_peerstore()
+                assert peer_id_a in dummy_peerstore.peer_ids()
+
+            nursery.cancel_scope.cancel()
+
+
+@pytest.mark.trio
+async def test_all_peers_receive_identify_push_with_semaphore_under_high_peer_load(
+    security_protocol,
+):
+    dummy_peers = []
+
+    async with host_pair_factory(security_protocol=security_protocol) as (host_a, _):
+        # Create dummy peers
+        # Breaking with more than 500 peers
+        # Trio have a async tasks limit of 1000
+        for _ in range(499):
+            key_pair = create_new_key_pair()
+            dummy_host = new_host(key_pair=key_pair)
+            dummy_host.set_stream_handler(
+                ID_PUSH, identify_push_handler_for(dummy_host)
+            )
+            listen_addr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
+            dummy_peers.append((dummy_host, listen_addr))
+
+        async with trio.open_nursery() as nursery:
+            # Start all dummy hosts
+            for host, listen_addr in dummy_peers:
+                nursery.start_soon(run_host_forever, host, listen_addr)
+
+            # Wait for all hosts to finish setting up listeners
+            for host, _ in dummy_peers:
+                await wait_until_listening(host)
+
+            # Now connect host_a → dummy peers
+            for host, _ in dummy_peers:
+                await host_a.connect(info_from_p2p_addr(host.get_addrs()[0]))
+
+            await push_identify_to_peers(
+                host_a,
+            )
+
+            await trio.sleep(0.5)
+
+            peer_id_a = host_a.get_id()
+            for host, _ in dummy_peers:
+                dummy_peerstore = host.get_peerstore()
+                assert peer_id_a in dummy_peerstore.peer_ids()
+
+            nursery.cancel_scope.cancel()

tests/core/peer/test_peerdata.py

@@ -6,10 +6,12 @@ from multiaddr import Multiaddr
 from libp2p.crypto.secp256k1 import (
     create_new_key_pair,
 )
+from libp2p.peer.id import ID
 from libp2p.peer.peerdata import (
     PeerData,
     PeerDataError,
 )
+from libp2p.peer.peerstore import PeerStore
 
 MOCK_ADDR = Multiaddr("/ip4/127.0.0.1/tcp/4001")
 MOCK_KEYPAIR = create_new_key_pair()
@@ -39,6 +41,59 @@ def test_set_protocols():
     assert peer_data.get_protocols() == protocols
 
 
+# Test case when removing protocols:
+def test_remove_protocols():
+    peer_data = PeerData()
+    protocols: Sequence[str] = ["protocol1", "protocol2"]
+    peer_data.set_protocols(protocols)
+
+    peer_data.remove_protocols(["protocol1"])
+    assert peer_data.get_protocols() == ["protocol2"]
+
+
+# Test case when clearing the protocol list:
+def test_clear_protocol_data():
+    peer_data = PeerData()
+    protocols: Sequence[str] = ["protocol1", "protocol2"]
+    peer_data.set_protocols(protocols)
+
+    peer_data.clear_protocol_data()
+    assert peer_data.get_protocols() == []
+
+
+# Test case when supports protocols:
+def test_supports_protocols():
+    peer_data = PeerData()
+    peer_data.set_protocols(["protocol1", "protocol2", "protocol3"])
+
+    input_protocols = ["protocol1", "protocol4", "protocol2"]
+    supported = peer_data.supports_protocols(input_protocols)
+
+    assert supported == ["protocol1", "protocol2"]
+
+
+# Test case for first supported protocol is found
+def test_first_supported_protocol_found():
+    peer_data = PeerData()
+    peer_data.set_protocols(["protocolA", "protocolB"])
+
+    input_protocols = ["protocolC", "protocolB", "protocolA"]
+    first = peer_data.first_supported_protocol(input_protocols)
+
+    assert first == "protocolB"
+
+
+# Test case for first supported protocol not found
+def test_first_supported_protocol_none():
+    peer_data = PeerData()
+    peer_data.set_protocols(["protocolX", "protocolY"])
+
+    input_protocols = ["protocolA", "protocolB"]
+    first = peer_data.first_supported_protocol(input_protocols)
+
+    assert first == "None supported"
+
+
 # Test case when adding addresses
 def test_add_addrs():
     peer_data = PeerData()
@@ -81,6 +136,15 @@ def test_get_metadata_key_not_found():
         peer_data.get_metadata("nonexistent_key")
 
 
+# Test case for clearing metadata
+def test_clear_metadata():
+    peer_data = PeerData()
+    peer_data.metadata = {"key1": "value1", "key2": "value2"}
+
+    peer_data.clear_metadata()
+    assert peer_data.metadata == {}
+
+
 # Test case for adding public key
 def test_add_pubkey():
     peer_data = PeerData()
@@ -107,3 +171,71 @@ def test_get_privkey_not_found():
     peer_data = PeerData()
     with pytest.raises(PeerDataError):
         peer_data.get_privkey()
+
+
+# Test case for returning all the peers with stored keys
+def test_peer_with_keys():
+    peer_store = PeerStore()
+    peer_id_1 = ID(b"peer1")
+    peer_id_2 = ID(b"peer2")
+
+    peer_data_1 = PeerData()
+    peer_data_2 = PeerData()
+
+    peer_data_1.pubkey = MOCK_PUBKEY
+    peer_data_2.pubkey = None
+
+    peer_store.peer_data_map = {
+        peer_id_1: peer_data_1,
+        peer_id_2: peer_data_2,
+    }
+
+    assert peer_store.peer_with_keys() == [peer_id_1]
+
+
+# Test case for clearing the key book
+def test_clear_keydata():
+    peer_store = PeerStore()
+    peer_id = ID(b"peer123")
+    peer_data = PeerData()
+
+    peer_data.pubkey = MOCK_PUBKEY
+    peer_data.privkey = MOCK_PRIVKEY
+    peer_store.peer_data_map = {peer_id: peer_data}
+
+    peer_store.clear_keydata(peer_id)
+
+    assert peer_data.pubkey is None
+    assert peer_data.privkey is None
+
+
+# Test case for recording latency for the first time
+def test_record_latency_initial():
+    peer_data = PeerData()
+    assert peer_data.latency_EWMA() == 0
+
+    peer_data.record_latency(100.0)
+    assert peer_data.latency_EWMA() == 100.0
+
+
+# Test case for updating latency
+def test_record_latency_updates_ewma():
+    peer_data = PeerData()
+    peer_data.record_latency(100.0)  # first measurement
+    first = peer_data.latency_EWMA()
+
+    peer_data.record_latency(50.0)  # second measurement
+    second = peer_data.latency_EWMA()
+
+    assert second < first  # EWMA should have smoothed downward
+    assert second > 50.0  # Not as low as the new latency
+    assert second != first
+
+
+def test_clear_metrics():
+    peer_data = PeerData()
+    peer_data.record_latency(200.0)
+    assert peer_data.latency_EWMA() == 200.0
+
+    peer_data.clear_metrics()
+    assert peer_data.latency_EWMA() == 0

tests/core/peer/test_peerstore.py

@@ -2,6 +2,7 @@ import time
 
 import pytest
 from multiaddr import Multiaddr
+import trio
 
 from libp2p.peer.id import ID
 from libp2p.peer.peerstore import (
@@ -89,3 +90,33 @@ def test_peers():
     store.add_addr(ID(b"peer3"), Multiaddr("/ip4/127.0.0.1/tcp/4001"), 10)
 
     assert set(store.peer_ids()) == {ID(b"peer1"), ID(b"peer2"), ID(b"peer3")}
+
+
+@pytest.mark.trio
+async def test_addr_stream_yields_new_addrs():
+    store = PeerStore()
+    peer_id = ID(b"peer1")
+    addr1 = Multiaddr("/ip4/127.0.0.1/tcp/4001")
+    addr2 = Multiaddr("/ip4/127.0.0.1/tcp/4002")
+
+    collected = []
+
+    async def consume_addrs():
+        async for addr in store.addr_stream(peer_id):
+            collected.append(addr)
+            if len(collected) == 2:
+                break
+
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(consume_addrs)
+        await trio.sleep(2)  # Give time for the stream to start
+
+        store.add_addr(peer_id, addr1, ttl=10)
+        await trio.sleep(0.2)
+        store.add_addr(peer_id, addr2, ttl=10)
+        await trio.sleep(0.2)
+
+        # After collecting expected addresses, cancel the stream
+        nursery.cancel_scope.cancel()
+
+    assert collected == [addr1, addr2]

tests/core/stream_muxer/test_yamux_interleaving.py

@@ -0,0 +1,199 @@
+import logging
+
+import pytest
+import trio
+from trio.testing import (
+    memory_stream_pair,
+)
+
+from libp2p.abc import IRawConnection
+from libp2p.crypto.ed25519 import (
+    create_new_key_pair,
+)
+from libp2p.peer.id import (
+    ID,
+)
+from libp2p.security.insecure.transport import (
+    InsecureTransport,
+)
+from libp2p.stream_muxer.yamux.yamux import (
+    Yamux,
+    YamuxStream,
+)
+
+
+class TrioStreamAdapter(IRawConnection):
+    """Adapter to make trio memory streams work with libp2p."""
+
+    def __init__(self, send_stream, receive_stream, is_initiator=False):
+        self.send_stream = send_stream
+        self.receive_stream = receive_stream
+        self.is_initiator = is_initiator
+
+    async def write(self, data: bytes) -> None:
+        logging.debug(f"Attempting to write {len(data)} bytes")
+        with trio.move_on_after(2):
+            await self.send_stream.send_all(data)
+
+    async def read(self, n: int | None = None) -> bytes:
+        if n is None or n <= 0:
+            raise ValueError("Reading unbounded or zero bytes not supported")
+        logging.debug(f"Attempting to read {n} bytes")
+        with trio.move_on_after(2):
+            data = await self.receive_stream.receive_some(n)
+            logging.debug(f"Read {len(data)} bytes")
+            return data
+
+    async def close(self) -> None:
+        logging.debug("Closing stream")
+        await self.send_stream.aclose()
+        await self.receive_stream.aclose()
+
+    def get_remote_address(self) -> tuple[str, int] | None:
+        """Return None since this is a test adapter without real network info."""
+        return None
+
+
+@pytest.fixture
+def key_pair():
+    return create_new_key_pair()
+
+
+@pytest.fixture
+def peer_id(key_pair):
+    return ID.from_pubkey(key_pair.public_key)
+
+
+@pytest.fixture
+async def secure_conn_pair(key_pair, peer_id):
+    """Create a pair of secure connections for testing."""
+    logging.debug("Setting up secure_conn_pair")
+    client_send, server_receive = memory_stream_pair()
+    server_send, client_receive = memory_stream_pair()
+
+    client_rw = TrioStreamAdapter(client_send, client_receive)
+    server_rw = TrioStreamAdapter(server_send, server_receive)
+
+    insecure_transport = InsecureTransport(key_pair)
+
+    async def run_outbound(nursery_results):
+        with trio.move_on_after(5):
+            client_conn = await insecure_transport.secure_outbound(client_rw, peer_id)
+            logging.debug("Outbound handshake complete")
+            nursery_results["client"] = client_conn
+
+    async def run_inbound(nursery_results):
+        with trio.move_on_after(5):
+            server_conn = await insecure_transport.secure_inbound(server_rw)
+            logging.debug("Inbound handshake complete")
+            nursery_results["server"] = server_conn
+
+    nursery_results = {}
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(run_outbound, nursery_results)
+        nursery.start_soon(run_inbound, nursery_results)
+        await trio.sleep(0.1)  # Give tasks a chance to finish
+
+    client_conn = nursery_results.get("client")
+    server_conn = nursery_results.get("server")
+
+    if client_conn is None or server_conn is None:
+        raise RuntimeError("Handshake failed: client_conn or server_conn is None")
+
+    logging.debug("secure_conn_pair setup complete")
+    return client_conn, server_conn
+
+
+@pytest.fixture
+async def yamux_pair(secure_conn_pair, peer_id):
+    """Create a pair of Yamux multiplexers for testing."""
+    logging.debug("Setting up yamux_pair")
+    client_conn, server_conn = secure_conn_pair
+    client_yamux = Yamux(client_conn, peer_id, is_initiator=True)
+    server_yamux = Yamux(server_conn, peer_id, is_initiator=False)
+    async with trio.open_nursery() as nursery:
+        with trio.move_on_after(5):
+            nursery.start_soon(client_yamux.start)
+            nursery.start_soon(server_yamux.start)
+            await trio.sleep(0.1)
+            logging.debug("yamux_pair started")
+        yield client_yamux, server_yamux
+    logging.debug("yamux_pair cleanup")
+
+
+@pytest.mark.trio
+async def test_yamux_race_condition_without_locks(yamux_pair):
+    """
+    Test for race-around/interleaving in Yamux streams,when reading in
+    segments of data.
+    This launches concurrent writers/readers on both sides of a stream.
+    If there is no proper locking, the received data may be interleaved
+    or corrupted.
+
+    The test creates structured messages and verifies they are received
+    intact and in order.
+    Without proper locking, concurrent read/write operations could cause
+    data corruption
+    or message interleaving, which this test will catch.
+    """
+    client_yamux, server_yamux = yamux_pair
+    client_stream: YamuxStream = await client_yamux.open_stream()
+    server_stream: YamuxStream = await server_yamux.accept_stream()
+    MSG_COUNT = 10
+    MSG_SIZE = 256 * 1024  # At max,only DEFAULT_WINDOW_SIZE bytes can be read
+    client_msgs = [
+        f"CLIENT-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"C") for i in range(MSG_COUNT)
+    ]
+    server_msgs = [
+        f"SERVER-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"S") for i in range(MSG_COUNT)
+    ]
+    client_received = []
+    server_received = []
+
+    async def writer(stream, msgs, name):
+        """Write messages with minimal delays to encourage race conditions."""
+        for i, msg in enumerate(msgs):
+            await stream.write(msg)
+            # Yield control frequently to encourage interleaving
+            if i % 5 == 0:
+                await trio.sleep(0.005)
+
+    async def reader(stream, received, name):
+        """Read messages and store them for verification."""
+        for i in range(MSG_COUNT):
+            data = await stream.read(MSG_SIZE)
+            received.append(data)
+            if i % 3 == 0:
+                await trio.sleep(0.001)
+
+    # Running all operations concurrently
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(writer, client_stream, client_msgs, "client")
+        nursery.start_soon(writer, server_stream, server_msgs, "server")
+        nursery.start_soon(reader, client_stream, client_received, "client")
+        nursery.start_soon(reader, server_stream, server_received, "server")
+
+    assert len(client_received) == MSG_COUNT, (
+        f"Client received {len(client_received)} messages, expected {MSG_COUNT}"
+    )
+    assert len(server_received) == MSG_COUNT, (
+        f"Server received {len(server_received)} messages, expected {MSG_COUNT}"
+    )
+    assert client_received == server_msgs, (
+        "Client did not receive server messages in order or intact!"
+    )
+    assert server_received == client_msgs, (
+        "Server did not receive client messages in order or intact!"
+    )
+    for i, msg in enumerate(client_received):
+        assert len(msg) == MSG_SIZE, (
+            f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}"
+        )
+
+    for i, msg in enumerate(server_received):
+        assert len(msg) == MSG_SIZE, (
+            f"Server message {i} has wrong size: {len(msg)} != {MSG_SIZE}"
+        )
+
+    await client_stream.close()
+    await server_stream.close()

tests/core/stream_muxer/test_yamux_interleaving_EOF.py

@@ -0,0 +1,195 @@
+import logging
+
+import pytest
+import trio
+from trio.testing import (
+    memory_stream_pair,
+)
+
+from libp2p.abc import IRawConnection
+from libp2p.crypto.ed25519 import (
+    create_new_key_pair,
+)
+from libp2p.peer.id import (
+    ID,
+)
+from libp2p.security.insecure.transport import (
+    InsecureTransport,
+)
+from libp2p.stream_muxer.exceptions import MuxedStreamEOF
+from libp2p.stream_muxer.yamux.yamux import (
+    Yamux,
+    YamuxStream,
+)
+
+
+class TrioStreamAdapter(IRawConnection):
+    """Adapter to make trio memory streams work with libp2p."""
+
+    def __init__(self, send_stream, receive_stream, is_initiator=False):
+        self.send_stream = send_stream
+        self.receive_stream = receive_stream
+        self.is_initiator = is_initiator
+
+    async def write(self, data: bytes) -> None:
+        logging.debug(f"Attempting to write {len(data)} bytes")
+        with trio.move_on_after(2):
+            await self.send_stream.send_all(data)
+
+    async def read(self, n: int | None = None) -> bytes:
+        if n is None or n <= 0:
+            raise ValueError("Reading unbounded or zero bytes not supported")
+        logging.debug(f"Attempting to read {n} bytes")
+        with trio.move_on_after(2):
+            data = await self.receive_stream.receive_some(n)
+            logging.debug(f"Read {len(data)} bytes")
+            return data
+
+    async def close(self) -> None:
+        logging.debug("Closing stream")
+        await self.send_stream.aclose()
+        await self.receive_stream.aclose()
+
+    def get_remote_address(self) -> tuple[str, int] | None:
+        """Return None since this is a test adapter without real network info."""
+        return None
+
+
+@pytest.fixture
+def key_pair():
+    return create_new_key_pair()
+
+
+@pytest.fixture
+def peer_id(key_pair):
+    return ID.from_pubkey(key_pair.public_key)
+
+
+@pytest.fixture
+async def secure_conn_pair(key_pair, peer_id):
+    """Create a pair of secure connections for testing."""
+    logging.debug("Setting up secure_conn_pair")
+    client_send, server_receive = memory_stream_pair()
+    server_send, client_receive = memory_stream_pair()
+
+    client_rw = TrioStreamAdapter(client_send, client_receive)
+    server_rw = TrioStreamAdapter(server_send, server_receive)
+
+    insecure_transport = InsecureTransport(key_pair)
+
+    async def run_outbound(nursery_results):
+        with trio.move_on_after(5):
+            client_conn = await insecure_transport.secure_outbound(client_rw, peer_id)
+            logging.debug("Outbound handshake complete")
+            nursery_results["client"] = client_conn
+
+    async def run_inbound(nursery_results):
+        with trio.move_on_after(5):
+            server_conn = await insecure_transport.secure_inbound(server_rw)
+            logging.debug("Inbound handshake complete")
+            nursery_results["server"] = server_conn
+
+    nursery_results = {}
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(run_outbound, nursery_results)
+        nursery.start_soon(run_inbound, nursery_results)
+        await trio.sleep(0.1)  # Give tasks a chance to finish
+
+    client_conn = nursery_results.get("client")
+    server_conn = nursery_results.get("server")
+
+    if client_conn is None or server_conn is None:
+        raise RuntimeError("Handshake failed: client_conn or server_conn is None")
+
+    logging.debug("secure_conn_pair setup complete")
+    return client_conn, server_conn
+
+
+@pytest.fixture
+async def yamux_pair(secure_conn_pair, peer_id):
+    """Create a pair of Yamux multiplexers for testing."""
+    logging.debug("Setting up yamux_pair")
+    client_conn, server_conn = secure_conn_pair
+    client_yamux = Yamux(client_conn, peer_id, is_initiator=True)
+    server_yamux = Yamux(server_conn, peer_id, is_initiator=False)
+    async with trio.open_nursery() as nursery:
+        with trio.move_on_after(5):
+            nursery.start_soon(client_yamux.start)
+            nursery.start_soon(server_yamux.start)
+            await trio.sleep(0.1)
+            logging.debug("yamux_pair started")
+        yield client_yamux, server_yamux
+    logging.debug("yamux_pair cleanup")
+
+
+@pytest.mark.trio
+async def test_yamux_race_condition_without_locks(yamux_pair):
+    """
+    Test for race-around/interleaving in Yamux streams,when reading till
+    EOF is being used.
+    This launches concurrent writers/readers on both sides of a stream.
+    If there is no proper locking, the received data may be interleaved
+    or corrupted.
+
+    The test creates structured messages and verifies they are received
+    intact and in order.
+    Without proper locking, concurrent read/write operations could cause
+    data corruption
+    or message interleaving, which this test will catch.
+    """
+    client_yamux, server_yamux = yamux_pair
+    client_stream: YamuxStream = await client_yamux.open_stream()
+    server_stream: YamuxStream = await server_yamux.accept_stream()
+    MSG_COUNT = 1
+    MSG_SIZE = 512 * 1024
+    client_msgs = [
+        f"CLIENT-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"C") for i in range(MSG_COUNT)
+    ]
+    server_msgs = [
+        f"SERVER-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"S") for i in range(MSG_COUNT)
+    ]
+    client_received = []
+    server_received = []
+
+    async def writer(stream, msgs, name):
+        """Write messages with minimal delays to encourage race conditions."""
+        for i, msg in enumerate(msgs):
+            await stream.write(msg)
+            # Yield control frequently to encourage interleaving
+            if i % 5 == 0:
+                await trio.sleep(0.005)
+
+    async def reader(stream, received, name):
+        """Read messages and store them for verification."""
+        try:
+            data = await stream.read()
+            if data:
+                received.append(data)
+        except MuxedStreamEOF:
+            pass
+
+    # Running all operations concurrently
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(writer, client_stream, client_msgs, "client")
+        nursery.start_soon(writer, server_stream, server_msgs, "server")
+        nursery.start_soon(reader, client_stream, client_received, "client")
+        nursery.start_soon(reader, server_stream, server_received, "server")
+
+    assert client_received == server_msgs, (
+        "Client did not receive server messages in order or intact!"
+    )
+    assert server_received == client_msgs, (
+        "Server did not receive client messages in order or intact!"
+    )
+    for i, msg in enumerate(client_received):
+        assert len(msg) == MSG_SIZE, (
+            f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}"
+        )
+
+    for i, msg in enumerate(server_received):
+        assert len(msg) == MSG_SIZE, (
+            f"Server message {i} has wrong size: {len(msg)} != {MSG_SIZE}"
+        )
+
+    await client_stream.close()
+    await server_stream.close()

tests/utils/utils.py

@@ -2,13 +2,30 @@ from unittest.mock import (
     MagicMock,
 )
 
+import trio
 
-def create_mock_connections() -> dict:
+from libp2p.abc import IHost
+
+
+def create_mock_connections(count: int = 50) -> dict:
     connections = {}
 
-    for i in range(1, 31):
+    for i in range(1, count):
         peer_id = f"peer-{i}"
         mock_conn = MagicMock(name=f"INetConn-{i}")
         connections[peer_id] = mock_conn
 
     return connections
+
+
+async def run_host_forever(host: IHost, addr):
+    async with host.run([addr]):
+        await trio.sleep_forever()
+
+
+async def wait_until_listening(host, timeout=3):
+    with trio.move_on_after(timeout):
+        while not host.get_addrs():
+            await trio.sleep(0.05)
+        return
+    raise RuntimeError("Timed out waiting for host to get an address")