Lib/test/test_asyncio/test_queues.py

"""Tests for queues.py"""

import asyncio
import unittest
from types import GenericAlias


def tearDownModule():
    asyncio.set_event_loop(None)


class QueueBasicTests(unittest.IsolatedAsyncioTestCase):

    async def _test_repr_or_str(self, fn, expect_id):
        """Test Queue's repr or str.

        fn is repr or str. expect_id is True if we expect the Queue's id to
        appear in fn(Queue()).
        """
        q = asyncio.Queue()
        self.assertStartsWith(fn(q), '<Queue')
        id_is_present = hex(id(q)) in fn(q)
        self.assertEqual(expect_id, id_is_present)

        # getters
        q = asyncio.Queue()
        async with asyncio.TaskGroup() as tg:
            # Start a task that waits to get.
            getter = tg.create_task(q.get())
            # Let it start waiting.
            await asyncio.sleep(0)
            self.assertTrue('_getters[1]' in fn(q))
            # resume q.get coroutine to finish generator
            q.put_nowait(0)

        self.assertEqual(0, await getter)

        # putters
        q = asyncio.Queue(maxsize=1)
        async with asyncio.TaskGroup() as tg:
            q.put_nowait(1)
            # Start a task that waits to put.
            putter = tg.create_task(q.put(2))
            # Let it start waiting.
            await asyncio.sleep(0)
            self.assertTrue('_putters[1]' in fn(q))
            # resume q.put coroutine to finish generator
            q.get_nowait()

        self.assertTrue(putter.done())

        q = asyncio.Queue()
        q.put_nowait(1)
        self.assertTrue('_queue=[1]' in fn(q))

    async def test_repr(self):
        await self._test_repr_or_str(repr, True)

    async def test_str(self):
        await self._test_repr_or_str(str, False)

    def test_generic_alias(self):
        q = asyncio.Queue[int]
        self.assertEqual(q.__args__, (int,))
        self.assertIsInstance(q, GenericAlias)

    async def test_empty(self):
        q = asyncio.Queue()
        self.assertTrue(q.empty())
        await q.put(1)
        self.assertFalse(q.empty())
        self.assertEqual(1, await q.get())
        self.assertTrue(q.empty())

    async def test_full(self):
        q = asyncio.Queue()
        self.assertFalse(q.full())

        q = asyncio.Queue(maxsize=1)
        await q.put(1)
        self.assertTrue(q.full())

    async def test_order(self):
        q = asyncio.Queue()
        for i in [1, 3, 2]:
            await q.put(i)

        items = [await q.get() for _ in range(3)]
        self.assertEqual([1, 3, 2], items)

    async def test_maxsize(self):
        q = asyncio.Queue(maxsize=2)
        self.assertEqual(2, q.maxsize)
        have_been_put = []

        async def putter():
            for i in range(3):
                await q.put(i)
                have_been_put.append(i)
            return True

        t = asyncio.create_task(putter())
        for i in range(2):
            await asyncio.sleep(0)

        # The putter is blocked after putting two items.
        self.assertEqual([0, 1], have_been_put)
        self.assertEqual(0, await q.get())

        # Let the putter resume and put last item.
        await asyncio.sleep(0)
        self.assertEqual([0, 1, 2], have_been_put)
        self.assertEqual(1, await q.get())
        self.assertEqual(2, await q.get())

        self.assertTrue(t.done())
        self.assertTrue(t.result())


class QueueGetTests(unittest.IsolatedAsyncioTestCase):

    async def test_blocking_get(self):
        q = asyncio.Queue()
        q.put_nowait(1)

        self.assertEqual(1, await q.get())

    async def test_get_with_putters(self):
        loop = asyncio.get_running_loop()

        q = asyncio.Queue(1)
        await q.put(1)

        waiter = loop.create_future()
        q._putters.append(waiter)

        self.assertEqual(1, await q.get())
        self.assertTrue(waiter.done())
        self.assertIsNone(waiter.result())

    async def test_blocking_get_wait(self):
        loop = asyncio.get_running_loop()
        q = asyncio.Queue()
        started = asyncio.Event()
        finished = False

        async def queue_get():
            nonlocal finished
            started.set()
            res = await q.get()
            finished = True
            return res

        queue_get_task = asyncio.create_task(queue_get())
        await started.wait()
        self.assertFalse(finished)
        loop.call_later(0.01, q.put_nowait, 1)
        res = await queue_get_task
        self.assertTrue(finished)
        self.assertEqual(1, res)

    def test_nonblocking_get(self):
        q = asyncio.Queue()
        q.put_nowait(1)
        self.assertEqual(1, q.get_nowait())

    def test_nonblocking_get_exception(self):
        q = asyncio.Queue()
        self.assertRaises(asyncio.QueueEmpty, q.get_nowait)

    async def test_get_cancelled_race(self):
        q = asyncio.Queue()

        t1 = asyncio.create_task(q.get())
        t2 = asyncio.create_task(q.get())

        await asyncio.sleep(0)
        t1.cancel()
        await asyncio.sleep(0)
        self.assertTrue(t1.done())
        await q.put('a')
        await asyncio.sleep(0)
        self.assertEqual('a', await t2)

    async def test_get_with_waiting_putters(self):
        q = asyncio.Queue(maxsize=1)
        asyncio.create_task(q.put('a'))
        asyncio.create_task(q.put('b'))
        self.assertEqual(await q.get(), 'a')
        self.assertEqual(await q.get(), 'b')

    async def test_why_are_getters_waiting(self):
        async def consumer(queue, num_expected):
            for _ in range(num_expected):
                await queue.get()

        async def producer(queue, num_items):
            for i in range(num_items):
                await queue.put(i)

        producer_num_items = 5

        q = asyncio.Queue(1)
        async with asyncio.TaskGroup() as tg:
            tg.create_task(producer(q, producer_num_items))
            tg.create_task(consumer(q, producer_num_items))

    async def test_cancelled_getters_not_being_held_in_self_getters(self):
        queue = asyncio.Queue(maxsize=5)

        with self.assertRaises(TimeoutError):
            await asyncio.wait_for(queue.get(), 0.1)

        self.assertEqual(len(queue._getters), 0)


class QueuePutTests(unittest.IsolatedAsyncioTestCase):

    async def test_blocking_put(self):
        q = asyncio.Queue()

        # No maxsize, won't block.
        await q.put(1)
        self.assertEqual(1, await q.get())

    async def test_blocking_put_wait(self):
        q = asyncio.Queue(maxsize=1)
        started = asyncio.Event()
        finished = False

        async def queue_put():
            nonlocal finished
            started.set()
            await q.put(1)
            await q.put(2)
            finished = True

        loop = asyncio.get_running_loop()
        loop.call_later(0.01, q.get_nowait)
        queue_put_task = asyncio.create_task(queue_put())
        await started.wait()
        self.assertFalse(finished)
        await queue_put_task
        self.assertTrue(finished)

    def test_nonblocking_put(self):
        q = asyncio.Queue()
        q.put_nowait(1)
        self.assertEqual(1, q.get_nowait())

    async def test_get_cancel_drop_one_pending_reader(self):
        q = asyncio.Queue()

        reader = asyncio.create_task(q.get())

        await asyncio.sleep(0)

        q.put_nowait(1)
        q.put_nowait(2)
        reader.cancel()

        try:
            await reader
        except asyncio.CancelledError:
            # try again
            reader = asyncio.create_task(q.get())
            await reader

        result = reader.result()
        # if we get 2, it means 1 got dropped!
        self.assertEqual(1, result)

    async def test_get_cancel_drop_many_pending_readers(self):
        q = asyncio.Queue()

        async with asyncio.TaskGroup() as tg:
            reader1 = tg.create_task(q.get())
            reader2 = tg.create_task(q.get())
            reader3 = tg.create_task(q.get())

            await asyncio.sleep(0)

            q.put_nowait(1)
            q.put_nowait(2)
            reader1.cancel()

            with self.assertRaises(asyncio.CancelledError):
                await reader1

            await reader3

        # It is undefined in which order concurrent readers receive results.
        self.assertEqual({reader2.result(), reader3.result()}, {1, 2})

    async def test_put_cancel_drop(self):
        q = asyncio.Queue(1)

        q.put_nowait(1)

        # putting a second item in the queue has to block (qsize=1)
        writer = asyncio.create_task(q.put(2))
        await asyncio.sleep(0)

        value1 = q.get_nowait()
        self.assertEqual(value1, 1)

        writer.cancel()
        try:
            await writer
        except asyncio.CancelledError:
            # try again
            writer = asyncio.create_task(q.put(2))
            await writer

        value2 = q.get_nowait()
        self.assertEqual(value2, 2)
        self.assertEqual(q.qsize(), 0)

    def test_nonblocking_put_exception(self):
        q = asyncio.Queue(maxsize=1, )
        q.put_nowait(1)
        self.assertRaises(asyncio.QueueFull, q.put_nowait, 2)

    async def test_float_maxsize(self):
        q = asyncio.Queue(maxsize=1.3, )
        q.put_nowait(1)
        q.put_nowait(2)
        self.assertTrue(q.full())
        self.assertRaises(asyncio.QueueFull, q.put_nowait, 3)

        q = asyncio.Queue(maxsize=1.3, )

        await q.put(1)
        await q.put(2)
        self.assertTrue(q.full())

    async def test_put_cancelled(self):
        q = asyncio.Queue()

        async def queue_put():
            await q.put(1)
            return True

        t = asyncio.create_task(queue_put())

        self.assertEqual(1, await q.get())
        self.assertTrue(t.done())
        self.assertTrue(t.result())

    async def test_put_cancelled_race(self):
        q = asyncio.Queue(maxsize=1)

        put_a = asyncio.create_task(q.put('a'))
        put_b = asyncio.create_task(q.put('b'))
        put_c = asyncio.create_task(q.put('X'))

        await asyncio.sleep(0)
        self.assertTrue(put_a.done())
        self.assertFalse(put_b.done())

        put_c.cancel()
        await asyncio.sleep(0)
        self.assertTrue(put_c.done())
        self.assertEqual(q.get_nowait(), 'a')
        await asyncio.sleep(0)
        self.assertEqual(q.get_nowait(), 'b')

        await put_b

    async def test_put_with_waiting_getters(self):
        q = asyncio.Queue()
        t = asyncio.create_task(q.get())
        await asyncio.sleep(0)
        await q.put('a')
        self.assertEqual(await t, 'a')

    async def test_why_are_putters_waiting(self):
        queue = asyncio.Queue(2)

        async def putter(item):
            await queue.put(item)

        async def getter():
            await asyncio.sleep(0)
            num = queue.qsize()
            for _ in range(num):
                queue.get_nowait()

        async with asyncio.TaskGroup() as tg:
            tg.create_task(getter())
            tg.create_task(putter(0))
            tg.create_task(putter(1))
            tg.create_task(putter(2))
            tg.create_task(putter(3))

    async def test_cancelled_puts_not_being_held_in_self_putters(self):
        # Full queue.
        queue = asyncio.Queue(maxsize=1)
        queue.put_nowait(1)

        # Task waiting for space to put an item in the queue.
        put_task = asyncio.create_task(queue.put(1))
        await asyncio.sleep(0)

        # Check that the putter is correctly removed from queue._putters when
        # the task is canceled.
        self.assertEqual(len(queue._putters), 1)
        put_task.cancel()
        with self.assertRaises(asyncio.CancelledError):
            await put_task
        self.assertEqual(len(queue._putters), 0)

    async def test_cancelled_put_silence_value_error_exception(self):
        # Full Queue.
        queue = asyncio.Queue(1)
        queue.put_nowait(1)

        # Task waiting for space to put a item in the queue.
        put_task = asyncio.create_task(queue.put(1))
        await asyncio.sleep(0)

        # get_nowait() remove the future of put_task from queue._putters.
        queue.get_nowait()
        # When canceled, queue.put is going to remove its future from
        # self._putters but it was removed previously by queue.get_nowait().
        put_task.cancel()

        # The ValueError exception triggered by queue._putters.remove(putter)
        # inside queue.put should be silenced.
        # If the ValueError is silenced we should catch a CancelledError.
        with self.assertRaises(asyncio.CancelledError):
            await put_task


class LifoQueueTests(unittest.IsolatedAsyncioTestCase):

    async def test_order(self):
        q = asyncio.LifoQueue()
        for i in [1, 3, 2]:
            await q.put(i)

        items = [await q.get() for _ in range(3)]
        self.assertEqual([2, 3, 1], items)


class PriorityQueueTests(unittest.IsolatedAsyncioTestCase):

    async def test_order(self):
        q = asyncio.PriorityQueue()
        for i in [1, 3, 2]:
            await q.put(i)

        items = [await q.get() for _ in range(3)]
        self.assertEqual([1, 2, 3], items)


class _QueueJoinTestMixin:

    q_class = None

    def test_task_done_underflow(self):
        q = self.q_class()
        self.assertRaises(ValueError, q.task_done)

    async def test_task_done(self):
        q = self.q_class()
        for i in range(100):
            q.put_nowait(i)

        accumulator = 0

        # Two workers get items from the queue and call task_done after each.
        # Join the queue and assert all items have been processed.
        running = True

        async def worker():
            nonlocal accumulator

            while running:
                item = await q.get()
                accumulator += item
                q.task_done()

        async with asyncio.TaskGroup() as tg:
            tasks = [tg.create_task(worker())
                     for index in range(2)]

            await q.join()
            self.assertEqual(sum(range(100)), accumulator)

            # close running generators
            running = False
            for i in range(len(tasks)):
                q.put_nowait(0)

    async def test_join_empty_queue(self):
        q = self.q_class()

        # Test that a queue join()s successfully, and before anything else
        # (done twice for insurance).

        await q.join()
        await q.join()

    async def test_format(self):
        q = self.q_class()
        self.assertEqual(q._format(), 'maxsize=0')

        q._unfinished_tasks = 2
        self.assertEqual(q._format(), 'maxsize=0 tasks=2')


class QueueJoinTests(_QueueJoinTestMixin, unittest.IsolatedAsyncioTestCase):
    q_class = asyncio.Queue


class LifoQueueJoinTests(_QueueJoinTestMixin, unittest.IsolatedAsyncioTestCase):
    q_class = asyncio.LifoQueue


class PriorityQueueJoinTests(_QueueJoinTestMixin, unittest.IsolatedAsyncioTestCase):
    q_class = asyncio.PriorityQueue


class _QueueShutdownTestMixin:
    q_class = None

    def assertRaisesShutdown(self, msg="Didn't appear to shut-down queue"):
        return self.assertRaises(asyncio.QueueShutDown, msg=msg)

    async def test_format(self):
        q = self.q_class()
        q.shutdown()
        self.assertEqual(q._format(), 'maxsize=0 shutdown')

    async def test_shutdown_empty(self):
        # Test shutting down an empty queue

        # Setup empty queue, and join() and get() tasks
        q = self.q_class()
        loop = asyncio.get_running_loop()
        get_task = loop.create_task(q.get())
        await asyncio.sleep(0)  # want get task pending before shutdown

        # Perform shut-down
        q.shutdown(immediate=False)  # unfinished tasks: 0 -> 0

        self.assertEqual(q.qsize(), 0)

        # Ensure join() task successfully finishes
        await q.join()

        # Ensure get() task is finished, and raised ShutDown
        await asyncio.sleep(0)
        self.assertTrue(get_task.done())
        with self.assertRaisesShutdown():
            await get_task

        # Ensure put() and get() raise ShutDown
        with self.assertRaisesShutdown():
            await q.put("data")
        with self.assertRaisesShutdown():
            q.put_nowait("data")

        with self.assertRaisesShutdown():
            await q.get()
        with self.assertRaisesShutdown():
            q.get_nowait()

    async def test_shutdown_nonempty(self):
        # Test shutting down a non-empty queue

        # Setup full queue with 1 item, and join() and put() tasks
        q = self.q_class(maxsize=1)
        loop = asyncio.get_running_loop()

        q.put_nowait("data")
        join_task = loop.create_task(q.join())
        put_task = loop.create_task(q.put("data2"))

        # Ensure put() task is not finished
        await asyncio.sleep(0)
        self.assertFalse(put_task.done())

        # Perform shut-down
        q.shutdown(immediate=False)  # unfinished tasks: 1 -> 1

        self.assertEqual(q.qsize(), 1)

        # Ensure put() task is finished, and raised ShutDown
        await asyncio.sleep(0)
        self.assertTrue(put_task.done())
        with self.assertRaisesShutdown():
            await put_task

        # Ensure get() succeeds on enqueued item
        self.assertEqual(await q.get(), "data")

        # Ensure join() task is not finished
        await asyncio.sleep(0)
        self.assertFalse(join_task.done())

        # Ensure put() and get() raise ShutDown
        with self.assertRaisesShutdown():
            await q.put("data")
        with self.assertRaisesShutdown():
            q.put_nowait("data")

        with self.assertRaisesShutdown():
            await q.get()
        with self.assertRaisesShutdown():
            q.get_nowait()

        # Ensure there is 1 unfinished task, and join() task succeeds
        q.task_done()

        await asyncio.sleep(0)
        self.assertTrue(join_task.done())
        await join_task

        with self.assertRaises(
            ValueError, msg="Didn't appear to mark all tasks done"
        ):
            q.task_done()

    async def test_shutdown_immediate(self):
        # Test immediately shutting down a queue

        # Setup queue with 1 item, and a join() task
        q = self.q_class()
        loop = asyncio.get_running_loop()
        q.put_nowait("data")
        join_task = loop.create_task(q.join())

        # Perform shut-down
        q.shutdown(immediate=True)  # unfinished tasks: 1 -> 0

        self.assertEqual(q.qsize(), 0)

        # Ensure join() task has successfully finished
        await asyncio.sleep(0)
        self.assertTrue(join_task.done())
        await join_task

        # Ensure put() and get() raise ShutDown
        with self.assertRaisesShutdown():
            await q.put("data")
        with self.assertRaisesShutdown():
            q.put_nowait("data")

        with self.assertRaisesShutdown():
            await q.get()
        with self.assertRaisesShutdown():
            q.get_nowait()

        # Ensure there are no unfinished tasks
        with self.assertRaises(
            ValueError, msg="Didn't appear to mark all tasks done"
        ):
            q.task_done()

    async def test_shutdown_immediate_with_unfinished(self):
        # Test immediately shutting down a queue with unfinished tasks

        # Setup queue with 2 items (1 retrieved), and a join() task
        q = self.q_class()
        loop = asyncio.get_running_loop()
        q.put_nowait("data")
        q.put_nowait("data")
        join_task = loop.create_task(q.join())
        self.assertEqual(await q.get(), "data")

        # Perform shut-down
        q.shutdown(immediate=True)  # unfinished tasks: 2 -> 1

        self.assertEqual(q.qsize(), 0)

        # Ensure join() task is not finished
        await asyncio.sleep(0)
        self.assertFalse(join_task.done())

        # Ensure put() and get() raise ShutDown
        with self.assertRaisesShutdown():
            await q.put("data")
        with self.assertRaisesShutdown():
            q.put_nowait("data")

        with self.assertRaisesShutdown():
            await q.get()
        with self.assertRaisesShutdown():
            q.get_nowait()

        # Ensure there is 1 unfinished task
        q.task_done()
        with self.assertRaises(
            ValueError, msg="Didn't appear to mark all tasks done"
        ):
            q.task_done()

        # Ensure join() task has successfully finished
        await asyncio.sleep(0)
        self.assertTrue(join_task.done())
        await join_task


class QueueShutdownTests(
    _QueueShutdownTestMixin, unittest.IsolatedAsyncioTestCase
):
    q_class = asyncio.Queue


class LifoQueueShutdownTests(
    _QueueShutdownTestMixin, unittest.IsolatedAsyncioTestCase
):
    q_class = asyncio.LifoQueue


class PriorityQueueShutdownTests(
    _QueueShutdownTestMixin, unittest.IsolatedAsyncioTestCase
):
    q_class = asyncio.PriorityQueue


if __name__ == '__main__':
    unittest.main()