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 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")