async def setup_node(self, peer_addresses, address, node_id): self.nodes: typing.Dict[int, Node] = {} self.advance = dht_mocks.get_time_accelerator(self.loop) self.instant_advance = dht_mocks.get_time_accelerator(self.loop) self.conf = Config() self.peer_manager = PeerManager(self.loop) self.node = Node(self.loop, self.peer_manager, node_id, 4444, 4444, 3333, address) await self.node.start_listening(address) await asyncio.gather(*[self.add_peer(node_id, address) for node_id, address in peer_addresses]) for first_peer in self.nodes.values(): for second_peer in self.nodes.values(): if first_peer == second_peer: continue self.add_peer_to_routing_table(first_peer, second_peer) self.add_peer_to_routing_table(second_peer, first_peer) await self.advance(0.1) # just to make pings go through self.node.joined.set() self.node._refresh_task = self.loop.create_task(self.node.refresh_node()) self.storage = SQLiteStorage(self.conf, ":memory:", self.loop, self.loop.time) await self.storage.open() self.blob_announcer = BlobAnnouncer(self.loop, self.node, self.storage)
async def setup_node(self, peer_addresses, address, node_id): self.nodes: typing.Dict[int, Node] = {} self.advance = dht_mocks.get_time_accelerator(self.loop, self.loop.time()) self.conf = Config() self.storage = SQLiteStorage(self.conf, ":memory:", self.loop, self.loop.time) await self.storage.open() self.peer_manager = PeerManager(self.loop) self.node = Node(self.loop, self.peer_manager, node_id, 4444, 4444, 3333, address) await self.node.start_listening(address) self.blob_announcer = BlobAnnouncer(self.loop, self.node, self.storage) for node_id, address in peer_addresses: await self.add_peer(node_id, address) self.node.joined.set()
async def test_ping_queue_discover(self): loop = asyncio.get_event_loop() loop.set_debug(False) peer_addresses = [ (constants.generate_id(1), '1.2.3.1'), (constants.generate_id(2), '1.2.3.2'), (constants.generate_id(3), '1.2.3.3'), (constants.generate_id(4), '1.2.3.4'), (constants.generate_id(5), '1.2.3.5'), (constants.generate_id(6), '1.2.3.6'), (constants.generate_id(7), '1.2.3.7'), (constants.generate_id(8), '1.2.3.8'), (constants.generate_id(9), '1.2.3.9'), ] with dht_mocks.mock_network_loop(loop): advance = dht_mocks.get_time_accelerator(loop) # start the nodes nodes: typing.Dict[int, Node] = { i: Node(loop, PeerManager(loop), node_id, 4444, 4444, 3333, address) for i, (node_id, address) in enumerate(peer_addresses) } for i, n in nodes.items(): n.start(peer_addresses[i][1], []) await advance(1) node_1 = nodes[0] # ping 8 nodes from node_1, this will result in a delayed return ping futs = [] for i in range(1, len(peer_addresses)): node = nodes[i] assert node.protocol.node_id != node_1.protocol.node_id peer = make_kademlia_peer(node.protocol.node_id, node.protocol.external_ip, udp_port=node.protocol.udp_port) futs.append(node_1.protocol.get_rpc_peer(peer).ping()) await advance(3) replies = await asyncio.gather(*tuple(futs)) self.assertTrue(all(map(lambda reply: reply == b"pong", replies))) # run for long enough for the delayed pings to have been sent by node 1 await advance(1000) # verify all of the previously pinged peers have node_1 in their routing tables for n in nodes.values(): peers = n.protocol.routing_table.get_peers() if n is node_1: self.assertEqual(8, len(peers)) # TODO: figure out why this breaks # else: # self.assertEqual(1, len(peers)) # self.assertEqual((peers[0].node_id, peers[0].address, peers[0].udp_port), # (node_1.protocol.node_id, node_1.protocol.external_ip, node_1.protocol.udp_port)) # run long enough for the refresh loop to run await advance(3600) # verify all the nodes know about each other for n in nodes.values(): if n is node_1: continue peers = n.protocol.routing_table.get_peers() self.assertEqual(8, len(peers)) self.assertSetEqual( { n_id[0] for n_id in peer_addresses if n_id[0] != n.protocol.node_id }, {c.node_id for c in peers}) self.assertSetEqual( { n_addr[1] for n_addr in peer_addresses if n_addr[1] != n.protocol.external_ip }, {c.address for c in peers}) # teardown for n in nodes.values(): n.stop()
async def test_losing_connection(self): async def wait_for(check_ok, insist, timeout=20): start = time.time() while time.time() - start < timeout: if check_ok(): break await asyncio.sleep(0) else: insist() loop = self.loop loop.set_debug(False) peer_addresses = [('1.2.3.4', 40000 + i) for i in range(10)] node_ids = [constants.generate_id(i) for i in range(10)] nodes = [ Node(loop, PeerManager(loop), node_id, udp_port, udp_port, 3333, address, storage=SQLiteStorage(Config(), ":memory:", self.loop, self.loop.time)) for node_id, (address, udp_port) in zip(node_ids, peer_addresses) ] dht_network = { peer_addresses[i]: node.protocol for i, node in enumerate(nodes) } num_seeds = 3 with dht_mocks.mock_network_loop(loop, dht_network): for i, n in enumerate(nodes): await n._storage.open() self.addCleanup(n.stop) n.start(peer_addresses[i][0], peer_addresses[:num_seeds]) await asyncio.gather(*[n.joined.wait() for n in nodes]) node = nodes[-1] advance = dht_mocks.get_time_accelerator(loop) await advance(500) # Join the network, assert that at least the known peers are in RT self.assertTrue(node.joined.is_set()) self.assertTrue( len(node.protocol.routing_table.get_peers()) >= num_seeds) # Refresh, so that the peers are persisted self.assertFalse( len(await node._storage.get_persisted_kademlia_peers()) > num_seeds) await advance(4000) self.assertTrue( len(await node._storage.get_persisted_kademlia_peers()) > num_seeds) # We lost internet connection - all the peers stop responding dht_network.pop( (node.protocol.external_ip, node.protocol.udp_port)) # The peers are cleared on refresh from RT and storage await advance(4000) self.assertListEqual([], await node._storage.get_persisted_kademlia_peers()) await wait_for( lambda: len(node.protocol.routing_table.get_peers()) == 0, lambda: self.assertListEqual( node.protocol.routing_table.get_peers(), [])) # Reconnect dht_network[(node.protocol.external_ip, node.protocol.udp_port)] = node.protocol # Check that node reconnects at least to them await advance(1000) await wait_for( lambda: len(node.protocol.routing_table.get_peers()) >= num_seeds, lambda: self.assertGreaterEqual( len(node.protocol.routing_table.get_peers()), num_seeds))