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)
