From a3b0ffb72a1e72fc3be0e96c7407ddff2ade67c9 Mon Sep 17 00:00:00 2001 From: Numerlor Date: Tue, 14 Jun 2022 00:37:07 +0200 Subject: Add decorator to block duplicate command invocations in a channel --- tests/botcore/utils/test_cooldown.py | 48 ++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/botcore/utils/test_cooldown.py (limited to 'tests') diff --git a/tests/botcore/utils/test_cooldown.py b/tests/botcore/utils/test_cooldown.py new file mode 100644 index 00000000..e7fe0f59 --- /dev/null +++ b/tests/botcore/utils/test_cooldown.py @@ -0,0 +1,48 @@ +import unittest +from unittest.mock import patch + +from botcore.utils.cooldown import _ArgsTuple, _CommandCooldownManager + + +def create_argument_tuple(*args, **kwargs) -> _ArgsTuple: + return (*args, *kwargs.items()) + + +class CommandCooldownManagerTests(unittest.IsolatedAsyncioTestCase): + test_call_args = ( + create_argument_tuple(0), + create_argument_tuple(a=0), + create_argument_tuple([]), + create_argument_tuple(a=[]), + create_argument_tuple(1, 2, 3, a=4, b=5, c=6), + create_argument_tuple([1], [2], [3], a=[4], b=[5], c=[6]), + create_argument_tuple([1], 2, [3], a=4, b=[5], c=6), + ) + + async def asyncSetUp(self): + self.cooldown_manager = _CommandCooldownManager(cooldown_duration=5) + + def test_no_cooldown_on_unset(self): + for call_args in self.test_call_args: + with self.subTest(arguments_tuple=call_args, channel=0): + self.assertFalse(self.cooldown_manager.is_on_cooldown(0, call_args)) + + for call_args in self.test_call_args: + with self.subTest(arguments_tuple=call_args, channel=1): + self.assertFalse(self.cooldown_manager.is_on_cooldown(1, call_args)) + + @patch("time.monotonic") + def test_cooldown_is_set(self, monotonic): + monotonic.side_effect = lambda: 0 + for call_args in self.test_call_args: + with self.subTest(arguments_tuple=call_args): + self.cooldown_manager.set_cooldown(0, call_args) + self.assertTrue(self.cooldown_manager.is_on_cooldown(0, call_args)) + + @patch("time.monotonic") + def test_cooldown_expires(self, monotonic): + for call_args in self.test_call_args: + monotonic.side_effect = (0, 1000) + with self.subTest(arguments_tuple=call_args): + self.cooldown_manager.set_cooldown(0, call_args) + self.assertFalse(self.cooldown_manager.is_on_cooldown(0, call_args)) -- cgit v1.2.3 From 805c60437e50e2153d61e484872b207affd8db1f Mon Sep 17 00:00:00 2001 From: Numerlor Date: Tue, 21 Jun 2022 19:20:03 +0200 Subject: stop cleanup task when manager is destroyed --- botcore/utils/cooldown.py | 8 +++++++- tests/botcore/utils/test_cooldown.py | 5 +++-- 2 files changed, 10 insertions(+), 3 deletions(-) (limited to 'tests') diff --git a/botcore/utils/cooldown.py b/botcore/utils/cooldown.py index 59b0722e..099edba0 100644 --- a/botcore/utils/cooldown.py +++ b/botcore/utils/cooldown.py @@ -6,6 +6,7 @@ import asyncio import random import time import typing +import weakref from collections.abc import Awaitable, Hashable, Iterable from contextlib import suppress from dataclasses import dataclass @@ -106,6 +107,7 @@ class _CommandCooldownManager: self._periodical_cleanup(random.uniform(0, 10)), name="CooldownManager cleanup", ) + weakref.finalize(self, self.cleanup_task.cancel) def set_cooldown(self, channel: Hashable, call_arguments: Iterable[object]) -> None: """Set `call_arguments` arguments on cooldown in `channel`.""" @@ -145,11 +147,15 @@ class _CommandCooldownManager: Delete stale items every hour after waiting for `initial_delay`. The `initial_delay` ensures cleanups are not running for every command at the same time. + A strong reference to self is only kept while cleanup is running. """ + weak_self = weakref.ref(self) + del self + await asyncio.sleep(initial_delay) while True: await asyncio.sleep(60 * 60) - self._delete_stale_items() + weak_self()._delete_stale_items() def _delete_stale_items(self) -> None: """Remove expired items from internal collections.""" diff --git a/tests/botcore/utils/test_cooldown.py b/tests/botcore/utils/test_cooldown.py index e7fe0f59..87c433ce 100644 --- a/tests/botcore/utils/test_cooldown.py +++ b/tests/botcore/utils/test_cooldown.py @@ -1,10 +1,11 @@ import unittest +from collections.abc import Iterable from unittest.mock import patch -from botcore.utils.cooldown import _ArgsTuple, _CommandCooldownManager +from botcore.utils.cooldown import _CommandCooldownManager -def create_argument_tuple(*args, **kwargs) -> _ArgsTuple: +def create_argument_tuple(*args, **kwargs) -> Iterable[object]: return (*args, *kwargs.items()) -- cgit v1.2.3 From 67003153c718925844447127f291501adddb49c0 Mon Sep 17 00:00:00 2001 From: Numerlor Date: Tue, 21 Jun 2022 19:48:36 +0200 Subject: ensure tuples from pos arg and kwarg tuples are differentiated --- botcore/utils/cooldown.py | 8 +++++++- tests/botcore/utils/test_cooldown.py | 26 +++++++++++++------------- 2 files changed, 20 insertions(+), 14 deletions(-) (limited to 'tests') diff --git a/botcore/utils/cooldown.py b/botcore/utils/cooldown.py index 099edba0..b9149b48 100644 --- a/botcore/utils/cooldown.py +++ b/botcore/utils/cooldown.py @@ -20,6 +20,8 @@ from botcore.utils.function import command_wraps __all__ = ["CommandOnCooldown", "block_duplicate_invocations", "P", "R"] +_KEYWORD_SEP_SENTINEL = object() + _ArgsList = list[object] _HashableArgsTuple = tuple[Hashable, ...] @@ -172,6 +174,10 @@ class _CommandCooldownManager: self._cooldowns[key] = filtered_cooldowns +def _create_argument_tuple(*args: object, **kwargs: object) -> Iterable[object]: + return (*args, _KEYWORD_SEP_SENTINEL, *kwargs.items()) + + def block_duplicate_invocations( *, cooldown_duration: float = 5, send_notice: bool = False ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]: @@ -194,7 +200,7 @@ def block_duplicate_invocations( @command_wraps(func) async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: - arg_tuple = (*args[2:], *kwargs.items()) # skip self and ctx from the command + arg_tuple = _create_argument_tuple(*args[2:], **kwargs) # skip self and ctx from the command ctx = typing.cast("Context[BotBase]", args[1]) channel = ctx.channel diff --git a/tests/botcore/utils/test_cooldown.py b/tests/botcore/utils/test_cooldown.py index 87c433ce..00e5a052 100644 --- a/tests/botcore/utils/test_cooldown.py +++ b/tests/botcore/utils/test_cooldown.py @@ -1,23 +1,18 @@ import unittest -from collections.abc import Iterable from unittest.mock import patch -from botcore.utils.cooldown import _CommandCooldownManager - - -def create_argument_tuple(*args, **kwargs) -> Iterable[object]: - return (*args, *kwargs.items()) +from botcore.utils.cooldown import _CommandCooldownManager, _create_argument_tuple class CommandCooldownManagerTests(unittest.IsolatedAsyncioTestCase): test_call_args = ( - create_argument_tuple(0), - create_argument_tuple(a=0), - create_argument_tuple([]), - create_argument_tuple(a=[]), - create_argument_tuple(1, 2, 3, a=4, b=5, c=6), - create_argument_tuple([1], [2], [3], a=[4], b=[5], c=[6]), - create_argument_tuple([1], 2, [3], a=4, b=[5], c=6), + _create_argument_tuple(0), + _create_argument_tuple(a=0), + _create_argument_tuple([]), + _create_argument_tuple(a=[]), + _create_argument_tuple(1, 2, 3, a=4, b=5, c=6), + _create_argument_tuple([1], [2], [3], a=[4], b=[5], c=[6]), + _create_argument_tuple([1], 2, [3], a=4, b=[5], c=6), ) async def asyncSetUp(self): @@ -47,3 +42,8 @@ class CommandCooldownManagerTests(unittest.IsolatedAsyncioTestCase): with self.subTest(arguments_tuple=call_args): self.cooldown_manager.set_cooldown(0, call_args) self.assertFalse(self.cooldown_manager.is_on_cooldown(0, call_args)) + + def test_keywords_and_tuples_differentiated(self): + self.cooldown_manager.set_cooldown(0, _create_argument_tuple(("a", 0))) + self.assertFalse(self.cooldown_manager.is_on_cooldown(0, _create_argument_tuple(a=0))) + self.assertTrue(self.cooldown_manager.is_on_cooldown(0, _create_argument_tuple(("a", 0)))) -- cgit v1.2.3