def setUp(self):
self.queue = Queue()
self.workers = [Worker(i, self.queue, False) for i in range(NUM_WORKERS)]
class TestWakeableQueue(unittest.TestCase):
def setUp(self):
self.queue = Queue()
self.workers = [Worker(i, self.queue, False) for i in range(NUM_WORKERS)]
def test_normal_processing(self):
logger.info("Runing test_normal_processing")
self.workers.append(Worker(NUM_WORKERS, self.queue, consume_messages=True))
for i in range(100):
self.queue.put([i+1, 2, 3, 4, 5])
WakeableQueueWorker.run_workers(self.workers, self.queue, STOP_MSG,
logger)
for w in self.workers:
self.assertTrue(w.status==WorkerStatus.STOPPED,
"Worker %s not stopped, status was %s" %
(w.worker_id, w.status))
def test_abort_processing(self):
logger.info("Runing test_abort_processing")
poison_pill = random.randint(25, 80)
for i in range(100):
if i == poison_pill:
logger.info("poison pill is at %d" % poison_pill)
self.queue.put([-1])
else:
self.queue.put([1, 2, 3, 4, 5])
try:
WakeableQueueWorker.run_workers(self.workers, self.queue, STOP_MSG,
logger)
self.assertTrue(False, "Did not get the expected abort request")
except AbortRequested:
logger.info("Master got AbortRequested, as expected")
for w in self.workers:
self.assertTrue(w.status==WorkerStatus.ABORTED or
w.status==WorkerStatus.ABORT_REQUESTED,
"Worker %s has status %s, expecting an abort" % (w.worker_id, w.status))
logger.info("All workers aborted as expected")
def test_exception_processing(self):
logger.info("Runing test_exception_processing")
poison_pill = random.randint(25, 80)
for i in range(100):
if i == poison_pill:
logger.info("poison pill is at %d" % poison_pill)
self.queue.put([-2])
else:
self.queue.put([1, 2, 3, 4, 5])
try:
WakeableQueueWorker.run_workers(self.workers, self.queue, STOP_MSG,
logger)
self.assertTrue(False, "Did not get the expected abort request")
except AbortRequested:
logger.info("Master got AbortRequested, as expected")
for w in self.workers:
self.assertTrue(w.status==WorkerStatus.ABORTED or
w.status==WorkerStatus.ABORT_REQUESTED,
"Worker %s has status %s, expecting an abort" % (w.worker_id, w.status))
logger.info("All workers aborted as expected")
def test_single_batch(self):
"""Test case where we have two workers, both competing for a single
batch.
"""
logger.info("Runing test_single_batch")
self.workers = [Worker(i, self.queue, True) for i in range(2)]
self.queue.put([1, 2, 3, 4, 5])
WakeableQueueWorker.run_workers(self.workers, self.queue, STOP_MSG,
logger)
for w in self.workers:
self.assertTrue(w.status==WorkerStatus.STOPPED,
"Worker %s not stopped, status was %s" %
(w.worker_id, w.status))
