mirror of
				https://github.com/python/cpython.git
				synced 2025-10-23 01:43:53 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			725 lines
		
	
	
	
		
			21 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			725 lines
		
	
	
	
		
			21 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
| """Tests for queues.py"""
 | |
| 
 | |
| 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):
 | |
|         """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.assertTrue(fn(q).startswith('<Queue'), fn(q))
 | |
|         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()
 | 
