# Some simple queue module tests, plus some failure conditions

# to ensure the Queue locks remain stable.

import Queue

import sys

import threading

import time

import unittest

from test import test_support

QUEUE_SIZE = 5

# A thread to run a function that unclogs a blocked Queue.

class _TriggerThread(threading.Thread):

    def __init__(self, fn, args):

        self.fn = fn

        self.args = args

        self.startedEvent = threading.Event()

        threading.Thread.__init__(self)

    def run(self):

        # The sleep isn't necessary, but is intended to give the blocking

        # function in the main thread a chance at actually blocking before

        # we unclog it.  But if the sleep is longer than the timeout-based

        # tests wait in their blocking functions, those tests will fail.

        # So we give them much longer timeout values compared to the

        # sleep here (I aimed at 10 seconds for blocking functions --

        # they should never actually wait that long - they should make

        # progress as soon as we call self.fn()).

        time.sleep(0.1)

        self.startedEvent.set()

        self.fn(*self.args)

# Execute a function that blocks, and in a separate thread, a function that

# triggers the release.  Returns the result of the blocking function.  Caution:

# block_func must guarantee to block until trigger_func is called, and

# trigger_func must guarantee to change queue state so that block_func can make

# enough progress to return.  In particular, a block_func that just raises an

# exception regardless of whether trigger_func is called will lead to

# timing-dependent sporadic failures, and one of those went rarely seen but

# undiagnosed for years.  Now block_func must be unexceptional.  If block_func

# is supposed to raise an exception, call do_exceptional_blocking_test()

# instead.

class BlockingTestMixin:

    def do_blocking_test(self, block_func, block_args, trigger_func, trigger_args):

        self.t = _TriggerThread(trigger_func, trigger_args)

        self.t.start()

        self.result = block_func(*block_args)

        # If block_func returned before our thread made the call, we failed!

        if not self.t.startedEvent.is_set():

            self.fail("blocking function '%r' appeared not to block" %

                      block_func)

        self.t.join(10) # make sure the thread terminates

        if self.t.is_alive():

            self.fail("trigger function '%r' appeared to not return" %

                      trigger_func)

        return self.result

    # Call this instead if block_func is supposed to raise an exception.

    def do_exceptional_blocking_test(self,block_func, block_args, trigger_func,

                                   trigger_args, expected_exception_class):

        self.t = _TriggerThread(trigger_func, trigger_args)

        self.t.start()

        try:

            try:

                block_func(*block_args)

            except expected_exception_class:

                raise

            else:

                self.fail("expected exception of kind %r" %

                                 expected_exception_class)

        finally:

            self.t.join(10) # make sure the thread terminates

            if self.t.is_alive():

                self.fail("trigger function '%r' appeared to not return" %

                                 trigger_func)

            if not self.t.startedEvent.is_set():

                self.fail("trigger thread ended but event never set")

class BaseQueueTest(unittest.TestCase, BlockingTestMixin):

    def setUp(self):

        self.cum = 0

        self.cumlock = threading.Lock()

    def simple_queue_test(self, q):

        if not q.empty():

            raise RuntimeError, "Call this function with an empty queue"

        # I guess we better check things actually queue correctly a little :)

        q.put(111)

        q.put(333)

        q.put(222)

        target_order = dict(Queue = [111, 333, 222],

                            LifoQueue = [222, 333, 111],

                            PriorityQueue = [111, 222, 333])

        actual_order = [q.get(), q.get(), q.get()]

        self.assertEquals(actual_order, target_order[q.__class__.__name__],

                          "Didn't seem to queue the correct data!")

        for i in range(QUEUE_SIZE-1):

            q.put(i)

            self.assert_(not q.empty(), "Queue should not be empty")

        self.assert_(not q.full(), "Queue should not be full")

        q.put("last")

        self.assert_(q.full(), "Queue should be full")

        try:

            q.put("full", block=0)

            self.fail("Didn't appear to block with a full queue")

        except Queue.Full:

            pass

        try:

            q.put("full", timeout=0.01)

            self.fail("Didn't appear to time-out with a full queue")

        except Queue.Full:

            pass

        # Test a blocking put

        self.do_blocking_test(q.put, ("full",), q.get, ())

        self.do_blocking_test(q.put, ("full", True, 10), q.get, ())

        # Empty it

        for i in range(QUEUE_SIZE):

            q.get()

        self.assert_(q.empty(), "Queue should be empty")

        try:

            q.get(block=0)

            self.fail("Didn't appear to block with an empty queue")

        except Queue.Empty:

            pass

        try:

            q.get(timeout=0.01)

            self.fail("Didn't appear to time-out with an empty queue")

        except Queue.Empty:

            pass

        # Test a blocking get

        self.do_blocking_test(q.get, (), q.put, ('empty',))

        self.do_blocking_test(q.get, (True, 10), q.put, ('empty',))

    def worker(self, q):

        while True:

            x = q.get()

            if x is None:

                q.task_done()

                return

            with self.cumlock:

                self.cum += x

            q.task_done()

    def queue_join_test(self, q):

        self.cum = 0

        for i in (0,1):

            threading.Thread(target=self.worker, args=(q,)).start()

        for i in xrange(100):

            q.put(i)

        q.join()

        self.assertEquals(self.cum, sum(range(100)),

                          "q.join() did not block until all tasks were done")

        for i in (0,1):

            q.put(None)         # instruct the threads to close

        q.join()                # verify that you can join twice

    def test_queue_task_done(self):

        # Test to make sure a queue task completed successfully.

        q = self.type2test()

        try:

            q.task_done()

        except ValueError:

            pass

        else:

            self.fail("Did not detect task count going negative")

    def test_queue_join(self):

        # Test that a queue join()s successfully, and before anything else

        # (done twice for insurance).

        q = self.type2test()

        self.queue_join_test(q)

        self.queue_join_test(q)

        try:

            q.task_done()

        except ValueError:

            pass

        else:

            self.fail("Did not detect task count going negative")

    def test_simple_queue(self):

        # Do it a couple of times on the same queue.

        # Done twice to make sure works with same instance reused.

        q = self.type2test(QUEUE_SIZE)

        self.simple_queue_test(q)

        self.simple_queue_test(q)

class QueueTest(BaseQueueTest):

    type2test = Queue.Queue

class LifoQueueTest(BaseQueueTest):

    type2test = Queue.LifoQueue

class PriorityQueueTest(BaseQueueTest):

    type2test = Queue.PriorityQueue

# A Queue subclass that can provoke failure at a moment's notice :)

class FailingQueueException(Exception):

    pass

class FailingQueue(Queue.Queue):

    def __init__(self, *args):

        self.fail_next_put = False

        self.fail_next_get = False

        Queue.Queue.__init__(self, *args)

    def _put(self, item):

        if self.fail_next_put:

            self.fail_next_put = False

            raise FailingQueueException, "You Lose"

        return Queue.Queue._put(self, item)

    def _get(self):

        if self.fail_next_get:

            self.fail_next_get = False

            raise FailingQueueException, "You Lose"

        return Queue.Queue._get(self)

class FailingQueueTest(unittest.TestCase, BlockingTestMixin):

    def failing_queue_test(self, q):

        if not q.empty():

            raise RuntimeError, "Call this function with an empty queue"

        for i in range(QUEUE_SIZE-1):

            q.put(i)

        # Test a failing non-blocking put.

        q.fail_next_put = True

        try:

            q.put("oops", block=0)

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        q.fail_next_put = True

        try:

            q.put("oops", timeout=0.1)

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        q.put("last")

        self.assert_(q.full(), "Queue should be full")

        # Test a failing blocking put

        q.fail_next_put = True

        try:

            self.do_blocking_test(q.put, ("full",), q.get, ())

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        # Check the Queue isn't damaged.

        # put failed, but get succeeded - re-add

        q.put("last")

        # Test a failing timeout put

        q.fail_next_put = True

        try:

            self.do_exceptional_blocking_test(q.put, ("full", True, 10), q.get, (),

                                              FailingQueueException)

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        # Check the Queue isn't damaged.

        # put failed, but get succeeded - re-add

        q.put("last")

        self.assert_(q.full(), "Queue should be full")

        q.get()

        self.assert_(not q.full(), "Queue should not be full")

        q.put("last")

        self.assert_(q.full(), "Queue should be full")

        # Test a blocking put

        self.do_blocking_test(q.put, ("full",), q.get, ())

        # Empty it

        for i in range(QUEUE_SIZE):

            q.get()

        self.assert_(q.empty(), "Queue should be empty")

        q.put("first")

        q.fail_next_get = True

        try:

            q.get()

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        self.assert_(not q.empty(), "Queue should not be empty")

        q.fail_next_get = True

        try:

            q.get(timeout=0.1)

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        self.assert_(not q.empty(), "Queue should not be empty")

        q.get()

        self.assert_(q.empty(), "Queue should be empty")

        q.fail_next_get = True

        try:

            self.do_exceptional_blocking_test(q.get, (), q.put, ('empty',),

                                              FailingQueueException)

            self.fail("The queue didn't fail when it should have")

        except FailingQueueException:

            pass

        # put succeeded, but get failed.

        self.assert_(not q.empty(), "Queue should not be empty")

        q.get()

        self.assert_(q.empty(), "Queue should be empty")

    def test_failing_queue(self):

        # Test to make sure a queue is functioning correctly.

        # Done twice to the same instance.

        q = FailingQueue(QUEUE_SIZE)

        self.failing_queue_test(q)

        self.failing_queue_test(q)

def test_main():

    test_support.run_unittest(QueueTest, LifoQueueTest, PriorityQueueTest,

                              FailingQueueTest)

if __name__ == "__main__":

    test_main()