"""Test the wakeable_queue.Queue class
"""

import unittest
import os
import sys
import random
import logging
import threading
from time import sleep

logger = logging.getLogger(__name__)

from wakeable_queue import Queue, AbortRequested, WakeableQueueWorker, WorkerStatus

random.seed()

STOP_MSG = []

class Worker(WakeableQueueWorker):
    def __init__(self, worker_id, queue, consume_messages=False):
        WakeableQueueWorker.__init__(self, worker_id, queue, STOP_MSG, logger)
        self.consume_messages = consume_messages
        if consume_messages:
            logger.info("Worker %s will consume messages" % worker_id)
        
    def process_batch(self, data):
        if data[0]>=0:
            #logger.info("worker %s got message %s" % (self.worker_id, data[0]))
            # normal processing
            sleep(0.1)
            if not self.consume_messages:
                #logger.info("worker %s got %s" % (self.worker_id, data[0]))
                self.queue.put(data)
            else:
                logger.info("worker %s consuming message %s" % (self.worker_id, data[0]))
        elif data[0]==(-1): # a -1 means we should signal the abort
            logger.info("Worker %s requesting an abort" % self.worker_id)
            self.status = WorkerStatus.ABORT_REQUESTED
            self.queue.abort_request()
        elif data[0]==(-2): # a -2 means we should throw an exception
            raise Exception("Expected error")
        else:
            assert 0
            
            

NUM_WORKERS = 4

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


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    unittest.main(argv=sys.argv)
                           
