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_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