def test_send_worker(caplog): TEST_FILE_NAME = "test_send_file.txt" class TestSendWorker(SendWorker): def startup(self): self.send_file = open(TEST_FILE_NAME, "w") send_q = MPQueue() send_q.put(EventMessage("TEST", "OBSERVATION", "SOME DATA 1")) send_q.put(EventMessage("TEST", "OBSERVATION", "SOME DATA 2")) send_q.put(EventMessage("TEST", "OBSERVATION", "SOME DATA 3")) send_q.put(EventMessage("TEST", "OBSERVATION", "SOME DATA 4")) send_q.put(EventMessage("TEST", "OBSERVATION", "SOME DATA 5")) try: items = _proc_worker_wrapper_helper(caplog, TestSendWorker, args=(send_q,), expect_shutdown_evt=True, alarm_secs=1) assert items == [] with open("test_send_file.txt", "r") as f: for idx, line in enumerate(f): assert line == f"OBSERVATION::SOME DATA {idx + 1}\n" finally: send_q.safe_close() if os.path.exists(TEST_FILE_NAME): os.remove(TEST_FILE_NAME)
def _proc_worker_wrapper_helper(caplog, worker_class, args=None, expect_shutdown_evt=True, alarm_secs=1.0): startup_evt = mp.Event() shutdown_evt = mp.Event() event_q = MPQueue() if args is None: args = () def alarm_handler(signal_num, current_stack_frame): shutdown_evt.set() if alarm_secs: signal.signal(signal.SIGALRM, alarm_handler) signal.setitimer(signal.ITIMER_REAL, alarm_secs) caplog.set_level(logging.DEBUG) exitcode = proc_worker_wrapper(worker_class, "TEST", startup_evt, shutdown_evt, event_q, *args) assert startup_evt.is_set() assert shutdown_evt.is_set() == expect_shutdown_evt items = list(event_q.drain()) assert items last_item = items[-1] assert last_item.msg_src == "TEST" assert last_item.msg_type == "SHUTDOWN" assert last_item.msg == "Normal" assert exitcode == 0 return items[:-1]
def test_queue_put(): Q = MPQueue(2) assert Q.safe_put("ITEM1") assert Q.safe_put("ITEM2") assert not Q.safe_put("ITEM3") num_left = Q.safe_close() assert num_left == 2
def test_proc_worker_no_main_func(caplog): startup_evt = mp.Event() shutdown_evt = mp.Event() event_q = MPQueue() try: caplog.set_level(logging.INFO) pw = ProcWorker("TEST", startup_evt, shutdown_evt, event_q) with pytest.raises(NotImplementedError): pw.main_func() finally: event_q.safe_close()
def test_queue_proc_worker(caplog): work_q = MPQueue() work_q.put(1) work_q.put(2) work_q.put(3) work_q.put(4) work_q.put("END") work_q.put(5) items = _proc_worker_wrapper_helper(caplog, QueueProcWorkerTest, args=(work_q, ), expect_shutdown_evt=False) assert len(items) == 4 assert items == [f'DONE {idx + 1}' for idx in range(4)]
def test_proc_full_stop(caplog): shutdown_evt = mp.Event() event_q = MPQueue() caplog.set_level(logging.INFO) proc = Proc("TEST", TimerProcWorkerTest, shutdown_evt, event_q) for idx in range(4): item = event_q.safe_get(1.0) assert item, f"idx: {idx}" assert item.startswith(f'TIMER {idx + 1} [') item = event_q.safe_get(1.0) assert item.msg_src == "TEST" assert item.msg_type == "SHUTDOWN" assert item.msg == "Normal" proc.full_stop(wait_time=0.5) assert not proc.proc.is_alive()
def test_proc_full_stop_need_terminate(caplog): class NeedTerminateWorker(ProcWorker): def main_loop(self): while True: time.sleep(1.0) shutdown_evt = mp.Event() event_q = MPQueue() caplog.set_level(logging.INFO) proc = Proc("TEST", NeedTerminateWorker, shutdown_evt, event_q) proc.full_stop(wait_time=0.1)
def test_proc_worker_run(caplog): startup_evt = mp.Event() shutdown_evt = mp.Event() event_q = MPQueue() caplog.set_level(logging.INFO) pw = ProcWorkerTest("TEST", startup_evt, shutdown_evt, event_q, "ARG1", "ARG2") assert not startup_evt.is_set() assert not shutdown_evt.is_set() pw.run() assert startup_evt.is_set() assert shutdown_evt.is_set() item = event_q.safe_get() assert item assert item.msg_src == "TEST" assert item.msg_type == "SHUTDOWN" assert item.msg == "Normal" assert f"MAIN_FUNC: ('ARG1', 'ARG2')" in caplog.text
def test_proc_worker_exception(caplog): class ProcWorkerException(ProcWorker): def main_func(self): raise NameError("Because this doesn't happen often") startup_evt = mp.Event() shutdown_evt = mp.Event() event_q = MPQueue() caplog.set_level(logging.INFO) with pytest.raises(SystemExit): proc_worker_wrapper(ProcWorkerException, "TEST", startup_evt, shutdown_evt, event_q) assert startup_evt.is_set() assert not shutdown_evt.is_set() item = event_q.safe_get() assert item assert item.msg_src == "TEST" assert item.msg_type == "FATAL" assert item.msg == "Because this doesn't happen often" assert f"Exception Shutdown" in caplog.text
def test_mpqueue_get(): Q = MPQueue() item = Q.safe_get(None) assert item is None Q.put("ITEM1") Q.put("ITEM2") assert Q.safe_get(0.02) == "ITEM1" assert Q.safe_get(0.02) == "ITEM2" assert Q.safe_get(0.02) is None assert Q.safe_get(None) is None num_left = Q.safe_close() assert num_left == 0
def test_proc_start_hangs(caplog): class StartHangWorker(ProcWorker): def startup(self): while True: time.sleep(1.0) shutdown_evt = mp.Event() event_q = MPQueue() caplog.set_level(logging.INFO) Proc.STARTUP_WAIT_SECS = 0.2 try: with pytest.raises(RuntimeError): Proc("TEST", StartHangWorker, shutdown_evt, event_q) finally: Proc.STARTUP_WAIT_SECS = 3.0
def test_drain_queue(): Q = MPQueue() items = list(Q.drain()) assert items == [] expected = [f"ITEM{idx}" for idx in range(10)] for item in expected: Q.put(item) items = list(Q.drain()) assert items == expected num_left = Q.safe_close() assert num_left == 0
def test_listen_worker(caplog): class TestListenWorker(ListenWorker): def _test_hook(self): request = self.event_q.get() assert request.msg == f"REQUEST {self._test_hook_idx + 1}" self.reply_q.put(request.msg.replace("REQUEST", "REPLY")) startup_evt = mp.Event() shutdown_evt = mp.Event() event_q = MPQueue() reply_q = MPQueue() lw = TestListenWorker('TEST', startup_evt, shutdown_evt, event_q, reply_q) try: lw.startup() NUM_TO_PROCESS = 5 def thread_worker(): for idx in range(NUM_TO_PROCESS): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('127.0.0.1', 9999)) try: sock.send(f"REQUEST {idx + 1}".encode("utf-8")) reply = sock.recv(1500).decode() assert reply == f"REPLY {idx + 1}" finally: sock.close() t = threading.Thread(target=thread_worker) t.start() for idx in range(NUM_TO_PROCESS): lw._test_hook_idx = idx lw.main_func() t.join() finally: lw.shutdown() event_q.safe_close() reply_q.safe_close()