def test_main_loop_adds_pubsub_messages_to_event_queues():
"""
PUBSUB messages should be added to event queues.
"""
mock_ctx = mock.MagicMock()
event_q = MPQueue()
event_q.put(EventMessage("TEST", "PUBSUB", msg="1"))
event_q.put(EventMessage("TEST", "PUBSUB", msg="2"))
event_q.put(EventMessage("TEST", "PUBSUB", msg="3"))
event_q.put(EventMessage("TEST", "END", msg="foo"))
mock_ctx.event_queue = event_q
# one processing loop before shutdown in set, at which point the loop
# should exit with three messages still in the event queue
mock_ctx.shutdown_event.is_set.return_value = False
q1 = MPQueue()
q2 = MPQueue()
main_loop(mock_ctx, [q1, q2])
assert q1.safe_close() == 3
assert q2.safe_close() == 3
event_q.safe_close()
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_main_loop_ignores_and_logs_events_of_unknown_types():
"""
Loop should log events it doesn't know how to handle.
"""
mock_ctx = mock.MagicMock()
event_q = MPQueue()
event_q.put(EventMessage("TEST", "FOO", msg="1"))
mock_ctx.event_queue = event_q
# one processing loop before shutdown in set, at which point the loop
# should exit with three messages still in the event queue
mock_ctx.shutdown_event.is_set.side_effect = [False, True]
main_loop(mock_ctx, [])
event_q.safe_close()
mock_ctx.log.assert_called_once()
assert "Unknown Event" in mock_ctx.log.call_args[0][1]
def test_queue_put():
# Create MPQueue of max size 2
q = MPQueue(2)
# Putting two items should succeed and return True
assert q.safe_put("ITEM1")
assert q.safe_put("ITEM2")
# But once full, puts should return False
assert not q.safe_put("ITEM3")
num_left = q.safe_close()
assert num_left == 2
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_main_loop_ends_on_fatal_message():
"""
Main loop should terminate when fatal messsage is received.
"""
mock_ctx = mock.MagicMock()
event_q = MPQueue()
event_q.put(EventMessage("TEST", "FATAL", msg="foo"))
event_q.put(EventMessage("TEST", "END", msg="foo"))
mock_ctx.event_queue = event_q
mock_ctx.shutdown_event.is_set.return_value = False
main_loop(mock_ctx, [])
assert event_q.safe_close() == 1
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_main_loop_checks_shutdown_event_after_every_queue_get():
"""
Loop should regularly check shutdown event,
"""
mock_ctx = mock.MagicMock()
event_q = MPQueue()
mock_ctx.event_queue.safe_get.side_effect = [
False,
False,
EventMessage("TEST", "END", msg="foo"),
]
# loop won't exit as a result of shutdown_event being True
mock_ctx.shutdown_event.is_set.side_effect = [False, False, False, False, False]
main_loop(mock_ctx, [])
assert event_q.safe_close() == 0
assert mock_ctx.shutdown_event.is_set.call_count == 3
def test_main_loop_ends_when_shutdown_event_is_set():
"""
Main loop should terminate when shutdown event is set.
"""
mock_ctx = mock.MagicMock()
event_q = MPQueue()
event_q.put(EventMessage("TEST", "PUBSUB", msg="foo"))
event_q.put(EventMessage("TEST", "PUBSUB", msg="foo"))
event_q.put(EventMessage("TEST", "PUBSUB", msg="foo"))
event_q.put(EventMessage("TEST", "END", msg="foo"))
mock_ctx.event_queue = event_q
# one processing loop before shutdown in set, at which point the loop
# should exit with two messages still in the event queue
mock_ctx.shutdown_event.is_set.side_effect = [False, False, True]
main_loop(mock_ctx, [])
assert event_q.safe_close() == 2