Implemented addr_stream in the peerstore

This commit is contained in:
lla-dane
2025-07-04 14:53:44 +05:30
parent 51c08de1bc
commit d1c31483bd
2 changed files with 64 additions and 0 deletions

View File

@ -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,36 @@ 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")
# 🔧 Pre-initialize peer in peer_data_map
# store.add_addr(peer_id, Multiaddr("/ip4/127.0.0.1/tcp/0"), ttl=1)
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]