Merge pull request #196 from NIC619/fix_peer_id

Refactor `peer.ID` class and only takes in `bytes` type argument

libp2p/__init__.py

@@ -4,7 +4,7 @@ from Crypto.PublicKey import RSA
 
 from .security.insecure_security import InsecureTransport
 from .peer.peerstore import PeerStore
-from .peer.id import id_from_public_key
+from .peer.id import ID
 from .network.swarm import Swarm
 from .host.basic_host import BasicHost
 from .transport.upgrader import TransportUpgrader
@@ -29,7 +29,7 @@ async def cleanup_done_tasks():
 
 def generate_id():
     new_key = RSA.generate(2048, e=65537)
-    new_id = id_from_public_key(new_key.publickey())
+    new_id = ID.from_pubkey(new_key.publickey())
     # private_key = new_key.exportKey("PEM")
     return new_id
 
@@ -47,7 +47,7 @@ def initialize_default_kademlia_router(ksize=20, alpha=3, id_opt=None, storage=N
     if not id_opt:
         id_opt = generate_id()
 
-    node_id = id_opt.get_raw_id()
+    node_id = id_opt.to_bytes()
     server = KademliaServer(ksize=ksize, alpha=alpha, node_id=node_id, storage=storage)
     return KadmeliaPeerRouter(server)
 

libp2p/kademlia/crawling.py

@@ -61,7 +61,7 @@ class SpiderCrawl:
 
         dicts = {}
         for peer in self.nearest.get_uncontacted()[:count]:
-            dicts[peer.peer_id] = rpcmethod(peer, self.node)
+            dicts[peer.peer_id_bytes] = rpcmethod(peer, self.node)
             self.nearest.mark_contacted(peer)
         found = await gather_dict(dicts)
         return await self._nodes_found(found)
@@ -125,7 +125,7 @@ class ValueSpiderCrawl(SpiderCrawl):
 
         peer = self.nearest_without_value.popleft()
         if peer:
-            await self.protocol.call_store(peer, self.node.peer_id, value)
+            await self.protocol.call_store(peer, self.node.peer_id_bytes, value)
         return value
 
 

libp2p/kademlia/kad_peerinfo.py

@@ -16,9 +16,8 @@ class KadPeerInfo(PeerInfo):
     def __init__(self, peer_id, peer_data=None):
         super(KadPeerInfo, self).__init__(peer_id, peer_data)
 
-        self.peer_id_obj = peer_id
+        self.peer_id_bytes = peer_id.to_bytes()
-        self.peer_id = peer_id.get_raw_id()
+        self.xor_id = peer_id.xor_id
-        self.xor_id = peer_id.get_xor_id()
 
         self.addrs = peer_data.get_addrs() if peer_data else None
 
@@ -39,17 +38,17 @@ class KadPeerInfo(PeerInfo):
         """
         Enables use of Node as a tuple - i.e., tuple(node) works.
         """
-        return iter([self.peer_id, self.ip, self.port])
+        return iter([self.peer_id_bytes, self.ip, self.port])
 
     def __repr__(self):
-        return repr([self.xor_id, self.ip, self.port, self.peer_id])
+        return repr([self.xor_id, self.ip, self.port, self.peer_id_bytes])
 
     def __str__(self):
         return "%s:%s" % (self.ip, str(self.port))
 
     def encode(self):
         return (
-            str(self.peer_id)
+            str(self.peer_id_bytes)
             + "\n"
             + str("/ip4/" + str(self.ip) + "/udp/" + str(self.port))
         )
@@ -85,13 +84,13 @@ class KadPeerHeap:
             return
         nheap = []
         for distance, node in self.heap:
-            if node.peer_id not in peers:
+            if node.peer_id_bytes not in peers:
                 heapq.heappush(nheap, (distance, node))
         self.heap = nheap
 
     def get_node(self, node_id):
         for _, node in self.heap:
-            if node.peer_id == node_id:
+            if node.peer_id_bytes == node_id:
                 return node
         return None
 
@@ -99,10 +98,10 @@ class KadPeerHeap:
         return len(self.get_uncontacted()) == 0
 
     def get_ids(self):
-        return [n.peer_id for n in self]
+        return [n.peer_id_bytes for n in self]
 
     def mark_contacted(self, node):
-        self.contacted.add(node.peer_id)
+        self.contacted.add(node.peer_id_bytes)
 
     def popleft(self):
         return heapq.heappop(self.heap)[1] if self else None
@@ -130,16 +129,18 @@ class KadPeerHeap:
 
     def __contains__(self, node):
         for _, other in self.heap:
-            if node.peer_id == other.peer_id:
+            if node.peer_id_bytes == other.peer_id_bytes:
                 return True
         return False
 
     def get_uncontacted(self):
-        return [n for n in self if n.peer_id not in self.contacted]
+        return [n for n in self if n.peer_id_bytes not in self.contacted]
 
 
-def create_kad_peerinfo(raw_node_id=None, sender_ip=None, sender_port=None):
+def create_kad_peerinfo(node_id_bytes=None, sender_ip=None, sender_port=None):
-    node_id = ID(raw_node_id) if raw_node_id else ID(digest(random.getrandbits(255)))
+    node_id = (
+        ID(node_id_bytes) if node_id_bytes else ID(digest(random.getrandbits(255)))
+    )
     peer_data = None
     if sender_ip and sender_port:
         peer_data = PeerData()  # pylint: disable=no-value-for-parameter

libp2p/kademlia/network.py

@@ -130,7 +130,7 @@ class KademliaServer:
         return await spider.find()
 
     async def bootstrap_node(self, addr):
-        result = await self.protocol.ping(addr, self.node.peer_id)
+        result = await self.protocol.ping(addr, self.node.peer_id_bytes)
         return create_kad_peerinfo(result[1], addr[0], addr[1]) if result[0] else None
 
     async def get(self, key):
@@ -170,7 +170,7 @@ class KademliaServer:
         """
         neighbors = self.protocol.router.find_neighbors(self.node)
         return [
-            await self.protocol.call_add_provider(n, key, self.node.peer_id)
+            await self.protocol.call_add_provider(n, key, self.node.peer_id_bytes)
             for n in neighbors
         ]
 
@@ -214,7 +214,7 @@ class KademliaServer:
         data = {
             "ksize": self.ksize,
             "alpha": self.alpha,
-            "id": self.node.peer_id,
+            "id": self.node.peer_id_bytes,
             "neighbors": self.bootstrappable_neighbors(),
         }
         if not data["neighbors"]:

libp2p/kademlia/protocol.py

@@ -44,7 +44,7 @@ class KademliaProtocol(RPCProtocol):
         source = create_kad_peerinfo(nodeid, sender[0], sender[1])
 
         self.welcome_if_new(source)
-        return self.source_node.peer_id
+        return self.source_node.peer_id_bytes
 
     def rpc_store(self, sender, nodeid, key, value):
         source = create_kad_peerinfo(nodeid, sender[0], sender[1])
@@ -106,39 +106,39 @@ class KademliaProtocol(RPCProtocol):
         keynode = create_kad_peerinfo(key)
         neighbors = self.router.find_neighbors(keynode)
         for neighbor in neighbors:
-            if neighbor.peer_id != record:
+            if neighbor.peer_id_bytes != record:
-                providers.append(neighbor.peer_id)
+                providers.append(neighbor.peer_id_bytes)
 
         return providers
 
     async def call_find_node(self, node_to_ask, node_to_find):
         address = (node_to_ask.ip, node_to_ask.port)
         result = await self.find_node(
-            address, self.source_node.peer_id, node_to_find.peer_id
+            address, self.source_node.peer_id_bytes, node_to_find.peer_id_bytes
         )
         return self.handle_call_response(result, node_to_ask)
 
     async def call_find_value(self, node_to_ask, node_to_find):
         address = (node_to_ask.ip, node_to_ask.port)
         result = await self.find_value(
-            address, self.source_node.peer_id, node_to_find.peer_id
+            address, self.source_node.peer_id_bytes, node_to_find.peer_id_bytes
         )
         return self.handle_call_response(result, node_to_ask)
 
     async def call_ping(self, node_to_ask):
         address = (node_to_ask.ip, node_to_ask.port)
-        result = await self.ping(address, self.source_node.peer_id)
+        result = await self.ping(address, self.source_node.peer_id_bytes)
         return self.handle_call_response(result, node_to_ask)
 
     async def call_store(self, node_to_ask, key, value):
         address = (node_to_ask.ip, node_to_ask.port)
-        result = await self.store(address, self.source_node.peer_id, key, value)
+        result = await self.store(address, self.source_node.peer_id_bytes, key, value)
         return self.handle_call_response(result, node_to_ask)
 
     async def call_add_provider(self, node_to_ask, key, provider_id):
         address = (node_to_ask.ip, node_to_ask.port)
         result = await self.add_provider(
-            address, self.source_node.peer_id, key, provider_id
+            address, self.source_node.peer_id_bytes, key, provider_id
         )
 
         return self.handle_call_response(result, node_to_ask)

libp2p/kademlia/routing.py

@@ -35,24 +35,24 @@ class KBucket:
         two = KBucket(midpoint + 1, self.range[1], self.ksize)
         for node in self.nodes.values():
             bucket = one if node.xor_id <= midpoint else two
-            bucket.nodes[node.peer_id] = node
+            bucket.nodes[node.peer_id_bytes] = node
         return (one, two)
 
     def remove_node(self, node):
-        if node.peer_id not in self.nodes:
+        if node.peer_id_bytes not in self.nodes:
             return
 
         # delete node, and see if we can add a replacement
-        del self.nodes[node.peer_id]
+        del self.nodes[node.peer_id_bytes]
         if self.replacement_nodes:
             newnode = self.replacement_nodes.pop()
-            self.nodes[newnode.peer_id] = newnode
+            self.nodes[newnode.peer_id_bytes] = newnode
 
     def has_in_range(self, node):
         return self.range[0] <= node.xor_id <= self.range[1]
 
     def is_new_node(self, node):
-        return node.peer_id not in self.nodes
+        return node.peer_id_bytes not in self.nodes
 
     def add_node(self, node):
         """
@@ -62,11 +62,11 @@ class KBucket:
         If the bucket is full, keep track of node in a replacement list,
         per section 4.1 of the paper.
         """
-        if node.peer_id in self.nodes:
+        if node.peer_id_bytes in self.nodes:
-            del self.nodes[node.peer_id]
+            del self.nodes[node.peer_id_bytes]
-            self.nodes[node.peer_id] = node
+            self.nodes[node.peer_id_bytes] = node
         elif len(self) < self.ksize:
-            self.nodes[node.peer_id] = node
+            self.nodes[node.peer_id_bytes] = node
         else:
             self.replacement_nodes.push(node)
             return False
@@ -74,7 +74,7 @@ class KBucket:
 
     def depth(self):
         vals = self.nodes.values()
-        sprefix = shared_prefix([bytes_to_bit_string(n.peer_id) for n in vals])
+        sprefix = shared_prefix([bytes_to_bit_string(n.peer_id_bytes) for n in vals])
         return len(sprefix)
 
     def head(self):
@@ -186,7 +186,7 @@ class RoutingTable:
         nodes = []
         for neighbor in TableTraverser(self, node):
             notexcluded = exclude is None or not neighbor.same_home_as(exclude)
-            if neighbor.peer_id != node.peer_id and notexcluded:
+            if neighbor.peer_id_bytes != node.peer_id_bytes and notexcluded:
                 heapq.heappush(nodes, (node.distance_to(neighbor), neighbor))
             if len(nodes) == k:
                 break

libp2p/network/network_interface.py

@@ -58,9 +58,9 @@ class INetwork(ABC):
         """
 
     @abstractmethod
-    async def listen(self, *args: Sequence[Multiaddr]) -> bool:
+    async def listen(self, *multiaddrs: Sequence[Multiaddr]) -> bool:
         """
-        :param *args: one or many multiaddrs to start listening on
+        :param multiaddrs: one or many multiaddrs to start listening on
         :return: True if at least one success
         """
 

libp2p/network/swarm.py

@@ -3,7 +3,7 @@ from typing import Awaitable, Callable, Dict, List, Sequence
 
 from multiaddr import Multiaddr
 
-from libp2p.peer.id import ID, id_b58_decode
+from libp2p.peer.id import ID
 from libp2p.peer.peerstore import PeerStore
 from libp2p.protocol_muxer.multiselect import Multiselect
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
@@ -163,12 +163,12 @@ class Swarm(INetwork):
 
         return net_stream
 
-    async def listen(self, *args: Sequence[Multiaddr]) -> bool:
+    async def listen(self, *multiaddrs: Sequence[Multiaddr]) -> bool:
         """
-        :param *args: one or many multiaddrs to start listening on
+        :param multiaddrs: one or many multiaddrs to start listening on
         :return: true if at least one success
 
-        For each multiaddr in args
+        For each multiaddr
             Check if a listener for multiaddr exists already
             If listener already exists, continue
             Otherwise:
@@ -177,7 +177,7 @@ class Swarm(INetwork):
                 Call listener listen with the multiaddr
                 Map multiaddr to listener
         """
-        for multiaddr in args:
+        for multiaddr in multiaddrs:
             if str(multiaddr) in self.listeners:
                 return True
 
@@ -185,7 +185,7 @@ class Swarm(INetwork):
                 reader: asyncio.StreamReader, writer: asyncio.StreamWriter
             ) -> None:
                 # Read in first message (should be peer_id of initiator) and ack
-                peer_id = id_b58_decode((await reader.read(1024)).decode())
+                peer_id = ID.from_base58((await reader.read(1024)).decode())
 
                 writer.write("received peer id".encode())
                 await writer.drain()

libp2p/peer/id.py

@@ -18,69 +18,68 @@ MAX_INLINE_KEY_LENGTH = 42
 
 class ID:
 
-    _id_str: bytes
+    _bytes: bytes
+    _xor_id: int = None
+    _b58_str: str = None
 
-    def __init__(self, id_str: bytes) -> None:
+    def __init__(self, peer_id_bytes: bytes) -> None:
-        self._id_str = id_str
+        self._bytes = peer_id_bytes
+
+    @property
+    def xor_id(self) -> int:
+        if not self._xor_id:
+            self._xor_id = int(digest(self._bytes).hex(), 16)
+        return self._xor_id
 
     def to_bytes(self) -> bytes:
-        return self._id_str
+        return self._bytes
 
-    def get_raw_id(self) -> bytes:
+    def to_base58(self) -> str:
-        return self._id_str
+        if not self._b58_str:
+            self._b58_str = base58.b58encode(self._bytes).decode()
+        return self._b58_str
 
-    def pretty(self) -> str:
+    def __bytes__(self) -> bytes:
-        return base58.b58encode(self._id_str).decode()
+        return self._bytes
 
-    def get_xor_id(self) -> int:
+    __repr__ = __str__ = pretty = to_string = to_base58
-        return int(digest(self.get_raw_id()).hex(), 16)
-
-    def __str__(self) -> str:
-        pid = self.pretty()
-        return pid
-
-    __repr__ = __str__
 
     def __eq__(self, other: object) -> bool:
-        # pylint: disable=protected-access
+        # pylint: disable=protected-access, no-else-return
-        if not isinstance(other, ID):
+        if isinstance(other, str):
+            return self.to_base58() == other
+        elif isinstance(other, bytes):
+            return self._bytes == other
+        elif isinstance(other, ID):
+            return self._bytes == other._bytes
+        else:
             return NotImplemented
-        return self._id_str == other._id_str
 
     def __hash__(self) -> int:
-        return hash(self._id_str)
+        return hash(self._bytes)
 
+    @classmethod
+    def from_base58(cls, b58_encoded_peer_id_str: str) -> "ID":
+        peer_id_bytes = base58.b58decode(b58_encoded_peer_id_str)
+        pid = ID(peer_id_bytes)
+        return pid
 
-def id_b58_encode(peer_id: ID) -> str:
+    @classmethod
-    """
+    def from_pubkey(cls, key: RsaKey) -> "ID":
-    return a b58-encoded string
+        # export into binary format
-    """
+        key_bin = key.exportKey("DER")
-    # pylint: disable=protected-access
-    return base58.b58encode(peer_id.get_raw_id()).decode()
 
+        algo: int = multihash.Func.sha2_256
+        # TODO: seems identity is not yet supported in pymultihash
+        # if len(b) <= MAX_INLINE_KEY_LENGTH:
+        #     algo multihash.func.identity
 
-def id_b58_decode(peer_id_str: str) -> ID:
+        mh_digest: multihash.Multihash = multihash.digest(key_bin, algo)
-    """
+        return cls(mh_digest.encode())
-    return a base58-decoded peer ID
-    """
-    return ID(base58.b58decode(peer_id_str))
 
-
+    @classmethod
-def id_from_public_key(key: RsaKey) -> ID:
+    def from_privkey(cls, key: RsaKey) -> "ID":
-    # export into binary format
+        return cls.from_pubkey(key.publickey())
-    key_bin = key.exportKey("DER")
-
-    algo: int = multihash.Func.sha2_256
-    # TODO: seems identity is not yet supported in pymultihash
-    # if len(b) <= MAX_INLINE_KEY_LENGTH:
-    #     algo multihash.func.identity
-
-    mh_digest: multihash.Multihash = multihash.digest(key_bin, algo)
-    return ID(mh_digest.encode())
-
-
-def id_from_private_key(key: RsaKey) -> ID:
-    return id_from_public_key(key.publickey())
 
 
 def digest(data: Union[str, bytes]) -> bytes:

libp2p/peer/peerinfo.py

@@ -2,7 +2,7 @@ from typing import List
 
 import multiaddr
 
-from .id import ID, id_b58_decode
+from .id import ID
 from .peerdata import PeerData
 
 
@@ -39,7 +39,7 @@ def info_from_p2p_addr(addr: multiaddr.Multiaddr) -> PeerInfo:
 
     # make sure the /p2p value parses as a peer.ID
     peer_id_str: str = p2p_part.value_for_protocol(multiaddr.protocols.P_P2P)
-    peer_id: ID = id_b58_decode(peer_id_str)
+    peer_id: ID = ID.from_base58(peer_id_str)
 
     # we might have received just an / p2p part, which means there's no addr.
     if len(parts) > 1:

libp2p/pubsub/floodsub.py

@@ -1,6 +1,6 @@
 from typing import Iterable, List, Sequence
 
-from libp2p.peer.id import ID, id_b58_decode
+from libp2p.peer.id import ID
 
 from .pb import rpc_pb2
 from .pubsub import Pubsub
@@ -73,7 +73,7 @@ class FloodSub(IPubsubRouter):
         )
         rpc_msg = rpc_pb2.RPC(publish=[pubsub_msg])
         for peer_id in peers_gen:
-            stream = self.pubsub.peers[str(peer_id)]
+            stream = self.pubsub.peers[peer_id]
             # FIXME: We should add a `WriteMsg` similar to write delimited messages.
             #   Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/master/comm.go#L107
             await stream.write(rpc_msg.SerializeToString())
@@ -105,11 +105,9 @@ class FloodSub(IPubsubRouter):
         for topic in topic_ids:
             if topic not in self.pubsub.peer_topics:
                 continue
-            for peer_id_str in self.pubsub.peer_topics[topic]:
+            for peer_id in self.pubsub.peer_topics[topic]:
-                peer_id = id_b58_decode(peer_id_str)
                 if peer_id in (msg_forwarder, origin):
                     continue
-                # FIXME: Should change `self.pubsub.peers` to Dict[PeerID, ...]
+                if peer_id not in self.pubsub.peers:
-                if str(peer_id) not in self.pubsub.peers:
                     continue
                 yield peer_id

libp2p/pubsub/gossipsub.py

@@ -3,7 +3,7 @@ import asyncio
 import random
 from typing import Any, Dict, Iterable, List, Set, Sequence
 
-from libp2p.peer.id import ID, id_b58_decode
+from libp2p.peer.id import ID
 
 from .mcache import MessageCache
 from .pb import rpc_pb2
@@ -25,20 +25,15 @@ class GossipSub(IPubsubRouter):
 
     time_to_live: int
 
-    # FIXME: Should be changed to `Dict[str, List[ID]]`
+    mesh: Dict[str, List[ID]]
-    mesh: Dict[str, List[str]]
+    fanout: Dict[str, List[ID]]
-    # FIXME: Should be changed to `Dict[str, List[ID]]`
-    fanout: Dict[str, List[str]]
 
-    # FIXME: Should be changed to `Dict[ID, str]`
+    peers_to_protocol: Dict[ID, str]
-    peers_to_protocol: Dict[str, str]
 
     time_since_last_publish: Dict[str, int]
 
-    # FIXME: Should be changed to List[ID]
+    peers_gossipsub: List[ID]
-    peers_gossipsub: List[str]
+    peers_floodsub: List[ID]
-    # FIXME: Should be changed to List[ID]
-    peers_floodsub: List[str]
 
     mcache: MessageCache
 
@@ -115,27 +110,25 @@ class GossipSub(IPubsubRouter):
 
         # Add peer to the correct peer list
         peer_type = GossipSub.get_peer_type(protocol_id)
-        peer_id_str = str(peer_id)
 
-        self.peers_to_protocol[peer_id_str] = protocol_id
+        self.peers_to_protocol[peer_id] = protocol_id
 
         if peer_type == "gossip":
-            self.peers_gossipsub.append(peer_id_str)
+            self.peers_gossipsub.append(peer_id)
         elif peer_type == "flood":
-            self.peers_floodsub.append(peer_id_str)
+            self.peers_floodsub.append(peer_id)
 
     def remove_peer(self, peer_id: ID) -> None:
         """
         Notifies the router that a peer has been disconnected
         :param peer_id: id of peer to remove
         """
-        peer_id_str = str(peer_id)
+        del self.peers_to_protocol[peer_id]
-        del self.peers_to_protocol[peer_id_str]
 
-        if peer_id_str in self.peers_gossipsub:
+        if peer_id in self.peers_gossipsub:
-            self.peers_gossipsub.remove(peer_id_str)
+            self.peers_gossipsub.remove(peer_id)
-        if peer_id_str in self.peers_gossipsub:
+        if peer_id in self.peers_gossipsub:
-            self.peers_floodsub.remove(peer_id_str)
+            self.peers_floodsub.remove(peer_id)
 
     async def handle_rpc(self, rpc: rpc_pb2.Message, sender_peer_id: ID) -> None:
         """
@@ -145,21 +138,20 @@ class GossipSub(IPubsubRouter):
         :param sender_peer_id: id of the peer who sent the message
         """
         control_message = rpc.control
-        sender_peer_id_str = str(sender_peer_id)
 
         # Relay each rpc control message to the appropriate handler
         if control_message.ihave:
             for ihave in control_message.ihave:
-                await self.handle_ihave(ihave, sender_peer_id_str)
+                await self.handle_ihave(ihave, sender_peer_id)
         if control_message.iwant:
             for iwant in control_message.iwant:
-                await self.handle_iwant(iwant, sender_peer_id_str)
+                await self.handle_iwant(iwant, sender_peer_id)
         if control_message.graft:
             for graft in control_message.graft:
-                await self.handle_graft(graft, sender_peer_id_str)
+                await self.handle_graft(graft, sender_peer_id)
         if control_message.prune:
             for prune in control_message.prune:
-                await self.handle_prune(prune, sender_peer_id_str)
+                await self.handle_prune(prune, sender_peer_id)
 
     async def publish(self, msg_forwarder: ID, pubsub_msg: rpc_pb2.Message) -> None:
         # pylint: disable=too-many-locals
@@ -175,7 +167,7 @@ class GossipSub(IPubsubRouter):
         )
         rpc_msg = rpc_pb2.RPC(publish=[pubsub_msg])
         for peer_id in peers_gen:
-            stream = self.pubsub.peers[str(peer_id)]
+            stream = self.pubsub.peers[peer_id]
             # FIXME: We should add a `WriteMsg` similar to write delimited messages.
             #   Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/master/comm.go#L107
             # TODO: Go use `sendRPC`, which possibly piggybacks gossip/control messages.
@@ -197,16 +189,14 @@ class GossipSub(IPubsubRouter):
                 continue
 
             # floodsub peers
-            for peer_id_str in self.pubsub.peer_topics[topic]:
+            for peer_id in self.pubsub.peer_topics[topic]:
                 # FIXME: `gossipsub.peers_floodsub` can be changed to `gossipsub.peers` in go.
                 #   This will improve the efficiency when searching for a peer's protocol id.
-                if peer_id_str in self.peers_floodsub:
+                if peer_id in self.peers_floodsub:
-                    peer_id = id_b58_decode(peer_id_str)
                     send_to.add(peer_id)
 
             # gossipsub peers
-            # FIXME: Change `str` to `ID`
+            in_topic_gossipsub_peers: List[ID] = None
-            in_topic_gossipsub_peers: List[str] = None
             # TODO: Do we need to check `topic in self.pubsub.my_topics`?
             if topic in self.mesh:
                 in_topic_gossipsub_peers = self.mesh[topic]
@@ -222,8 +212,8 @@ class GossipSub(IPubsubRouter):
                         topic, self.degree, []
                     )
                 in_topic_gossipsub_peers = self.fanout[topic]
-            for peer_id_str in in_topic_gossipsub_peers:
+            for peer_id in in_topic_gossipsub_peers:
-                send_to.add(id_b58_decode(peer_id_str))
+                send_to.add(peer_id)
         # Excludes `msg_forwarder` and `origin`
         yield from send_to.difference([msg_forwarder, origin])
 
@@ -241,8 +231,7 @@ class GossipSub(IPubsubRouter):
         self.mesh[topic] = []
 
         topic_in_fanout: bool = topic in self.fanout
-        # FIXME: Should be changed to `List[ID]`
+        fanout_peers: List[ID] = self.fanout[topic] if topic_in_fanout else []
-        fanout_peers: List[str] = self.fanout[topic] if topic_in_fanout else []
         fanout_size = len(fanout_peers)
         if not topic_in_fanout or (topic_in_fanout and fanout_size < self.degree):
             # There are less than D peers (let this number be x)
@@ -290,14 +279,8 @@ class GossipSub(IPubsubRouter):
             return "flood"
         return "unknown"
 
-    # FIXME: type of `peers` should be changed to `List[ID]`
-    # FIXME: type of `msg_sender` and `origin_id` should be changed to `ID`
     async def deliver_messages_to_peers(
-        self,
+        self, peers: List[ID], msg_sender: ID, origin_id: ID, serialized_packet: bytes
-        peers: List[str],
-        msg_sender: str,
-        origin_id: str,
-        serialized_packet: bytes,
     ) -> None:
         for peer_id_in_topic in peers:
             # Forward to all peers that are not the
@@ -338,8 +321,7 @@ class GossipSub(IPubsubRouter):
                     topic, self.degree - num_mesh_peers_in_topic, self.mesh[topic]
                 )
 
-                # FIXME: Should be changed to `List[ID]`
+                fanout_peers_not_in_mesh: List[ID] = [
-                fanout_peers_not_in_mesh: List[str] = [
                     peer for peer in selected_peers if peer not in self.mesh[topic]
                 ]
                 for peer in fanout_peers_not_in_mesh:
@@ -351,7 +333,6 @@ class GossipSub(IPubsubRouter):
 
             if num_mesh_peers_in_topic > self.degree_high:
                 # Select |mesh[topic]| - D peers from mesh[topic]
-                # FIXME: Should be changed to `List[ID]`
                 selected_peers = GossipSub.select_from_minus(
                     num_mesh_peers_in_topic - self.degree, self.mesh[topic], []
                 )
@@ -461,15 +442,13 @@ class GossipSub(IPubsubRouter):
 
         return selection
 
-    # FIXME: type of `minus` should be changed to type `Sequence[ID]`
-    # FIXME: return type should be changed to type `List[ID]`
     def _get_in_topic_gossipsub_peers_from_minus(
-        self, topic: str, num_to_select: int, minus: Sequence[str]
+        self, topic: str, num_to_select: int, minus: Sequence[ID]
-    ) -> List[str]:
+    ) -> List[ID]:
         gossipsub_peers_in_topic = [
-            peer_str
+            peer_id
-            for peer_str in self.pubsub.peer_topics[topic]
+            for peer_id in self.pubsub.peer_topics[topic]
-            if peer_str in self.peers_gossipsub
+            if peer_id in self.peers_gossipsub
         ]
         return self.select_from_minus(
             num_to_select, gossipsub_peers_in_topic, list(minus)
@@ -478,15 +457,13 @@ class GossipSub(IPubsubRouter):
     # RPC handlers
 
     async def handle_ihave(
-        self, ihave_msg: rpc_pb2.Message, sender_peer_id: str
+        self, ihave_msg: rpc_pb2.Message, sender_peer_id: ID
     ) -> None:
         """
         Checks the seen set and requests unknown messages with an IWANT message.
         """
         # from_id_bytes = ihave_msg.from_id
 
-        from_id_str = sender_peer_id
-
         # Get list of all seen (seqnos, from) from the (seqno, from) tuples in seen_messages cache
         seen_seqnos_and_peers = [
             seqno_and_from for seqno_and_from in self.pubsub.seen_messages.keys()
@@ -503,16 +480,14 @@ class GossipSub(IPubsubRouter):
 
         # Request messages with IWANT message
         if msg_ids_wanted:
-            await self.emit_iwant(msg_ids_wanted, from_id_str)
+            await self.emit_iwant(msg_ids_wanted, sender_peer_id)
 
     async def handle_iwant(
-        self, iwant_msg: rpc_pb2.Message, sender_peer_id: str
+        self, iwant_msg: rpc_pb2.Message, sender_peer_id: ID
     ) -> None:
         """
         Forwards all request messages that are present in mcache to the requesting peer.
         """
-        from_id_str = sender_peer_id
-
         # FIXME: Update type of message ID
         # FIXME: Find a better way to parse the msg ids
         msg_ids: List[Any] = [literal_eval(msg) for msg in iwant_msg.messageIDs]
@@ -539,41 +514,36 @@ class GossipSub(IPubsubRouter):
         rpc_msg: bytes = packet.SerializeToString()
 
         # 3) Get the stream to this peer
-        # TODO: Should we pass in from_id or from_id_str here?
+        peer_stream = self.pubsub.peers[sender_peer_id]
-        peer_stream = self.pubsub.peers[from_id_str]
 
         # 4) And write the packet to the stream
         await peer_stream.write(rpc_msg)
 
     async def handle_graft(
-        self, graft_msg: rpc_pb2.Message, sender_peer_id: str
+        self, graft_msg: rpc_pb2.Message, sender_peer_id: ID
     ) -> None:
         topic: str = graft_msg.topicID
 
-        from_id_str = sender_peer_id
-
         # Add peer to mesh for topic
         if topic in self.mesh:
-            if from_id_str not in self.mesh[topic]:
+            if sender_peer_id not in self.mesh[topic]:
-                self.mesh[topic].append(from_id_str)
+                self.mesh[topic].append(sender_peer_id)
         else:
             # Respond with PRUNE if not subscribed to the topic
             await self.emit_prune(topic, sender_peer_id)
 
     async def handle_prune(
-        self, prune_msg: rpc_pb2.Message, sender_peer_id: str
+        self, prune_msg: rpc_pb2.Message, sender_peer_id: ID
     ) -> None:
         topic: str = prune_msg.topicID
 
-        from_id_str = sender_peer_id
-
         # Remove peer from mesh for topic, if peer is in topic
-        if topic in self.mesh and from_id_str in self.mesh[topic]:
+        if topic in self.mesh and sender_peer_id in self.mesh[topic]:
-            self.mesh[topic].remove(from_id_str)
+            self.mesh[topic].remove(sender_peer_id)
 
     # RPC emitters
 
-    async def emit_ihave(self, topic: str, msg_ids: Any, to_peer: str) -> None:
+    async def emit_ihave(self, topic: str, msg_ids: Any, to_peer: ID) -> None:
         """
         Emit ihave message, sent to to_peer, for topic and msg_ids
         """
@@ -587,7 +557,7 @@ class GossipSub(IPubsubRouter):
 
         await self.emit_control_message(control_msg, to_peer)
 
-    async def emit_iwant(self, msg_ids: Any, to_peer: str) -> None:
+    async def emit_iwant(self, msg_ids: Any, to_peer: ID) -> None:
         """
         Emit iwant message, sent to to_peer, for msg_ids
         """
@@ -600,7 +570,7 @@ class GossipSub(IPubsubRouter):
 
         await self.emit_control_message(control_msg, to_peer)
 
-    async def emit_graft(self, topic: str, to_peer: str) -> None:
+    async def emit_graft(self, topic: str, to_peer: ID) -> None:
         """
         Emit graft message, sent to to_peer, for topic
         """
@@ -613,7 +583,7 @@ class GossipSub(IPubsubRouter):
 
         await self.emit_control_message(control_msg, to_peer)
 
-    async def emit_prune(self, topic: str, to_peer: str) -> None:
+    async def emit_prune(self, topic: str, to_peer: ID) -> None:
         """
         Emit graft message, sent to to_peer, for topic
         """
@@ -627,7 +597,7 @@ class GossipSub(IPubsubRouter):
         await self.emit_control_message(control_msg, to_peer)
 
     async def emit_control_message(
-        self, control_msg: rpc_pb2.ControlMessage, to_peer: str
+        self, control_msg: rpc_pb2.ControlMessage, to_peer: ID
     ) -> None:
         # Add control message to packet
         packet: rpc_pb2.RPC = rpc_pb2.RPC()

libp2p/pubsub/pubsub.py

@@ -41,10 +41,8 @@ class Pubsub:
 
     my_topics: Dict[str, "asyncio.Queue[rpc_pb2.Message]"]
 
-    # FIXME: Should be changed to `Dict[str, List[ID]]`
+    peer_topics: Dict[str, List[ID]]
-    peer_topics: Dict[str, List[str]]
+    peers: Dict[ID, INetStream]
-    # FIXME: Should be changed to `Dict[ID, INetStream]`
-    peers: Dict[str, INetStream]
 
     # NOTE: Be sure it is increased atomically everytime.
     counter: int  # uint64
@@ -93,11 +91,9 @@ class Pubsub:
         self.my_topics = {}
 
         # Map of topic to peers to keep track of what peers are subscribed to
-        # FIXME: Should be changed to `Dict[str, ID]`
         self.peer_topics = {}
 
         # Create peers map, which maps peer_id (as string) to stream (to a given peer)
-        # FIXME: Should be changed to `Dict[ID, INetStream]`
         self.peers = {}
 
         self.counter = time.time_ns()
@@ -168,7 +164,7 @@ class Pubsub:
         # Add peer
         # Map peer to stream
         peer_id: ID = stream.mplex_conn.peer_id
-        self.peers[str(peer_id)] = stream
+        self.peers[peer_id] = stream
         self.router.add_peer(peer_id, stream.get_protocol())
 
         # Send hello packet
@@ -198,7 +194,7 @@ class Pubsub:
 
             # Add Peer
             # Map peer to stream
-            self.peers[str(peer_id)] = stream
+            self.peers[peer_id] = stream
             self.router.add_peer(peer_id, stream.get_protocol())
 
             # Send hello packet
@@ -223,17 +219,16 @@ class Pubsub:
         :param origin_id: id of the peer who subscribe to the message
         :param sub_message: RPC.SubOpts
         """
-        origin_id_str = str(origin_id)
         if sub_message.subscribe:
             if sub_message.topicid not in self.peer_topics:
-                self.peer_topics[sub_message.topicid] = [origin_id_str]
+                self.peer_topics[sub_message.topicid] = [origin_id]
-            elif origin_id_str not in self.peer_topics[sub_message.topicid]:
+            elif origin_id not in self.peer_topics[sub_message.topicid]:
                 # Add peer to topic
-                self.peer_topics[sub_message.topicid].append(origin_id_str)
+                self.peer_topics[sub_message.topicid].append(origin_id)
         else:
             if sub_message.topicid in self.peer_topics:
-                if origin_id_str in self.peer_topics[sub_message.topicid]:
+                if origin_id in self.peer_topics[sub_message.topicid]:
-                    self.peer_topics[sub_message.topicid].remove(origin_id_str)
+                    self.peer_topics[sub_message.topicid].remove(origin_id)
 
     # FIXME(mhchia): Change the function name?
     # FIXME(mhchia): `publish_message` can be further type hinted with mypy_protobuf

libp2p/routing/kademlia/kademlia_peer_router.py

@@ -22,7 +22,7 @@ class KadmeliaPeerRouter(IPeerRouting):
         :return: KadPeerInfo of specified peer
         """
         # switching peer_id to xor_id used by kademlia as node_id
-        xor_id = peer_id.get_xor_id()
+        xor_id = peer_id.xor_id
         value = await self.server.get(xor_id)
         return decode_peerinfo(value)
 

libp2p/transport/tcp/tcp.py

@@ -3,7 +3,7 @@ import asyncio
 import multiaddr
 
 from libp2p.network.connection.raw_connection import RawConnection
-from libp2p.peer.id import id_b58_encode
+from libp2p.peer.id import ID
 
 from ..listener_interface import IListener
 from ..transport_interface import ITransport
@@ -74,7 +74,7 @@ class TCP(ITransport):
         reader, writer = await asyncio.open_connection(host, port)
 
         # First: send our peer ID so receiver knows it
-        writer.write(id_b58_encode(self_id).encode())
+        writer.write(self_id.to_base58().encode())
         await writer.drain()
 
         # Await ack for peer id

tests/kademlia/test_providers.py

@@ -28,4 +28,4 @@ async def test_example():
     first_providers = first_tuple[1]
     # [b'\xf9\xa1\xf5\x10a\xe5\xe0F']
     first_provider = first_providers[0]
-    assert node_b.node.peer_id == first_provider
+    assert node_b.node.peer_id_bytes == first_provider

tests/peer/test_peerid.py

@@ -3,25 +3,19 @@ import multihash
 import pytest
 import base58
 from Crypto.PublicKey import RSA
-from libp2p.peer.id import (
+from libp2p.peer.id import ID
-    ID,
-    id_b58_encode,
-    id_b58_decode,
-    id_from_public_key,
-    id_from_private_key,
-)
 
 
 ALPHABETS = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
 
 
-def test_init_():
+def test_init():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
-    peer_id = ID(random_id_string)
+    peer_id = ID(random_id_string.encode())
     # pylint: disable=protected-access
-    assert peer_id._id_str == random_id_string
+    assert peer_id == random_id_string.encode()
 
 
 def test_no_init_value():
@@ -34,7 +28,7 @@ def test_pretty():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
-    peer_id = ID(random_id_string)
+    peer_id = ID(random_id_string.encode())
     actual = peer_id.pretty()
     expected = base58.b58encode(random_id_string).decode()
 
@@ -45,9 +39,9 @@ def test_str_less_than_10():
     random_id_string = ""
     for _ in range(5):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
-    pid = base58.b58encode(random_id_string).decode()
+    peer_id = base58.b58encode(random_id_string).decode()
-    expected = pid
+    expected = peer_id
-    actual = ID(random_id_string).__str__()
+    actual = ID(random_id_string.encode()).__str__()
 
     assert actual == expected
 
@@ -56,9 +50,9 @@ def test_str_more_than_10():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
-    pid = base58.b58encode(random_id_string).decode()
+    peer_id = base58.b58encode(random_id_string).decode()
-    expected = pid
+    expected = peer_id
-    actual = ID(random_id_string).__str__()
+    actual = ID(random_id_string.encode()).__str__()
 
     assert actual == expected
 
@@ -67,21 +61,18 @@ def test_eq_true():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
-    other = ID(random_id_string)
+    peer_id = ID(random_id_string.encode())
 
-    expected = True
+    assert peer_id == base58.b58encode(random_id_string).decode()
-    actual = ID(random_id_string).__eq__(other)
+    assert peer_id == random_id_string.encode()
-
+    assert peer_id == ID(random_id_string.encode())
-    assert actual == expected
 
 
 def test_eq_false():
-    other = ID("efgh")
+    peer_id = ID("efgh")
+    other = ID("abcd")
 
-    expected = False
+    assert peer_id != other
-    actual = ID("abcd").__eq__(other)
-
-    assert actual == expected
 
 
 def test_hash():
@@ -89,28 +80,28 @@ def test_hash():
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
 
-    expected = hash(random_id_string)
+    expected = hash(random_id_string.encode())
-    actual = ID(random_id_string).__hash__()
+    actual = ID(random_id_string.encode()).__hash__()
 
     assert actual == expected
 
 
-def test_id_b58_encode():
+def test_id_to_base58():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
     expected = base58.b58encode(random_id_string).decode()
-    actual = id_b58_encode(ID(random_id_string))
+    actual = ID(random_id_string.encode()).to_base58()
 
     assert actual == expected
 
 
-def test_id_b58_decode():
+def test_id_from_base58():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
     expected = ID(base58.b58decode(random_id_string))
-    actual = id_b58_decode(random_id_string)
+    actual = ID.from_base58(random_id_string.encode())
 
     assert actual == expected
 
@@ -122,13 +113,13 @@ def test_id_from_public_key():
     algo = multihash.Func.sha2_256
     mh_digest = multihash.digest(key_bin, algo)
     expected = ID(mh_digest.encode())
-    actual = id_from_public_key(key)
+    actual = ID.from_pubkey(key)
 
     assert actual == expected
 
 
 def test_id_from_private_key():
     key = RSA.generate(2048, e=65537)
-    id_from_pub = id_from_public_key(key.publickey())
+    id_from_pub = ID.from_pubkey(key.publickey())
-    id_from_priv = id_from_private_key(key)
+    id_from_priv = ID.from_privkey(key)
     assert id_from_pub == id_from_priv

tests/peer/test_peerinfo.py

@@ -19,7 +19,7 @@ def test_init_():
     random_id_string = ""
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
-    peer_id = ID(random_id_string)
+    peer_id = ID(random_id_string.encode())
     peer_info = PeerInfo(peer_id, peer_data)
 
     assert peer_info.peer_id == peer_id

tests/pubsub/floodsub_integration_test_settings.py

@@ -179,12 +179,12 @@ async def perform_test_from_obj(obj, router_factory):
     node_map = {}
     pubsub_map = {}
 
-    async def add_node(node_id: str) -> None:
+    async def add_node(node_id_str: str) -> None:
         pubsub_router = router_factory(protocols=obj["supported_protocols"])
         pubsub = PubsubFactory(router=pubsub_router)
         await pubsub.host.get_network().listen(LISTEN_MADDR)
-        node_map[node_id] = pubsub.host
+        node_map[node_id_str] = pubsub.host
-        pubsub_map[node_id] = pubsub
+        pubsub_map[node_id_str] = pubsub
 
     tasks_connect = []
     for start_node_id in adj_list:

tests/pubsub/test_gossipsub.py

@@ -54,11 +54,11 @@ async def test_join(num_hosts, hosts, gossipsubs, pubsubs_gsub):
 
     for i in hosts_indices:
         if i in subscribed_peer_indices:
-            assert str(hosts[i].get_id()) in gossipsubs[central_node_index].mesh[topic]
+            assert hosts[i].get_id() in gossipsubs[central_node_index].mesh[topic]
-            assert str(hosts[central_node_index].get_id()) in gossipsubs[i].mesh[topic]
+            assert hosts[central_node_index].get_id() in gossipsubs[i].mesh[topic]
         else:
             assert (
-                str(hosts[i].get_id()) not in gossipsubs[central_node_index].mesh[topic]
+                hosts[i].get_id() not in gossipsubs[central_node_index].mesh[topic]
             )
             assert topic not in gossipsubs[i].mesh
 
@@ -89,9 +89,9 @@ async def test_leave(pubsubs_gsub):
 @pytest.mark.asyncio
 async def test_handle_graft(pubsubs_gsub, hosts, gossipsubs, event_loop, monkeypatch):
     index_alice = 0
-    id_alice = str(hosts[index_alice].get_id())
+    id_alice = hosts[index_alice].get_id()
     index_bob = 1
-    id_bob = str(hosts[index_bob].get_id())
+    id_bob = hosts[index_bob].get_id()
     await connect(hosts[index_alice], hosts[index_bob])
 
     # Wait 2 seconds for heartbeat to allow mesh to connect
@@ -141,9 +141,9 @@ async def test_handle_graft(pubsubs_gsub, hosts, gossipsubs, event_loop, monkeyp
 @pytest.mark.asyncio
 async def test_handle_prune(pubsubs_gsub, hosts, gossipsubs):
     index_alice = 0
-    id_alice = str(hosts[index_alice].get_id())
+    id_alice = hosts[index_alice].get_id()
     index_bob = 1
-    id_bob = str(hosts[index_bob].get_id())
+    id_bob = hosts[index_bob].get_id()
 
     topic = "test_handle_prune"
     for pubsub in pubsubs_gsub:

tests/pubsub/test_pubsub.py

@@ -60,11 +60,11 @@ async def test_peers_subscribe(pubsubs_fsub):
     await pubsubs_fsub[0].subscribe(TESTING_TOPIC)
     # Yield to let 0 notify 1
     await asyncio.sleep(0.1)
-    assert str(pubsubs_fsub[0].my_id) in pubsubs_fsub[1].peer_topics[TESTING_TOPIC]
+    assert pubsubs_fsub[0].my_id in pubsubs_fsub[1].peer_topics[TESTING_TOPIC]
     await pubsubs_fsub[0].unsubscribe(TESTING_TOPIC)
     # Yield to let 0 notify 1
     await asyncio.sleep(0.1)
-    assert str(pubsubs_fsub[0].my_id) not in pubsubs_fsub[1].peer_topics[TESTING_TOPIC]
+    assert pubsubs_fsub[0].my_id not in pubsubs_fsub[1].peer_topics[TESTING_TOPIC]
 
 
 @pytest.mark.parametrize("num_hosts", (1,))
@@ -212,23 +212,23 @@ def test_handle_subscription(pubsubs_fsub):
         and TESTING_TOPIC in pubsubs_fsub[0].peer_topics
     )
     assert len(pubsubs_fsub[0].peer_topics[TESTING_TOPIC]) == 1
-    assert str(peer_ids[0]) in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
+    assert peer_ids[0] in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
     # Test: Another peer is subscribed
     pubsubs_fsub[0].handle_subscription(peer_ids[1], sub_msg_0)
     assert len(pubsubs_fsub[0].peer_topics) == 1
     assert len(pubsubs_fsub[0].peer_topics[TESTING_TOPIC]) == 2
-    assert str(peer_ids[1]) in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
+    assert peer_ids[1] in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
     # Test: Subscribe to another topic
     another_topic = "ANOTHER_TOPIC"
     sub_msg_1 = rpc_pb2.RPC.SubOpts(subscribe=True, topicid=another_topic)
     pubsubs_fsub[0].handle_subscription(peer_ids[0], sub_msg_1)
     assert len(pubsubs_fsub[0].peer_topics) == 2
     assert another_topic in pubsubs_fsub[0].peer_topics
-    assert str(peer_ids[0]) in pubsubs_fsub[0].peer_topics[another_topic]
+    assert peer_ids[0] in pubsubs_fsub[0].peer_topics[another_topic]
     # Test: unsubscribe
     unsub_msg = rpc_pb2.RPC.SubOpts(subscribe=False, topicid=TESTING_TOPIC)
     pubsubs_fsub[0].handle_subscription(peer_ids[0], unsub_msg)
-    assert str(peer_ids[0]) not in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
+    assert peer_ids[0] not in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
 
 
 @pytest.mark.parametrize("num_hosts", (1,))
@@ -261,7 +261,7 @@ async def test_handle_talk(pubsubs_fsub):
 @pytest.mark.asyncio
 async def test_message_all_peers(pubsubs_fsub, monkeypatch):
     peer_ids = [ID(b"\x12\x20" + i.to_bytes(32, "big")) for i in range(10)]
-    mock_peers = {str(peer_id): FakeNetStream() for peer_id in peer_ids}
+    mock_peers = {peer_id: FakeNetStream() for peer_id in peer_ids}
     monkeypatch.setattr(pubsubs_fsub[0], "peers", mock_peers)
 
     empty_rpc = rpc_pb2.RPC()

tests/routing/test_kad_peer_router.py

@@ -1,5 +1,6 @@
 import pytest
 
+from libp2p.peer.id import ID
 from libp2p.kademlia.network import KademliaServer
 from libp2p.routing.kademlia.kademlia_peer_router import KadmeliaPeerRouter
 
@@ -17,7 +18,7 @@ async def test_simple_two_nodes():
     await node_a.set(node_a_kad_peerinfo.xor_id, repr(node_a_kad_peerinfo))
 
     router = KadmeliaPeerRouter(node_b)
-    returned_info = await router.find_peer(node_a_kad_peerinfo.peer_id_obj)
+    returned_info = await router.find_peer(ID(node_a_kad_peerinfo.peer_id_bytes))
     print(repr(returned_info))
     print(repr(node_a_kad_peerinfo))
     assert repr(returned_info) == repr(node_a_kad_peerinfo)
@@ -41,7 +42,7 @@ async def test_simple_three_nodes():
     await node_a.set(node_a_kad_peerinfo.xor_id, repr(node_a_kad_peerinfo))
 
     router = KadmeliaPeerRouter(node_c)
-    returned_info = await router.find_peer(node_a_kad_peerinfo.peer_id_obj)
+    returned_info = await router.find_peer(ID(node_a_kad_peerinfo.peer_id_bytes))
     assert str(returned_info) == str(node_a_kad_peerinfo)
 
 
@@ -69,5 +70,5 @@ async def test_simple_four_nodes():
     await node_b.set(node_a_kad_peerinfo.xor_id, repr(node_a_kad_peerinfo))
 
     router = KadmeliaPeerRouter(node_d)
-    returned_info = await router.find_peer(node_a_kad_peerinfo.peer_id_obj)
+    returned_info = await router.find_peer(ID(node_a_kad_peerinfo.peer_id_bytes))
     assert str(returned_info) == str(node_a_kad_peerinfo)