async def test_broadcast(self, loop_communicator): BROADCAST = { 'body': 'present', 'sender': 'Martin', 'subject': 'sup', 'correlation_id': 420 } broadcast_future = plumpy.Future() loop = asyncio.get_event_loop() def get_broadcast(_comm, body, sender, subject, correlation_id): assert loop is asyncio.get_event_loop() broadcast_future.set_result({ 'body': body, 'sender': sender, 'subject': subject, 'correlation_id': correlation_id }) loop_communicator.add_broadcast_subscriber(get_broadcast) loop_communicator.broadcast_send(**BROADCAST) result = await broadcast_future assert result == BROADCAST
def test_call_on_process_finish(create_runner): """Test call on calculation finish.""" runner = create_runner() loop = runner.loop proc = Proc(runner=runner) future = plumpy.Future() event = threading.Event() def calc_done(): if event.is_set(): future.set_exc_info(AssertionError('the callback was called twice, which should never happen')) future.set_result(True) event.set() loop.stop() runner.call_on_process_finish(proc.node.pk, calc_done) # Run the calculation runner.loop.add_callback(proc.step_until_terminated) loop.call_later(5, the_hans_klok_comeback, runner.loop) loop.start() assert not future.exc_info() assert future.result()
def interrupt(self, reason): """Interrupt the `Waiting` state by calling interrupt on the transport task `InterruptableFuture`.""" if self._task is not None: self._task.interrupt(reason) if isinstance(reason, plumpy.KillInterruption): if self._killing is None: self._killing = plumpy.Future() return self._killing
async def test_rpc(self, loop_communicator): MSG = 'rpc this' rpc_future = plumpy.Future() loop = asyncio.get_event_loop() def get_rpc(_comm, msg): assert loop is asyncio.get_event_loop() rpc_future.set_result(msg) loop_communicator.add_rpc_subscriber(get_rpc, 'rpc') loop_communicator.rpc_send('rpc', MSG) result = await rpc_future assert result == MSG
async def test_task(self, loop_communicator): TASK = 'task this' task_future = plumpy.Future() loop = asyncio.get_event_loop() def get_task(_comm, msg): assert loop is asyncio.get_event_loop() task_future.set_result(msg) loop_communicator.add_task_subscriber(get_task) loop_communicator.task_send(TASK) result = await task_future assert result == TASK
def test_call_on_calculation_finish(self): loop = self.runner.loop proc = Proc(runner=self.runner) future = plumpy.Future() def calc_done(pk): self.assertEqual(pk, proc.node.pk) loop.stop() future.set_result(True) self.runner.call_on_calculation_finish(proc.node.pk, calc_done) # Run the calculation self.runner.loop.add_callback(proc.step_until_terminated) self._run_loop_for(5.) self.assertTrue(future.result())
def run_until_paused(proc): """ Set up a future that will be resolved on entering the WAITING state """ listener = plumpy.ProcessListener() paused = plumpy.Future() if proc.paused: paused.set_result(True) else: def on_paused(_paused_proc): paused.set_result(True) proc.remove_process_listener(listener) listener.on_process_paused = on_paused proc.add_process_listener(listener) return paused
def run_until_waiting(proc): """ Set up a future that will be resolved on entering the WAITING state """ from aiida.work import ProcessState listener = plumpy.ProcessListener() in_waiting = plumpy.Future() if proc.state == ProcessState.WAITING: in_waiting.set_result(True) else: def on_waiting(waiting_proc): in_waiting.set_result(True) proc.remove_process_listener(listener) listener.on_process_waiting = on_waiting proc.add_process_listener(listener) return in_waiting
def test_call_on_wf_finish(self): loop = self.runner.loop future = plumpy.Future() # Need to start() so it's stored wf = WorkflowDemo() wf.start() def wf_done(pk): self.assertEqual(pk, wf.pk) loop.stop() future.set_result(True) self.runner.call_on_legacy_workflow_finish(wf.pk, wf_done) # Run the wf while wf.is_running(): execute_steps() self._run_loop_for(10.) self.assertTrue(future.result())
def _continue(self, task): """ Continue the task Note that the task may already have been completed, as indicated from the corresponding the node, in which case it is not continued, but the corresponding future is reconstructed and returned. This scenario may occur when the Process was already completed by another worker that however failed to send the acknowledgment. :param task: the task to continue :raises plumpy.TaskRejected: if the node corresponding to the task cannot be loaded """ from aiida.common import exceptions from aiida.orm import load_node, Data try: node = load_node(pk=task[PID_KEY]) except (exceptions.MultipleObjectsError, exceptions.NotExistent) as exception: raise plumpy.TaskRejected( 'Cannot continue process: {}'.format(exception)) if node.is_terminated: future = plumpy.Future() if node.is_finished: future.set_result( dict(node.get_outputs(node_type=Data, also_labels=True))) elif node.is_excepted: future.set_exception(PastException(node.exception)) elif node.is_killed: future.set_exception(plumpy.KilledError()) return future return super(ProcessLauncher, self)._continue(task)