class TestWaitOnProcessStateEvent(TestCase):
def setUp(self):
super(TestWaitOnProcessStateEvent, self).setUp()
self.executor = ThreadExecutor()
def tearDown(self):
self.executor.shutdown()
def test_already_in_state(self):
p = DummyProcess.new()
self.assertTrue(WaitOnProcessState(p, ProcessState.CREATED).wait(timeout=2.))
def test_state_messages(self):
tp = PythonThreadPoolExecutor(max_workers=1)
for state in (ProcessState.RUNNING, ProcessState.STOPPED):
p = DummyProcess.new()
waiton = WaitOnProcessState(p, state)
future = tp.submit(waiton.wait)
while not future.running():
pass
p.play()
self.assertTrue(future.result(timeout=2.))
def test_waiting_state(self):
tp = PythonThreadPoolExecutor(max_workers=1)
p = WaitForSignalProcess.new()
waiton = WaitOnProcessState(p, ProcessState.WAITING)
future = tp.submit(waiton.wait)
self.executor.play(p)
self.assertTrue(future.result(timeout=2.))
self.assertTrue(p.abort(timeout=2.))
def test_interrupt(self):
tp = PythonThreadPoolExecutor(max_workers=1)
p = DummyProcess.new()
waiton = WaitOnProcessState(p, ProcessState.STOPPED)
future = tp.submit(waiton.wait)
while not future.running():
pass
with self.assertRaises(Interrupted):
waiton.interrupt()
future.result(timeout=2.)
def test_interrupt_not_waiting(self):
"""
If you interrupt when it's not waiting then nothing happens.
"""
p = DummyProcess.new()
waiton = WaitOnProcessState(p, ProcessState.STOPPED)
waiton.interrupt()
def test_wait_until(self):
p = WaitForSignalProcess.new()
self.executor.play(p)
self.assertTrue(wait_until(p, ProcessState.WAITING, timeout=1.))
class TestProcess(TestCase):
def setUp(self):
super(TestProcess, self).setUp()
self.events_tester = ProcessListenerTester()
self.proc = DummyProcessWithOutput.new()
self.proc.add_process_listener(self.events_tester)
self.executor = ThreadExecutor()
def tearDown(self):
self.proc.remove_process_listener(self.events_tester)
self.executor.shutdown()
super(TestProcess, self).tearDown()
def test_spec(self):
"""
Check that the references to specs are doing the right thing...
"""
dp = DummyProcess.new()
self.assertIsNot(DummyProcess.spec(), Process.spec())
self.assertIs(dp.spec(), DummyProcess.spec())
class Proc(DummyProcess):
pass
self.assertIsNot(Proc.spec(), Process.spec())
self.assertIsNot(Proc.spec(), DummyProcess.spec())
p = Proc.new()
self.assertIs(p.spec(), Proc.spec())
def test_dynamic_inputs(self):
class NoDynamic(Process):
def _run(self, **kwargs):
pass
class WithDynamic(Process):
@classmethod
def define(cls, spec):
super(WithDynamic, cls).define(spec)
spec.dynamic_input()
def _run(self, **kwargs):
pass
with self.assertRaises(ValueError):
NoDynamic.run(a=5)
WithDynamic.run(a=5)
def test_inputs(self):
class Proc(Process):
@classmethod
def define(cls, spec):
super(Proc, cls).define(spec)
spec.input('a')
def _run(self, a):
pass
p = Proc.new({'a': 5})
# Check that we can access the inputs after creating
self.assertEqual(p.raw_inputs.a, 5)
with self.assertRaises(AttributeError):
p.raw_inputs.b
def test_inputs_default(self):
class Proc(DummyProcess):
@classmethod
def define(cls, spec):
super(Proc, cls).define(spec)
spec.input("input", default=5, required=False)
# Supply a value
p = Proc.new(inputs={'input': 2})
self.assertEqual(p.inputs['input'], 2)
# Don't supply, use default
p = Proc.new()
self.assertEqual(p.inputs['input'], 5)
def test_run(self):
p = DummyProcessWithOutput.new()
p.play()
self.assertTrue(p.has_finished())
self.assertEqual(p.state, ProcessState.STOPPED)
self.assertEqual(p.outputs, {'default': 5})
def test_run_from_class(self):
# Test running through class method
results = DummyProcessWithOutput.run()
self.assertEqual(results['default'], 5)
def test_forget_to_call_parent(self):
for event in ('start', 'run', 'finish', 'stop'):
with self.assertRaises(AssertionError):
ForgetToCallParent.run(forget_on=event)
def test_pid(self):
# Test auto generation of pid
p = DummyProcessWithOutput.new()
self.assertIsNotNone(p.pid)
# Test using integer as pid
p = DummyProcessWithOutput.new(pid=5)
self.assertEquals(p.pid, 5)
# Test using string as pid
p = DummyProcessWithOutput.new(pid='a')
self.assertEquals(p.pid, 'a')
def test_exception(self):
proc = ExceptionProcess.new()
with self.assertRaises(RuntimeError):
proc.play()
self.assertEqual(proc.state, ProcessState.FAILED)
def test_get_description(self):
# Not all that much we can test for, but check if it's a string at
# least
for ProcClass in TEST_PROCESSES:
desc = ProcClass.get_description()
self.assertIsInstance(desc, str)
# Dummy process should at least use the docstring as part of the
# description and so it shouldn't be empty
desc = DummyProcess.get_description()
self.assertNotEqual(desc, "")
def test_created_bundle(self):
"""
Check that the bundle after just creating a process is as we expect
:return:
"""
proc = DummyProcessWithOutput.new()
b = Bundle()
proc.save_instance_state(b)
self.assertIsNone(b.get('inputs', None))
self.assertEqual(len(b['outputs']), 0)
def test_instance_state(self):
proc = DummyProcessWithOutput.new()
saver = ProcessSaver(proc)
proc.play()
for info, outputs in zip(saver.snapshots, saver.outputs):
state, bundle = info
# Check that it is a copy
self.assertIsNot(
outputs, bundle[Process.BundleKeys.OUTPUTS.value].get_dict())
# Check the contents are the same
self.assertEqual(
outputs, bundle[Process.BundleKeys.OUTPUTS.value].get_dict())
self.assertIsNot(
proc.outputs,
saver.snapshots[-1][1][Process.BundleKeys.OUTPUTS.value])
def test_saving_each_step(self):
for ProcClass in TEST_PROCESSES:
proc = ProcClass.new()
saver = ProcessSaver(proc)
proc.play()
self.assertEqual(proc.state, ProcessState.STOPPED)
self.assertTrue(
check_process_against_snapshots(ProcClass, saver.snapshots))
def test_saving_each_step_interleaved(self):
for ProcClass in TEST_PROCESSES:
proc = ProcClass.new()
ps = ProcessSaver(proc)
try:
proc.play()
except BaseException:
pass
self.assertTrue(
check_process_against_snapshots(ProcClass, ps.snapshots))
def test_logging(self):
class LoggerTester(Process):
def _run(self, **kwargs):
self.logger.info("Test")
# TODO: Test giving a custom logger to see if it gets used
p = LoggerTester.new()
p.play()
def test_abort(self):
proc = DummyProcess.new()
proc.abort()
self.assertTrue(proc.has_aborted())
self.assertEqual(proc.state, ProcessState.STOPPED)
def test_wait_continue(self):
p = WaitForSignalProcess.new()
self.executor.play(p)
self.assertTrue(wait_until(p, ProcessState.WAITING, timeout=1.))
p.continue_()
self.assertTrue(p.wait(timeout=2.))
self.assertTrue(p.has_finished())
def test_wait_pause_continue_play(self):
p = WaitForSignalProcess.new()
# Play the process and wait until it is waiting
self.executor.play(p)
# Wait
self.assertTrue(wait_until(p, ProcessState.WAITING, 1.))
# Pause
self.assertTrue(p.pause(timeout=1.))
# Continue
p.continue_()
# Play
p.play()
self.assertTrue(p.wait(timeout=1.))
self.assertEqual(p.state, ProcessState.STOPPED)
def test_wait_pause_play_continue(self):
p = WaitForSignalProcess.new()
fut = self.executor.play(p)
# Wait
self.assertTrue(wait_until(p, ProcessState.WAITING, 1.))
self.assertTrue(p.is_playing())
# Pause
self.assertTrue(p.pause(timeout=1.))
self.assertFalse(p.is_playing())
# Play
fut.play()
self.assertEqual(p.state, ProcessState.WAITING)
# Continue
p.continue_()
self.assertTrue(wait_until(p, ProcessState.STOPPED), 1.)
def test_exc_info(self):
p = ExceptionProcess.new()
try:
p.play()
except BaseException:
import sys
exc_info = sys.exc_info()
p_exc_info = p.get_exc_info()
self.assertEqual(p_exc_info[0], exc_info[0])
self.assertEqual(p_exc_info[1], exc_info[1])
def test_exception_in_on_playing(self):
class P(DummyProcess):
def on_playing(self):
raise RuntimeError("Cope with this")
p = P.new()
with self.assertRaises(RuntimeError):
p.play()
def test_exception_in_done_playing(self):
class P(DummyProcess):
def on_done_playing(self):
raise RuntimeError("Cope with this")
p = P.new()
with self.assertRaises(RuntimeError):
p.play()
def test_restart(self):
p = _RestartProcess.new()
future = self.executor.play(p)
self.assertTrue(wait_until(p, ProcessState.WAITING, timeout=2.))
# Save the state of the process
bundle = Bundle()
p.save_instance_state(bundle)
self.assertTrue(future.abort(timeout=2.))
# Load a process from the saved state
p = _RestartProcess.create_from(bundle)
self.assertEqual(p.state, ProcessState.WAITING)
# Now play it
future = self.executor.play(p)
p.continue_()
self.assertEqual(future.result(timeout=1.0), {'finished': True})
def test_wait(self):
p = DummyProcess.new()
self.assertTrue(p.wait(timeout=2.),
"Not running process didn't return from wait")
self.executor.play(p)
self.assertTrue(p.wait(timeout=2.),
"Process failed to return from wait when done")
def test_wait_pause_play(self):
p = WaitForSignalProcess.new()
self.executor.play(p)
# Wait
self.assertTrue(wait_until(p, ProcessState.WAITING, timeout=2.))
# Pause
self.assertTrue(p.pause(timeout=1.))
# Play
self.executor.play(p)
self.assertTrue(wait_until(p, ProcessState.WAITING, timeout=2.))
def test_pause_play(self):
"""
Pausing a process that is not playing should have no effect and next
time it is played it should continue normally.
"""
p = DummyProcess.new()
self.assertTrue(p.pause(timeout=1.))
p.play()
self.assertEqual(p.state, ProcessState.STOPPED)
def test_play_terminated(self):
p = DummyProcess()
p.play()
self.assertTrue(p.has_terminated())
p.play()
def _check_process_against_snapshot(self, snapshot, proc):
self.assertEqual(snapshot.state, proc.state)
new_bundle = Bundle()
proc.save_instance_state(new_bundle)
self.assertEqual(
snapshot.bundle, new_bundle,
"Bundle mismatch with process class {}\n"
"Snapshot:\n{}\n"
"Loaded:\n{}".format(proc.__class__, snapshot.bundle, new_bundle))
self.assertEqual(
snapshot.outputs, proc.outputs,
"Outputs mismatch with process class {}\n"
"Snapshot:\n{}\n"
"Loaded:\n{}".format(proc.__class__, snapshot.outputs,
proc.outputs))
class ProcessController(ProcessListener):
def __init__(self, executor=None, remove_on_terminate=True):
"""
:param executor: The executor used to play processes
:param remove_on_terminate: Automatically remove any process that terminates
from the control of this class.
"""
if executor is None:
self._executor = ThreadExecutor()
else:
self._executor = executor
self._processes = {}
self._remove_on_terminate = remove_on_terminate
def insert(self, process):
"""
Insert a process into the controller
:param process: The process
:type process: :class:`plum.process.Process`
"""
self._processes[process.pid] = process
process.add_process_listener(self)
def insert_and_play(self, process):
self.insert(process)
self.play(process.pid)
def remove(self, pid, timeout=None):
try:
if not self._executor.pause(pid, timeout):
return False
except ValueError:
pass
self._processes[pid].remove_process_listener(self)
del self._processes[pid]
return True
def remove_all(self, timeout=None):
num_removed = 0
time_left = timeout
t0 = time.time()
for pid in self._processes.keys():
if not self.remove(pid, time_left):
return num_removed
num_removed += 1
if time_left is not None:
time_left = timeout - (time.time() - t0)
return num_removed
def play(self, pid):
self._executor.play(self._processes[pid])
def pause(self, pid, timeout=None):
try:
return self._executor.pause(pid, timeout)
except ValueError:
return False
def pause_all(self, timeout=None):
return self._executor.pause_all(timeout)
def abort(self, pid, message=None, timeout=None):
try:
return self._executor.abort(pid, message, timeout)
except ValueError:
return False
def abort_all(self, msg=None, timeout=None):
return self._executor.abort_all(msg, timeout)
def on_process_done_playing(self, process):
if self._remove_on_terminate and process.has_terminated():
self.remove(process.pid)
def get_process(self, pid):
return self._processes[pid]
def get_processes(self):
return self._processes.values()
def get_num_processes(self):
return len(self._processes)
class TestPicklePersistence(TestCase):
def setUp(self):
import tempfile
super(TestPicklePersistence, self).setUp()
self.store_dir = tempfile.mkdtemp()
self.pickle_persistence = PicklePersistence(
running_directory=self.store_dir)
self.procman = ThreadExecutor()
def tearDown(self):
super(TestPicklePersistence, self).tearDown()
self.pickle_persistence.clear_all_persisted()
self._empty_directory()
def test_store_directory(self):
self.assertEqual(self.store_dir,
self.pickle_persistence.store_directory)
def test_on_create_process(self):
proc = ProcessWithCheckpoint.new()
self.pickle_persistence.persist_process(proc)
save_path = self.pickle_persistence.get_running_path(proc.pid)
self.assertTrue(os.path.isfile(save_path))
def test_on_waiting_process(self):
proc = WaitForSignalProcess.new()
self.pickle_persistence.persist_process(proc)
save_path = self.pickle_persistence.get_running_path(proc.pid)
future = self.procman.play(proc)
# Check the file exists
self.assertTrue(os.path.isfile(save_path))
self.assertTrue(future.abort(timeout=2.))
def test_on_finishing_process(self):
proc = ProcessWithCheckpoint.new()
pid = proc.pid
self.pickle_persistence.persist_process(proc)
running_path = self.pickle_persistence.get_running_path(proc.pid)
self.assertTrue(os.path.isfile(running_path))
proc.play()
self.assertFalse(os.path.isfile(running_path))
finished_path = \
os.path.join(self.store_dir,
self.pickle_persistence.finished_directory,
self.pickle_persistence.pickle_filename(pid))
self.assertTrue(os.path.isfile(finished_path))
def test_load_all_checkpoints(self):
self._empty_directory()
# Create some processes
for i in range(0, 3):
proc = ProcessWithCheckpoint.new(pid=i)
self.pickle_persistence.save(proc)
# Check that the number of checkpoints matches we expected
num_cps = len(self.pickle_persistence.load_all_checkpoints())
self.assertEqual(num_cps, 3)
def test_save(self):
proc = ProcessWithCheckpoint.new()
running_path = self.pickle_persistence.get_running_path(proc.pid)
self.pickle_persistence.save(proc)
self.assertTrue(os.path.isfile(running_path))
def test_persist_twice(self):
proc = WaitForSignalProcess.new()
self.pickle_persistence.persist_process(proc)
future = self.procman.play(proc)
# Try persisting the process again using another persistence manager
try:
PicklePersistence(
running_directory=self.store_dir).persist_process(proc)
except LockError:
pass
self.assertTrue(proc.abort(timeout=1.))
def _empty_directory(self):
import shutil
if os.path.isdir(self.store_dir):
shutil.rmtree(self.store_dir)
class TestThreadExecutor(TestCase):
def setUp(self):
self.assertEqual(len(MONITOR.get_pids()), 0)
self.executor = ThreadExecutor()
def tearDown(self):
self.executor.shutdown()
def test_launch_simple(self):
class Tester(ProcessMonitorListener):
def __init__(self):
self.proc_class = None
self.stopped = False
def on_monitored_process_registered(self, process):
self.proc_class = process.__class__
def on_monitored_process_stopped(self, process):
self.stopped = True
t = Tester()
with MONITOR.listen(t):
self.executor.launch(DummyProcess)
while not t.stopped:
pass
self.assertIs(t.proc_class, DummyProcess)
def test_play(self):
p = DummyProcess.new()
self.assertFalse(p.has_finished())
fut = self.executor.play(p)
self.assertTrue(fut.wait(timeout=1.))
self.assertTrue(p.has_terminated())
self.assertTrue(p.has_finished())
def test_pause_all(self):
procs = []
# Launch a bunch of processes
for i in range(0, 9):
procs.append(WaitForSignalProcess.new())
self.executor.play(procs[-1])
self.assertTrue(wait_until(procs, ProcessState.WAITING, timeout=5))
# Check they are all in state we expect
for p in procs:
self.assertTrue(p.is_playing())
# Now try and pause them all
self.executor.pause_all()
# Check they are all in state we expect
for p in procs:
self.assertEqual(p.state, ProcessState.WAITING)
self.assertFalse(p.is_playing())
def test_play_pause_abort(self):
num_procs = 10
procs = []
for i in range(0, num_procs):
procs.append(WaitForSignalProcess.new())
self.executor.play(procs[-1])
# Wait
self.assertTrue(wait_until(procs, ProcessState.WAITING))
self.assertEqual(self.executor.pause_all(timeout=3.), num_procs)
def test_future_pid(self):
p = DummyProcess.new()
future = self.executor.play(p)
self.assertEqual(future.pid, p.pid)
def test_future_abort(self):
p = WaitForSignalProcess.new()
future = self.executor.play(p)
# Wait
self.assertTrue(wait_until(p, ProcessState.WAITING))
self.assertTrue(p.is_playing())
# Abort
self.assertTrue(future.abort(timeout=3.))
self.assertTrue(p.has_aborted())
def test_future_pause_play(self):
p = WaitForSignalProcess.new()
future = self.executor.play(p)
# Wait
self.assertTrue(wait_until(p, ProcessState.WAITING))
self.assertTrue(p.is_playing())
# Pause
self.assertTrue(future.pause(timeout=3.))
self.assertFalse(p.is_playing())
# Play
future.play()
p.continue_()
self.assertTrue(future.wait(timeout=1.))
def test_abort(self):
"""
Test aborting a process through the process manager
"""
self.assertEqual(self.executor.get_num_processes(), 0)
proc = WaitForSignalProcess.new()
future = self.executor.play(proc)
self.assertTrue(future.abort(timeout=2.))
self.assertEqual(self.executor.get_num_processes(), 0)
def test_abort_interrupt(self):
"""
Test aborting a process through the process manager
"""
self.assertEqual(self.executor.get_num_processes(), 0)
proc = WaitForSignalProcess.new()
# Start a process and make sure it is waiting
future = self.executor.play(proc)
wait_until(proc, ProcessState.WAITING)
# Then interrupt by aborting
self.assertTrue(future.abort(timeout=2.))
self.assertEqual(self.executor.get_num_processes(), 0)
def test_abort_future(self):
"""
Test aborting a process through the future
"""
self.assertEqual(self.executor.get_num_processes(), 0)
proc = WaitForSignalProcess.new()
future = self.executor.play(proc)
wait_until(proc, ProcessState.WAITING)
self.assertTrue(future.abort(timeout=2.))
self.assertEqual(self.executor.get_num_processes(), 0)
def test_get_processes(self):
p = WaitForSignalProcess.new()
self.executor.play(p)
procs = self.executor.get_processes()
self.assertEqual(len(procs), 1)
self.assertIs(procs[0], p)
self.assertTrue(p.abort(timeout=2.), "Failed to abort process")