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