From 586023f72232cf077ec719442ecf91c49411d0aa Mon Sep 17 00:00:00 2001 From: Daniil Fajnberg Date: Mon, 7 Feb 2022 23:41:52 +0100 Subject: [PATCH] fixed unittests; minor changes --- setup.cfg | 2 +- src/asyncio_taskpool/pool.py | 10 ++++++++-- tests/test_pool.py | 21 +++++++++------------ 3 files changed, 18 insertions(+), 15 deletions(-) diff --git a/setup.cfg b/setup.cfg index a60e5f0..873a05b 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = asyncio-taskpool -version = 0.1.1 +version = 0.1.2 author = Daniil Fajnberg author_email = mail@daniil.fajnberg.de description = Dynamically manage pools of asyncio tasks diff --git a/src/asyncio_taskpool/pool.py b/src/asyncio_taskpool/pool.py index f6c75cc..e0d4f56 100644 --- a/src/asyncio_taskpool/pool.py +++ b/src/asyncio_taskpool/pool.py @@ -297,7 +297,8 @@ class BaseTaskPool: return_exceptions (optional): Passed directly into `gather`. """ results = await gather(*self._ended.values(), *self._cancelled.values(), return_exceptions=return_exceptions) - self._ended = self._cancelled = {} + self._ended.clear() + self._cancelled.clear() if self._interrupt_flag.is_set(): self._interrupt_flag.clear() return results @@ -333,7 +334,10 @@ class BaseTaskPool: await gather(*self._before_gathering) results = await gather(*self._ended.values(), *self._cancelled.values(), *self._running.values(), return_exceptions=return_exceptions) - self._ended = self._cancelled = self._running = {} + self._ended.clear() + self._cancelled.clear() + self._running.clear() + self._before_gathering.clear() if self._interrupt_flag.is_set(): self._interrupt_flag.clear() return results @@ -498,6 +502,8 @@ class TaskPool(BaseTaskPool): end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None: """ Creates coroutines with arguments from a supplied iterable and runs them as new tasks in the pool in batches. + TODO: If task groups are implemented, consider adding all tasks from one call of this method to the same group + and referring to "group size" rather than chunk/batch size. Each coroutine looks like `func(arg)`, `func(*arg)`, or `func(**arg)`, `arg` being an element from the iterable. This method blocks, **only if** there is not enough room in the pool for the first batch of new tasks. diff --git a/tests/test_pool.py b/tests/test_pool.py index 9bfa20e..110f74f 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -67,8 +67,7 @@ class BaseTaskPoolTestCase(IsolatedAsyncioTestCase): self.assertEqual(0, self.task_pool._num_ended) self.assertEqual(self.mock_idx, self.task_pool._idx) self.assertEqual(self.test_pool_name, self.task_pool._name) - self.assertIsInstance(self.task_pool._all_tasks_known_flag, asyncio.locks.Event) - self.assertTrue(self.task_pool._all_tasks_known_flag.is_set()) + self.assertListEqual(self.task_pool._before_gathering, EMPTY_LIST) self.assertIsInstance(self.task_pool._interrupt_flag, asyncio.locks.Event) self.assertFalse(self.task_pool._interrupt_flag.is_set()) self.mock__add_pool.assert_called_once_with(self.task_pool) @@ -342,11 +341,11 @@ class BaseTaskPoolTestCase(IsolatedAsyncioTestCase): self.assertFalse(self.task_pool._open) async def test_gather(self): - mock_wait = AsyncMock() - self.task_pool._all_tasks_known_flag = MagicMock(wait=mock_wait) test_exception = TestException() mock_ended_func, mock_cancelled_func = AsyncMock(return_value=FOO), AsyncMock(side_effect=test_exception) mock_running_func = AsyncMock(return_value=BAR) + mock_queue_join = AsyncMock() + self.task_pool._before_gathering = before_gather = [mock_queue_join()] self.task_pool._ended = ended = {123: mock_ended_func()} self.task_pool._cancelled = cancelled = {456: mock_cancelled_func()} self.task_pool._running = running = {789: mock_running_func()} @@ -358,25 +357,23 @@ class BaseTaskPoolTestCase(IsolatedAsyncioTestCase): self.assertDictEqual(self.task_pool._ended, ended) self.assertDictEqual(self.task_pool._cancelled, cancelled) self.assertDictEqual(self.task_pool._running, running) + self.assertListEqual(self.task_pool._before_gathering, before_gather) self.assertTrue(self.task_pool._interrupt_flag.is_set()) - mock_wait.assert_not_awaited() self.task_pool._open = False - def check_assertions() -> None: + def check_assertions(output) -> None: self.assertListEqual([FOO, test_exception, BAR], output) self.assertDictEqual(self.task_pool._ended, EMPTY_DICT) self.assertDictEqual(self.task_pool._cancelled, EMPTY_DICT) self.assertDictEqual(self.task_pool._running, EMPTY_DICT) + self.assertListEqual(self.task_pool._before_gathering, EMPTY_LIST) self.assertFalse(self.task_pool._interrupt_flag.is_set()) - mock_wait.assert_awaited_once_with() - output = await self.task_pool.gather(return_exceptions=True) - check_assertions() - mock_wait.reset_mock() + check_assertions(await self.task_pool.gather(return_exceptions=True)) + self.task_pool._before_gathering = [mock_queue_join()] self.task_pool._ended = {123: mock_ended_func()} self.task_pool._cancelled = {456: mock_cancelled_func()} self.task_pool._running = {789: mock_running_func()} - output = await self.task_pool.gather(return_exceptions=True) - check_assertions() + check_assertions(await self.task_pool.gather(return_exceptions=True))