diff --git a/tests/unit/dht/protocol/test_protocol.py b/tests/unit/dht/protocol/test_protocol.py index 8a19dce74..d5c4f2da8 100644 --- a/tests/unit/dht/protocol/test_protocol.py +++ b/tests/unit/dht/protocol/test_protocol.py @@ -93,3 +93,65 @@ class TestProtocol(AsyncioTestCase): peer2.stop() peer1.disconnect() peer2.disconnect() + + async def _make_protocol(self, other_peer, node_id, address, udp_port, tcp_port): + proto = KademliaProtocol( + self.loop, PeerManager(self.loop), node_id, address, udp_port, tcp_port + ) + await self.loop.create_datagram_endpoint(lambda: proto, (address, 4444)) + return proto, other_peer.peer_manager.get_kademlia_peer(node_id, address, udp_port=udp_port) + + async def test_add_peer_after_handle_request(self): + with dht_mocks.mock_network_loop(self.loop): + node_id1 = constants.generate_id() + node_id2 = constants.generate_id() + node_id3 = constants.generate_id() + node_id4 = constants.generate_id() + + peer1 = KademliaProtocol( + self.loop, PeerManager(self.loop), node_id1, '1.2.3.4', 4444, 3333 + ) + await self.loop.create_datagram_endpoint(lambda: peer1, ('1.2.3.4', 4444)) + + peer2, peer_2_from_peer_1 = await self._make_protocol(peer1, node_id2, '1.2.3.5', 4444, 3333) + peer3, peer_3_from_peer_1 = await self._make_protocol(peer1, node_id3, '1.2.3.6', 4444, 3333) + peer4, peer_4_from_peer_1 = await self._make_protocol(peer1, node_id4, '1.2.3.7', 4444, 3333) + + # peers who reply should be added + await peer1.get_rpc_peer(peer_2_from_peer_1).ping() + self.assertListEqual([peer_2_from_peer_1], peer1.routing_table.get_peers()) + peer1.routing_table.remove_peer(peer_2_from_peer_1) + + # peers not known by be good/bad should be enqueued to maybe-ping + peer1_from_peer3 = peer3.get_rpc_peer(peer3.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', 4444)) + self.assertEqual(0, len(peer1.ping_queue._pending_contacts)) + pong = await peer1_from_peer3.ping() + self.assertEqual(b'pong', pong) + self.assertEqual(1, len(peer1.ping_queue._pending_contacts)) + peer1.ping_queue._pending_contacts.clear() + + # peers who are already good should be added + peer1_from_peer4 = peer4.get_rpc_peer(peer4.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', 4444)) + peer1.peer_manager.update_contact_triple(node_id4,'1.2.3.7', 4444) + peer1.peer_manager.report_last_replied('1.2.3.7', 4444) + self.assertEqual(0, len(peer1.ping_queue._pending_contacts)) + pong = await peer1_from_peer4.ping() + self.assertEqual(b'pong', pong) + self.assertEqual(1, len(peer1.routing_table.get_peers())) + self.assertEqual(0, len(peer1.ping_queue._pending_contacts)) + peer1.routing_table.buckets[0].peers.clear() + + # peers who are known to be bad recently should not be added or maybe-pinged + peer1_from_peer4 = peer4.get_rpc_peer(peer4.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', 4444)) + peer1.peer_manager.update_contact_triple(node_id4,'1.2.3.7', 4444) + peer1.peer_manager.report_failure('1.2.3.7', 4444) + peer1.peer_manager.report_failure('1.2.3.7', 4444) + self.assertEqual(0, len(peer1.ping_queue._pending_contacts)) + pong = await peer1_from_peer4.ping() + self.assertEqual(b'pong', pong) + self.assertEqual(0, len(peer1.routing_table.get_peers())) + self.assertEqual(0, len(peer1.ping_queue._pending_contacts)) + + for p in [peer1, peer2, peer3, peer4]: + p.stop() + p.disconnect()