diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index bf54a494..725ca5a3 100644 --- a/libp2p/peer/peerdata.py +++ b/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 diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index 4539fc87..da18a5f0 100644 --- a/libp2p/peer/peerstore.py +++ b/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]: - """ - :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 + # ------METADATA--------- 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 diff --git a/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index 65e98959..500e19d4 100644 --- a/tests/core/peer/test_peerdata.py +++ b/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()