def test_topic_table(): table = discovery.TopicTable(logging.getLogger("test")) topic = b'topic' node = random_node() table.add_node(node, topic) assert len(table.get_nodes(topic)) == 1 assert table.get_nodes(topic)[0] == node node2 = random_node() table.add_node(node2, topic) assert len(table.get_nodes(topic)) == 2 assert table.get_nodes(topic)[1] == node2 # Adding the same node again won't cause a duplicate entry to be added. table.add_node(node, topic) assert len(table.get_nodes(topic)) == 2 # When we reach the max number of entries for a given topic, the first added items are evicted # to make room for the new ones. for _ in range(discovery.MAX_ENTRIES_PER_TOPIC + 2): table.add_node(random_node(), topic) assert node not in table.get_nodes(topic) assert node2 not in table.get_nodes(topic)
async def test_wait_neighbours(): proto = MockDiscoveryProtocol([]) node = random_node() # Schedule a call to proto.recv_neighbours_v4() simulating a neighbours response from the node # we expect. neighbours = (random_node(), random_node(), random_node()) neighbours_msg_payload = [ [n.address.to_endpoint() + [n.pubkey.to_bytes()] for n in neighbours], discovery._get_msg_expiration()] recv_neighbours_coroutine = asyncio.coroutine( lambda: proto.recv_neighbours_v4(node, neighbours_msg_payload, b'')) asyncio.ensure_future(recv_neighbours_coroutine()) received_neighbours = await proto.wait_neighbours(node) assert neighbours == received_neighbours # Ensure wait_neighbours() cleaned up after itself. assert node not in proto.neighbours_callbacks # If wait_neighbours() times out, we get an empty list of neighbours. received_neighbours = await proto.wait_neighbours(node) assert received_neighbours == tuple() assert node not in proto.neighbours_callbacks
async def test_wait_ping(echo): proto = MockDiscoveryProtocol([]) node = random_node() # Schedule a call to proto.recv_ping() simulating a ping from the node we expect. recv_ping_coroutine = asyncio.coroutine( lambda: proto.recv_ping_v4(node, echo, b'')) asyncio.ensure_future(recv_ping_coroutine()) got_ping = await proto.wait_ping(node) assert got_ping # Ensure wait_ping() cleaned up after itself. assert node not in proto.ping_callbacks # If we waited for a ping from a different node, wait_ping() would timeout and thus return # false. recv_ping_coroutine = asyncio.coroutine( lambda: proto.recv_ping_v4(node, echo, b'')) asyncio.ensure_future(recv_ping_coroutine()) node2 = random_node() got_ping = await proto.wait_ping(node2) assert not got_ping assert node2 not in proto.ping_callbacks
def _test_find_node_neighbours(use_v5): alice = get_discovery_protocol(b"alice") bob = get_discovery_protocol(b"bob") # Add some nodes to bob's routing table so that it has something to use when replying to # alice's find_node. for _ in range(kademlia.k_bucket_size * 2): bob.update_routing_table(random_node()) # Connect alice's and bob's transports directly so we don't need to deal with the complexities # of going over the wire. link_transports(alice, bob) # Collect all neighbours packets received by alice in a list for later inspection. received_neighbours = [] alice.recv_neighbours_v4 = lambda node, payload, hash_: received_neighbours.append((node, payload)) # noqa: E501 # Pretend that bob and alice have already bonded, otherwise bob will ignore alice's find_node. bob.update_routing_table(alice.this_node) if use_v5: alice.send_find_node_v5(bob.this_node, alice.this_node.id) else: alice.send_find_node_v4(bob.this_node, alice.this_node.id) # Bob should have sent two neighbours packets in order to keep the total packet size under the # 1280 bytes limit. assert len(received_neighbours) == 2 packet1, packet2 = received_neighbours neighbours = [] for packet in [packet1, packet2]: node, payload = packet assert node == bob.this_node neighbours.extend(discovery._extract_nodes_from_payload( node.address, payload[0], bob.logger)) assert len(neighbours) == kademlia.k_bucket_size
def test_update_routing_table(): proto = MockDiscoveryProtocol([]) node = random_node() assert proto.update_routing_table(node) is None assert node in proto.routing
async def test_wait_pong(): proto = MockDiscoveryProtocol([]) us = proto.this_node node = random_node() token = b'token' # Schedule a call to proto.recv_pong() simulating a pong from the node we expect. pong_msg_payload = [us.address.to_endpoint(), token, discovery._get_msg_expiration()] recv_pong_coroutine = asyncio.coroutine(lambda: proto.recv_pong_v4(node, pong_msg_payload, b'')) asyncio.ensure_future(recv_pong_coroutine()) got_pong = await proto.wait_pong_v4(node, token) assert got_pong # Ensure wait_pong() cleaned up after itself. pingid = proto._mkpingid(token, node) assert pingid not in proto.pong_callbacks # If the remote node echoed something different than what we expected, wait_pong() would # timeout. wrong_token = b"foo" pong_msg_payload = [us.address.to_endpoint(), wrong_token, discovery._get_msg_expiration()] recv_pong_coroutine = asyncio.coroutine(lambda: proto.recv_pong_v4(node, pong_msg_payload, b'')) asyncio.ensure_future(recv_pong_coroutine()) got_pong = await proto.wait_pong_v4(node, token) assert not got_pong assert pingid not in proto.pong_callbacks
async def test_protocol_bootstrap(): node1, node2 = [random_node(), random_node()] proto = MockDiscoveryProtocol([node1, node2]) async def bond(node): assert proto.routing.add_node(node) is None return True # Pretend we bonded successfully with our bootstrap nodes. proto.bond = bond await proto.bootstrap() assert len(proto.messages) == 2 # We don't care in which order the bootstrap nodes are contacted, nor which node_id was used # in the find_node request, so we just assert that we sent find_node msgs to both nodes. assert sorted([(node, cmd) for (node, cmd, _) in proto.messages ]) == sorted([(node1, 'find_node'), (node2, 'find_node')])
async def test_bond(): proto = MockDiscoveryProtocol([]) node = random_node() token = b'token' # Do not send pings, instead simply return the pingid we'd expect back together with the pong. proto.send_ping_v4 = lambda remote: token # Pretend we get a pong from the node we are bonding with. proto.wait_pong_v4 = asyncio.coroutine(lambda n, t: t == token and n == node) bonded = await proto.bond(node) assert bonded # If we try to bond with any other nodes we'll timeout and bond() will return False. node2 = random_node() bonded = await proto.bond(node2) assert not bonded
async def test_topic_query(event_loop): bob = await get_listening_discovery_protocol(event_loop) les_nodes = [random_node() for _ in range(10)] topic = b'les' for n in les_nodes: bob.topic_table.add_node(n, topic) alice = await get_listening_discovery_protocol(event_loop) echo = alice.send_topic_query(bob.this_node, topic) received_nodes = await alice.wait_topic_nodes(bob.this_node, echo) assert len(received_nodes) == 10 assert sorted(received_nodes) == sorted(les_nodes)
async def test_update_routing_table_triggers_bond_if_eviction_candidate(): proto = MockDiscoveryProtocol([]) old_node, new_node = random_node(), random_node() bond_called = False def bond(node): nonlocal bond_called bond_called = True assert node == old_node proto.bond = asyncio.coroutine(bond) # Pretend our routing table failed to add the new node by returning the least recently seen # node for an eviction check. proto.routing.add_node = lambda n: old_node proto.update_routing_table(new_node) assert new_node not in proto.routing # The update_routing_table() call above will have scheduled a future call to proto.bond() so # we need to yield here to give it a chance to run. await asyncio.sleep(0.001) assert bond_called