def test_tasks(self):
"""
Test task related functions
"""
comms = WorkerComms(mp.get_context('fork'), 2)
comms.init_comms(False, False)
# Nothing available yet
with self.assertRaises(queue.Empty):
comms._tasks_queue.get(block=False)
# Add a few tasks
comms.add_task(12)
comms.add_task('hello world')
comms.add_task({'foo': 'bar'})
self.assertListEqual([comms.get_task() for _ in range(3)],
[12, 'hello world', {
'foo': 'bar'
}])
[comms.task_done() for _ in range(3)]
# Throw in an exception. Should return None
comms.set_exception()
self.assertIsNone(comms.get_task())
# Should be joinable
comms.join_tasks_queue()
def test_worker_poison_pill(self):
"""
Test that a poison pill is inserted for every worker
"""
for n_jobs in [1, 2, 4]:
with self.subTest(n_jobs=n_jobs):
comms = WorkerComms(mp.get_context('fork'), n_jobs)
comms.init_comms(False, False)
comms.insert_poison_pill()
self.assertListEqual([comms.get_task() for _ in range(n_jobs)],
[POISON_PILL for _ in range(n_jobs)])
[comms.task_done() for _ in range(n_jobs)]
comms.join_tasks_queue()