Proto-Book: added tests

libp2p/peer/peerdata.py

@@ -46,6 +46,8 @@ class PeerData(IPeerData):
         self.ttl = 0
         self.latmap = 0
 
+    # --------PROTO-BOOK--------
+
     def get_protocols(self) -> list[str]:
         """
         :return: all protocols associated with given peer
@@ -94,6 +96,7 @@ class PeerData(IPeerData):
         """Clear all protocols"""
         self.protocols = []
 
+    # -------ADDR-BOOK---------
     def add_addrs(self, addrs: Sequence[Multiaddr]) -> None:
         """
         :param addrs: multiaddresses to add
@@ -112,6 +115,9 @@ class PeerData(IPeerData):
         """Clear all addresses."""
         self.addrs = []
 
+    # TODO! ADDRS_STREAM
+
+    # -------METADATA-----------
     def put_metadata(self, key: str, val: Any) -> None:
         """
         :param key: key in KV pair
@@ -133,6 +139,7 @@ class PeerData(IPeerData):
         """Clears metadata."""
         self.metadata = {}
 
+    # -------KEY-BOOK---------------
     def add_pubkey(self, pubkey: PublicKey) -> None:
         """
         :param pubkey:
@@ -168,6 +175,7 @@ class PeerData(IPeerData):
         self.pubkey = None
         self.privkey = None
 
+    # ----------METRICS--------------
     def record_latency(self, new_latency: float) -> None:
         """
         Records a new latency measurement for the given peer
@@ -196,6 +204,7 @@ class PeerData(IPeerData):
     def update_last_identified(self) -> None:
         self.last_identified = int(time.time())
 
+    # ----------TTL------------------
     def get_last_identified(self) -> int:
         """
         :return: last identified timestamp

libp2p/peer/peerstore.py

@@ -62,6 +62,20 @@ class PeerStore(IPeerStore):
     def clear_peerdata(self, peer_id: ID) -> None:
         """Clears the peer data of the peer"""
 
+    def valid_peer_ids(self) -> list[ID]:
+        """
+        :return: all of the valid peer IDs stored in peer store
+        """
+        valid_peer_ids: list[ID] = []
+        for peer_id, peer_data in self.peer_data_map.items():
+            if not peer_data.is_expired():
+                valid_peer_ids.append(peer_id)
+            else:
+                peer_data.clear_addrs()
+        return valid_peer_ids
+
+    # --------PROTO-BOOK--------
+
     def get_protocols(self, peer_id: ID) -> list[str]:
         """
         :param peer_id: peer ID to get protocols for
@@ -112,17 +126,7 @@ class PeerStore(IPeerStore):
         peer_data = self.peer_data_map[peer_id]
         peer_data.clear_protocol_data()
 
-    def valid_peer_ids(self) -> list[ID]:
+    # ------METADATA---------
-        """
-        :return: all of the valid peer IDs stored in peer store
-        """
-        valid_peer_ids: list[ID] = []
-        for peer_id, peer_data in self.peer_data_map.items():
-            if not peer_data.is_expired():
-                valid_peer_ids.append(peer_id)
-            else:
-                peer_data.clear_addrs()
-        return valid_peer_ids
 
     def get(self, peer_id: ID, key: str) -> Any:
         """
@@ -153,6 +157,8 @@ class PeerStore(IPeerStore):
         peer_data = self.peer_data_map[peer_id]
         peer_data.clear_metadata()
 
+    # -------ADDR-BOOK--------
+
     def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int = 0) -> None:
         """
         :param peer_id: peer ID to add address for
@@ -215,6 +221,8 @@ class PeerStore(IPeerStore):
         """addr_stream"""
         # TODO!
 
+    # -------KEY-BOOK---------
+
     def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None:
         """
         :param peer_id: peer ID to add public key for
@@ -288,6 +296,8 @@ class PeerStore(IPeerStore):
         peer_data = self.peer_data_map[peer_id]
         peer_data.clear_keydata()
 
+    # --------METRICS--------
+
     def record_latency(self, peer_id: ID, RTT: float) -> None:
         """
         Records a new latency measurement for the given peer

tests/core/peer/test_peerdata.py

@@ -39,6 +39,56 @@ 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 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"]
+
+
+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"
+
+
+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"
+
+
+def test_clear_protocol_data():
+    peer_data = PeerData()
+    peer_data.set_protocols(["proto1", "proto2"])
+
+    peer_data.clear_protocol_data()
+
+    assert peer_data.get_protocols() == []
+
+
 # Test case when adding addresses
 def test_add_addrs():
     peer_data = PeerData()