async def get_peer_and_receive_server(
request, event_loop, event_bus
) -> Tuple[BCCPeer, BCCRequestServer, BCCReceiveServer, asyncio.Queue]:
alice_chain = await get_fake_chain()
bob_chain = await get_fake_chain()
(alice, alice_peer_pool, bob, bob_peer_pool
) = await bcc_helpers.get_directly_linked_peers_in_peer_pools(
request,
event_loop,
alice_chain_db=alice_chain.chaindb,
bob_chain_db=bob_chain.chaindb,
)
msg_queue = asyncio.Queue()
orig_handle_msg = BCCReceiveServer._handle_msg
# Inject a queue to each `BCCReceiveServer`, which puts the message passed to `_handle_msg` to
# the queue, right after every `_handle_msg` finishes.
# This is crucial to make the test be able to wait until `_handle_msg` finishes.
async def _handle_msg(self, base_peer, cmd, msg):
task = asyncio.ensure_future(orig_handle_msg(self, base_peer, cmd,
msg))
def enqueue_msg(future, msg):
msg_queue.put_nowait(msg)
task.add_done_callback(functools.partial(enqueue_msg, msg=msg))
await task
BCCReceiveServer._handle_msg = _handle_msg
async with run_peer_pool_event_server(
event_bus, alice_peer_pool,
BCCPeerPoolEventServer), run_request_server(
event_bus, alice_chain.chaindb,
server_type=BCCRequestServer) as alice_req_server:
bob_recv_server = BCCReceiveServer(chain=bob_chain,
peer_pool=bob_peer_pool)
asyncio.ensure_future(bob_recv_server.run())
await bob_recv_server.events.started.wait()
def finalizer():
event_loop.run_until_complete(bob_recv_server.cancel())
request.addfinalizer(finalizer)
yield alice, alice_req_server, bob_recv_server, msg_queue
def _make_receive_server(self) -> BCCReceiveServer:
return BCCReceiveServer(
chain=cast(BeaconChain, self.chain),
peer_pool=self.peer_pool,
token=self.cancel_token,
)
async def get_peer_and_receive_server(
request, event_loop, event_bus
) -> Tuple[BCCPeer, BCCRequestServer, BCCReceiveServer, asyncio.Queue]:
alice_chain = await get_fake_chain()
bob_chain = await get_fake_chain()
alice_private_key = PrivateKeyFactory()
bob_private_key = PrivateKeyFactory()
alice_context = BeaconContextFactory(chain_db=alice_chain.chaindb)
bob_context = BeaconContextFactory(chain_db=bob_chain.chaindb)
peer_pair = BCCPeerPairFactory(
alice_peer_context=alice_context,
alice_private_key=alice_private_key,
bob_peer_context=bob_context,
bob_private_key=bob_private_key,
event_bus=event_bus,
)
async with peer_pair as (alice, bob):
alice_pool_ctx = BCCPeerPoolFactory.run_for_peer(
alice,
privkey=alice_private_key,
context=alice_context,
event_bus=event_bus,
)
bob_pool_ctx = BCCPeerPoolFactory.run_for_peer(
bob,
privkey=bob_private_key,
context=bob_context,
event_bus=event_bus,
)
async with alice_pool_ctx as alice_peer_pool, bob_pool_ctx as bob_peer_pool:
msg_queue = asyncio.Queue()
orig_handle_msg = BCCReceiveServer._handle_msg
# Inject a queue to each `BCCReceiveServer`, which puts the message
# passed to `_handle_msg` to the queue, right after every `_handle_msg`
# finishes. This is crucial to make the test be able to wait until
# `_handle_msg` finishes.
async def _handle_msg(self, base_peer, cmd, msg):
task = asyncio.ensure_future(
orig_handle_msg(self, base_peer, cmd, msg))
def enqueue_msg(future, msg):
msg_queue.put_nowait(msg)
task.add_done_callback(functools.partial(enqueue_msg, msg=msg))
await task
BCCReceiveServer._handle_msg = _handle_msg
async with run_peer_pool_event_server(
event_bus, alice_peer_pool,
BCCPeerPoolEventServer), run_request_server(
event_bus,
alice_chain.chaindb,
server_type=BCCRequestServer) as alice_req_server:
bob_recv_server = BCCReceiveServer(chain=bob_chain,
peer_pool=bob_peer_pool)
asyncio.ensure_future(bob_recv_server.run())
await bob_recv_server.events.started.wait()
def finalizer():
event_loop.run_until_complete(bob_recv_server.cancel())
request.addfinalizer(finalizer)
yield alice, alice_req_server, bob_recv_server, msg_queue