tests for queues py import asyncio import unittest from types import genericalias def teardownmodule asyncio seteventlooppolicynone class queuebasictestsunittest isolatedasynciotestcase async def testreprorstrself fn expectid q asyncio queue self asserttruefnq startswith queue fnq idispresent hexidq in fnq self assertequalexpectid idispresent getters q asyncio queue async with asyncio taskgroup as tg start a task that waits to get getter tg createtaskq get let it start waiting await asyncio sleep0 self asserttrue getters1 in fnq resume q get coroutine to finish generator q putnowait0 self assertequal0 await getter putters q asyncio queuemaxsize1 async with asyncio taskgroup as tg q putnowait1 start a task that waits to put putter tg createtaskq put2 let it start waiting await asyncio sleep0 self asserttrue putters1 in fnq resume q put coroutine to finish generator q getnowait self asserttrueputter done q asyncio queue q putnowait1 self asserttrue queue1 in fnq async def testreprself await self testreprorstrrepr true async def teststrself await self testreprorstrstr false def testgenericaliasself q asyncio queueint self assertequalq args int self assertisinstanceq genericalias async def testemptyself q asyncio queue self asserttrueq empty await q put1 self assertfalseq empty self assertequal1 await q get self asserttrueq empty async def testfullself q asyncio queue self assertfalseq full q asyncio queuemaxsize1 await q put1 self asserttrueq full async def testorderself q asyncio queue for i in 1 3 2 await q puti items await q get for in range3 self assertequal1 3 2 items async def testmaxsizeself q asyncio queuemaxsize2 self assertequal2 q maxsize havebeenput async def putter for i in range3 await q puti havebeenput appendi return true t asyncio createtaskputter for i in range2 await asyncio sleep0 the putter is blocked after putting two items self assertequal0 1 havebeenput self assertequal0 await q get let the putter resume and put last item await asyncio sleep0 self assertequal0 1 2 havebeenput self assertequal1 await q get self assertequal2 await q get self asserttruet done self asserttruet result class queuegettestsunittest isolatedasynciotestcase async def testblockinggetself q asyncio queue q putnowait1 self assertequal1 await q get async def testgetwithputtersself loop asyncio getrunningloop q asyncio queue1 await q put1 waiter loop createfuture q putters appendwaiter self assertequal1 await q get self asserttruewaiter done self assertisnonewaiter result async def testblockinggetwaitself loop asyncio getrunningloop q asyncio queue started asyncio event finished false async def queueget nonlocal finished started set res await q get finished true return res queuegettask asyncio createtaskqueueget await started wait self assertfalsefinished loop calllater0 01 q putnowait 1 res await queuegettask self asserttruefinished self assertequal1 res def testnonblockinggetself q asyncio queue q putnowait1 self assertequal1 q getnowait def testnonblockinggetexceptionself q asyncio queue self assertraisesasyncio queueempty q getnowait async def testgetcancelledraceself q asyncio queue t1 asyncio createtaskq get t2 asyncio createtaskq get await asyncio sleep0 t1 cancel await asyncio sleep0 self asserttruet1 done await q put a await asyncio sleep0 self assertequal a await t2 async def testgetwithwaitingputtersself q asyncio queuemaxsize1 asyncio createtaskq put a asyncio createtaskq put b self assertequalawait q get a self assertequalawait q get b async def testwhyaregetterswaitingself async def consumerqueue numexpected for in rangenumexpected await queue get async def producerqueue numitems for i in rangenumitems await queue puti producernumitems 5 q asyncio queue1 async with asyncio taskgroup as tg tg createtaskproducerq producernumitems tg createtaskconsumerq producernumitems async def testcancelledgettersnotbeingheldinselfgettersself queue asyncio queuemaxsize5 with self assertraisestimeouterror await asyncio waitforqueue get 0 1 self assertequallenqueue getters 0 class queueputtestsunittest isolatedasynciotestcase async def testblockingputself q asyncio queue no maxsize won t block await q put1 self assertequal1 await q get async def testblockingputwaitself q asyncio queuemaxsize1 started asyncio event finished false async def queueput nonlocal finished started set await q put1 await q put2 finished true loop asyncio getrunningloop loop calllater0 01 q getnowait queueputtask asyncio createtaskqueueput await started wait self assertfalsefinished await queueputtask self asserttruefinished def testnonblockingputself q asyncio queue q putnowait1 self assertequal1 q getnowait async def testgetcanceldroponependingreaderself q asyncio queue reader asyncio createtaskq get await asyncio sleep0 q putnowait1 q putnowait2 reader cancel try await reader except asyncio cancellederror try again reader asyncio createtaskq get await reader result reader result if we get 2 it means 1 got dropped self assertequal1 result async def testgetcanceldropmanypendingreadersself q asyncio queue async with asyncio taskgroup as tg reader1 tg createtaskq get reader2 tg createtaskq get reader3 tg createtaskq get await asyncio sleep0 q putnowait1 q putnowait2 reader1 cancel with self assertraisesasyncio cancellederror await reader1 await reader3 it is undefined in which order concurrent readers receive results self assertequalreader2 result reader3 result 1 2 async def testputcanceldropself q asyncio queue1 q putnowait1 putting a second item in the queue has to block qsize1 writer asyncio createtaskq put2 await asyncio sleep0 value1 q getnowait self assertequalvalue1 1 writer cancel try await writer except asyncio cancellederror try again writer asyncio createtaskq put2 await writer value2 q getnowait self assertequalvalue2 2 self assertequalq qsize 0 def testnonblockingputexceptionself q asyncio queuemaxsize1 q putnowait1 self assertraisesasyncio queuefull q putnowait 2 async def testfloatmaxsizeself q asyncio queuemaxsize1 3 q putnowait1 q putnowait2 self asserttrueq full self assertraisesasyncio queuefull q putnowait 3 q asyncio queuemaxsize1 3 await q put1 await q put2 self asserttrueq full async def testputcancelledself q asyncio queue async def queueput await q put1 return true t asyncio createtaskqueueput self assertequal1 await q get self asserttruet done self asserttruet result async def testputcancelledraceself q asyncio queuemaxsize1 puta asyncio createtaskq put a putb asyncio createtaskq put b putc asyncio createtaskq put x await asyncio sleep0 self asserttrueputa done self assertfalseputb done putc cancel await asyncio sleep0 self asserttrueputc done self assertequalq getnowait a await asyncio sleep0 self assertequalq getnowait b await putb async def testputwithwaitinggettersself q asyncio queue t asyncio createtaskq get await asyncio sleep0 await q put a self assertequalawait t a async def testwhyareputterswaitingself queue asyncio queue2 async def putteritem await queue putitem async def getter await asyncio sleep0 num queue qsize for in rangenum queue getnowait async with asyncio taskgroup as tg tg createtaskgetter tg createtaskputter0 tg createtaskputter1 tg createtaskputter2 tg createtaskputter3 async def testcancelledputsnotbeingheldinselfputtersself full queue queue asyncio queuemaxsize1 queue putnowait1 task waiting for space to put an item in the queue puttask asyncio createtaskqueue put1 await asyncio sleep0 check that the putter is correctly removed from queue putters when the task is canceled self assertequallenqueue putters 1 puttask cancel with self assertraisesasyncio cancellederror await puttask self assertequallenqueue putters 0 async def testcancelledputsilencevalueerrorexceptionself full queue queue asyncio queue1 queue putnowait1 task waiting for space to put a item in the queue puttask asyncio createtaskqueue put1 await asyncio sleep0 getnowait remove the future of puttask from queue putters queue getnowait when canceled queue put is going to remove its future from self putters but it was removed previously by queue getnowait puttask cancel the valueerror exception triggered by queue putters removeputter inside queue put should be silenced if the valueerror is silenced we should catch a cancellederror with self assertraisesasyncio cancellederror await puttask class lifoqueuetestsunittest isolatedasynciotestcase async def testorderself q asyncio lifoqueue for i in 1 3 2 await q puti items await q get for in range3 self assertequal2 3 1 items class priorityqueuetestsunittest isolatedasynciotestcase async def testorderself q asyncio priorityqueue for i in 1 3 2 await q puti items await q get for in range3 self assertequal1 2 3 items class queuejointestmixin qclass none def testtaskdoneunderflowself q self qclass self assertraisesvalueerror q taskdone async def testtaskdoneself q self qclass for i in range100 q putnowaiti accumulator 0 two workers get items from the queue and call taskdone 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 taskdone async with asyncio taskgroup as tg tasks tg createtaskworker for index in range2 await q join self assertequalsumrange100 accumulator close running generators running false for i in rangelentasks q putnowait0 async def testjoinemptyqueueself q self qclass test that a queue joins successfully and before anything else done twice for insurance await q join await q join async def testformatself q self qclass self assertequalq format maxsize0 q unfinishedtasks 2 self assertequalq format maxsize0 tasks2 class queuejointestsqueuejointestmixin unittest isolatedasynciotestcase qclass asyncio queue class lifoqueuejointestsqueuejointestmixin unittest isolatedasynciotestcase qclass asyncio lifoqueue class priorityqueuejointestsqueuejointestmixin unittest isolatedasynciotestcase qclass asyncio priorityqueue if name main unittest main 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 getters start a task that waits to get let it start waiting resume q get coroutine to finish generator putters start a task that waits to put let it start waiting resume q put coroutine to finish generator the putter is blocked after putting two items let the putter resume and put last item no maxsize won t block try again if we get 2 it means 1 got dropped it is undefined in which order concurrent readers receive results putting a second item in the queue has to block qsize 1 try again full queue task waiting for space to put an item in the queue check that the putter is correctly removed from queue _putters when the task is canceled full queue task waiting for space to put a item in the queue get_nowait remove the future of put_task from queue _putters when canceled queue put is going to remove its future from self _putters but it was removed previously by queue get_nowait 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 two workers get items from the queue and call task_done after each join the queue and assert all items have been processed close running generators test that a queue join s successfully and before anything else done twice for insurance
import asyncio
import unittest
from types import GenericAlias


def tearDownModule():
    asyncio.set_event_loop_policy(None)


class QueueBasicTests(unittest.IsolatedAsyncioTestCase):

    async def _test_repr_or_str(self, fn, expect_id):
        
        q = asyncio.Queue()
        self.assertTrue(fn(q).startswith('<Queue'), fn(q))
        id_is_present = hex(id(q)) in fn(q)
        self.assertEqual(expect_id, id_is_present)

        
        q = asyncio.Queue()
        async with asyncio.TaskGroup() as tg:
            
            getter = tg.create_task(q.get())
            
            await asyncio.sleep(0)
            self.assertTrue('_getters[1]' in fn(q))
            
            q.put_nowait(0)

        self.assertEqual(0, await getter)

        
        q = asyncio.Queue(maxsize=1)
        async with asyncio.TaskGroup() as tg:
            q.put_nowait(1)
            
            putter = tg.create_task(q.put(2))
            
            await asyncio.sleep(0)
            self.assertTrue('_putters[1]' in fn(q))
            
            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)

        
        self.assertEqual([0, 1], have_been_put)
        self.assertEqual(0, await q.get())

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

        
        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:
            
            reader = asyncio.create_task(q.get())
            await reader

        result = reader.result()
        
        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

        
        self.assertEqual({reader2.result(), reader3.result()}, {1, 2})

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

        q.put_nowait(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:
            
            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):
        
        queue = asyncio.Queue(maxsize=1)
        queue.put_nowait(1)

        
        put_task = asyncio.create_task(queue.put(1))
        await asyncio.sleep(0)

        
        
        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):
        
        queue = asyncio.Queue(1)
        queue.put_nowait(1)

        
        put_task = asyncio.create_task(queue.put(1))
        await asyncio.sleep(0)

        
        queue.get_nowait()
        
        
        put_task.cancel()

        
        
        
        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

        
        
        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)

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

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

        
        

        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


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