From 710a5ae21418fa89cb5dcc8a3a9cb7e299d83ff7 Mon Sep 17 00:00:00 2001 From: spacemanspiff2007 <10754716+spacemanspiff2007@users.noreply.github.com> Date: Wed, 20 Dec 2023 13:26:15 +0100 Subject: [PATCH] Update rate limiter --- docs/util.rst | 32 ++- src/HABApp/util/functions/min_max.py | 10 +- src/HABApp/util/rate_limiter/limiter.py | 95 ++++--- .../util/rate_limiter/limits/__init__.py | 3 + src/HABApp/util/rate_limiter/limits/base.py | 73 ++++++ .../util/rate_limiter/limits/fixed_window.py | 51 ++++ .../util/rate_limiter/limits/leaky_bucket.py | 50 ++++ src/HABApp/util/rate_limiter/parser.py | 27 ++ src/HABApp/util/rate_limiter/rate_limit.py | 73 ------ tests/test_utils/test_rate_limiter.py | 244 +++++++++++------- 10 files changed, 440 insertions(+), 218 deletions(-) create mode 100644 src/HABApp/util/rate_limiter/limits/__init__.py create mode 100644 src/HABApp/util/rate_limiter/limits/base.py create mode 100644 src/HABApp/util/rate_limiter/limits/fixed_window.py create mode 100644 src/HABApp/util/rate_limiter/limits/leaky_bucket.py create mode 100644 src/HABApp/util/rate_limiter/parser.py delete mode 100644 src/HABApp/util/rate_limiter/rate_limit.py diff --git a/docs/util.rst b/docs/util.rst index 53f4c663..560bbf43 100644 --- a/docs/util.rst +++ b/docs/util.rst @@ -81,6 +81,7 @@ A simple rate limiter implementation which can be used in rules. The limiter is not rule bound so the same limiter can be used in multiples files. It also works as expected across rule reloads. + Defining limits ^^^^^^^^^^^^^^^^^^ Limits can either be explicitly added or through a textual description. @@ -100,10 +101,10 @@ Examples: * ``300 / hour`` -Elastic expiry -^^^^^^^^^^^^^^^^^^ +Fixed window elastic expiry algorithm +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -The rate limiter implements a fixed window with elastic expiry. +This algorithm implements a fixed window with elastic expiry. That means if the limit is hit the interval time will be increased by the expiry time. For example ``3 per minute``: @@ -118,6 +119,14 @@ For example ``3 per minute``: will also get rejected and the intervall now goes from ``00:00:00`` - ``00:02:30``. +Leaky bucket algorithm +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The leaky bucket algorithm is based on the analogy of a bucket that leaks at a constant rate. +As long as the bucket is not full the hits will pass. If the bucket overflows the hits will get rejected. +Since the bucket leaks at a constant rate it will gradually get empty again thus allowing hits to pass again. + + Example ^^^^^^^^^^^^^^^^^^ @@ -128,10 +137,12 @@ Example # Create or get existing, name is case insensitive limiter = RateLimiter('MyRateLimiterName') - # define limits, duplicate limits will only be added once + # define limits, duplicate limits of the same algorithm will only be added once limiter.add_limit(5, 60) # add limits explicitly limiter.parse_limits('5 per minute').parse_limits('5 in 60s', '5/60seconds') # add limits through text + # add additional limit with leaky bucket algorithm + limiter.add_limit(10, 120, algorithm='leaky_bucket') # Test the limit without increasing the hits for _ in range(100): @@ -155,6 +166,19 @@ Documentation .. autoclass:: HABApp.util.rate_limiter.limiter.Limiter :members: + :inherited-members: + +.. autoclass:: HABApp.util.rate_limiter.limiter.LimiterInfo + :members: + :inherited-members: + +.. autoclass:: HABApp.util.rate_limiter.limiter.FixedWindowElasticExpiryLimitInfo + :members: + :inherited-members: + +.. autoclass:: HABApp.util.rate_limiter.limiter.LeakyBucketLimitInfo + :members: + :inherited-members: Statistics diff --git a/src/HABApp/util/functions/min_max.py b/src/HABApp/util/functions/min_max.py index d03284dc..92ab567b 100644 --- a/src/HABApp/util/functions/min_max.py +++ b/src/HABApp/util/functions/min_max.py @@ -2,8 +2,9 @@ from builtins import min as _min -def max(*args, default=None): - """Behaves like the built in max function but ignores any ``None`` values. e.g. ``max([1, None, 2]) == 2``. +# noinspection PyShadowingBuiltins +def max(*args, default=None): # noqa: A001 + """Behaves like the built-in max function but ignores any ``None`` values. e.g. ``max([1, None, 2]) == 2``. If the iterable is empty ``default`` will be returned. :param args: Single iterable or 1..n arguments @@ -16,8 +17,9 @@ def max(*args, default=None): ) -def min(*args, default=None): - """Behaves like the built in min function but ignores any ``None`` values. e.g. ``min([1, None, 2]) == 1``. +# noinspection PyShadowingBuiltins +def min(*args, default=None): # noqa: A001 + """Behaves like the built-in min function but ignores any ``None`` values. e.g. ``min([1, None, 2]) == 1``. If the iterable is empty ``default`` will be returned. :param args: Single iterable or 1..n arguments diff --git a/src/HABApp/util/rate_limiter/limiter.py b/src/HABApp/util/rate_limiter/limiter.py index 0c096091..cbed1ae3 100644 --- a/src/HABApp/util/rate_limiter/limiter.py +++ b/src/HABApp/util/rate_limiter/limiter.py @@ -1,82 +1,84 @@ -import re from dataclasses import dataclass -from time import monotonic -from typing import Final, List, Tuple +from typing import Final, List, Literal, Tuple, TypeAlias, Union -from .rate_limit import RateLimit, RateLimitInfo +from HABApp.core.const.const import StrEnum - -LIMIT_REGEX = re.compile( - r""" - \s* ([1-9][0-9]*) - \s* (/|per|in) - \s* ([1-9][0-9]*)? - \s* (s|sec|second|m|min|minute|h|hour|day|month|year)s? - \s*""", - re.IGNORECASE | re.VERBOSE, +from .limits import ( + BaseRateLimit, + FixedWindowElasticExpiryLimit, + FixedWindowElasticExpiryLimitInfo, + LeakyBucketLimit, + LeakyBucketLimitInfo, ) +from .parser import parse_limit -def parse_limit(text: str) -> Tuple[int, int]: - if not isinstance(text, str) or not (m := LIMIT_REGEX.fullmatch(text)): - msg = f'Invalid limit string: "{text:s}"' - raise ValueError(msg) - - count, per, factor, interval = m.groups() +class LimitTypeEnum(StrEnum): + LEAKY_BUCKET = 'leaky_bucket' + FIXED_WINDOW_ELASTIC_EXPIRY = 'fixed_window_elastic_expiry' - interval_secs = { - 's': 1, 'sec': 1, 'second': 1, 'm': 60, 'min': 60, 'minute': 60, 'hour': 3600, 'h': 3600, - 'day': 24 * 3600, 'month': 30 * 24 * 3600, 'year': 365 * 24 * 3600 - }[interval] - return int(count), int(1 if factor is None else factor) * interval_secs +LIMITER_ALGORITHM_HINT: TypeAlias = Literal[LimitTypeEnum.LEAKY_BUCKET, LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY] class Limiter: def __init__(self, name: str): self._name: Final = name - self._limits: Tuple[RateLimit, ...] = () - self._skipped = 0 + self._limits: Tuple[BaseRateLimit, ...] = () + self._skips = 0 def __repr__(self): return f'<{self.__class__.__name__} {self._name:s}>' - def add_limit(self, allowed: int, expiry: int) -> 'Limiter': + def add_limit(self, allowed: int, interval: int, + algorithm: LIMITER_ALGORITHM_HINT = LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY) -> 'Limiter': """Add a new rate limit :param allowed: How many hits are allowed - :param expiry: Interval in seconds + :param interval: Interval in seconds + :param algorithm: Which algorithm should this limit use """ if allowed <= 0 or not isinstance(allowed, int): msg = f'Allowed must be an int >= 0, is {allowed} ({type(allowed)})' raise ValueError(msg) - if expiry <= 0 or not isinstance(expiry, int): - msg = f'Expire time must be an int >= 0, is {expiry} ({type(expiry)})' + if interval <= 0 or not isinstance(interval, int): + msg = f'Expire time must be an int >= 0, is {interval} ({type(interval)})' raise ValueError(msg) + algo = LimitTypeEnum(algorithm) + if algo is LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY: + cls = FixedWindowElasticExpiryLimit + elif algo is LimitTypeEnum.LEAKY_BUCKET: + cls = LeakyBucketLimit + else: + raise ValueError() + + # Check if we have already added an algorithm with these parameters for window in self._limits: - if window.allowed == allowed and window.expiry == expiry: + if isinstance(window, cls) and window.allowed == allowed and window.interval == interval: return self - limit = RateLimit(allowed, expiry) - self._limits = tuple(sorted([*self._limits, limit], key=lambda x: x.expiry)) + limit = cls(allowed, interval) + self._limits = tuple(sorted([*self._limits, limit], key=lambda x: x.interval)) return self - def parse_limits(self, *text: str) -> 'Limiter': + def parse_limits(self, *text: str, + algorithm: LIMITER_ALGORITHM_HINT = LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY) -> 'Limiter': """Add one or more limits in textual form, e.g. ``5 in 60s``, ``10 per hour`` or ``10/15 mins``. If the limit does already exist it will not be added again. :param text: textual description of limit + :param algorithm: Which algorithm should these limits use """ for limit in [parse_limit(t) for t in text]: - self.add_limit(*limit) + self.add_limit(*limit, algorithm=algorithm) return self def allow(self) -> bool: """Test the limit. - :return: True if allowed, False if forbidden + :return: ``True`` if allowed, ``False`` if forbidden """ allow = True clear_skipped = True @@ -93,17 +95,17 @@ def allow(self) -> bool: clear_skipped = False if clear_skipped: - self._skipped = 0 + self._skips = 0 if not allow: - self._skipped += 1 + self._skips += 1 return allow def test_allow(self) -> bool: """Test the limit without hitting it. Calling this will not increase the hit counter. - :return: True if allowed, False if forbidden + :return: ``True`` if allowed, ``False`` if forbidden """ allow = True clear_skipped = True @@ -119,25 +121,20 @@ def test_allow(self) -> bool: clear_skipped = False if clear_skipped: - self._skipped = 0 + self._skips = 0 return allow def info(self) -> 'LimiterInfo': """Get some info about the limiter and the defined windows """ - now = monotonic() - remaining = max((w.stop for w in self._limits if w.hits), default=now) - now - if remaining <= 0: - remaining = 0 return LimiterInfo( - time_remaining=remaining, skipped=self._skipped, - limits=[limit.window_info() for limit in self._limits] + skips=self._skips, + limits=[limit.info() for limit in self._limits] ) @dataclass class LimiterInfo: - time_remaining: float #: time remaining until skipped will reset - skipped: int #: how many entries were skipped - limits: List['RateLimitInfo'] # Info for every window + skips: int #: How many entries were skipped + limits: List[Union[FixedWindowElasticExpiryLimitInfo, LeakyBucketLimitInfo]] #: Info for every limit diff --git a/src/HABApp/util/rate_limiter/limits/__init__.py b/src/HABApp/util/rate_limiter/limits/__init__.py new file mode 100644 index 00000000..2d7e52ce --- /dev/null +++ b/src/HABApp/util/rate_limiter/limits/__init__.py @@ -0,0 +1,3 @@ +from .base import BaseRateLimit +from .fixed_window import FixedWindowElasticExpiryLimit, FixedWindowElasticExpiryLimitInfo +from .leaky_bucket import LeakyBucketLimit, LeakyBucketLimitInfo diff --git a/src/HABApp/util/rate_limiter/limits/base.py b/src/HABApp/util/rate_limiter/limits/base.py new file mode 100644 index 00000000..ba68c156 --- /dev/null +++ b/src/HABApp/util/rate_limiter/limits/base.py @@ -0,0 +1,73 @@ +from dataclasses import dataclass +from typing import Final + + +@dataclass +class BaseRateLimitInfo: + hits: int #: Hits + skips: int #: Skips + limit: int #: Boundary + + @property + def hits_remaining(self) -> int: + return self.limit - self.hits + + +class BaseRateLimit: + def __init__(self, allowed: int, interval: int): + super().__init__() + assert allowed > 0, allowed + assert interval > 0, interval + + self.interval: Final = interval + self.allowed: Final = allowed + + self.hits: int = 0 + self.skips: int = 0 + + def repr_text(self) -> str: + return '' + + def __repr__(self): + return ( + f'<{self.__class__.__name__} hits={self.hits:d}/{self.allowed:d} interval={self.interval:d}s ' + f'{self.repr_text():s}>' + ) + + def do_test_allow(self): + raise NotImplementedError() + + def do_allow(self): + raise NotImplementedError() + + def do_deny(self): + raise NotImplementedError() + + def info(self) -> BaseRateLimitInfo: + raise NotImplementedError() + + def allow(self, weight: int = 1) -> bool: + if not isinstance(weight, int) or weight <= 0: + msg = f'weight must be an int > 0, is {weight}' + raise ValueError(msg) + + self.do_allow() + + self.hits += weight + if self.hits <= self.allowed: + return True + + self.skips += 1 + self.hits = self.allowed + + if self.do_deny: + self.do_deny() + return False + + def test_allow(self, weight: int = 1) -> bool: + if not isinstance(weight, int) or weight <= 0: + msg = f'weight must be an int > 0, is {weight}' + raise ValueError(msg) + + self.do_test_allow() + return self.hits + weight <= self.allowed diff --git a/src/HABApp/util/rate_limiter/limits/fixed_window.py b/src/HABApp/util/rate_limiter/limits/fixed_window.py new file mode 100644 index 00000000..48b15ba7 --- /dev/null +++ b/src/HABApp/util/rate_limiter/limits/fixed_window.py @@ -0,0 +1,51 @@ +from dataclasses import dataclass +from time import monotonic + +from .base import BaseRateLimit, BaseRateLimitInfo + + +@dataclass +class FixedWindowElasticExpiryLimitInfo(BaseRateLimitInfo): + time_remaining: float #: Time remaining until this window will reset + + +class FixedWindowElasticExpiryLimit(BaseRateLimit): + def __init__(self, allowed: int, interval: int): + super().__init__(allowed, interval) + + self.start: float = -1.0 + self.stop: float = -1.0 + + def repr_text(self): + return f'window={self.stop - self.start:.0f}s' + + def do_test_allow(self): + if self.stop <= monotonic(): + self.hits = 0 + self.skips = 0 + + def do_allow(self): + now = monotonic() + + if self.stop <= now: + self.hits = 0 + self.skips = 0 + self.start = now + self.stop = now + self.interval + + def do_deny(self): + self.stop = monotonic() + self.interval + + def info(self) -> FixedWindowElasticExpiryLimitInfo: + self.do_test_allow() + + remaining = self.stop - monotonic() + if remaining <= 0: + remaining = 0 + + if not remaining and not self.hits: + remaining = self.interval + + return FixedWindowElasticExpiryLimitInfo( + time_remaining=remaining, hits=self.hits, skips=self.skips, limit=self.allowed + ) diff --git a/src/HABApp/util/rate_limiter/limits/leaky_bucket.py b/src/HABApp/util/rate_limiter/limits/leaky_bucket.py new file mode 100644 index 00000000..9f61ed7b --- /dev/null +++ b/src/HABApp/util/rate_limiter/limits/leaky_bucket.py @@ -0,0 +1,50 @@ +from dataclasses import dataclass +from time import monotonic +from typing import Final + +from .base import BaseRateLimit, BaseRateLimitInfo + + +@dataclass +class LeakyBucketLimitInfo(BaseRateLimitInfo): + time_remaining: float #: Time remaining until the next drop + + +class LeakyBucketLimit(BaseRateLimit): + def __init__(self, allowed: int, interval: int): + super().__init__(allowed, interval) + + self.drop_interval: Final = interval / allowed + self.next_drop: float = -1.0 + + def repr_text(self): + return f'drop_interval={self.drop_interval:.1f}s' + + def do_test_allow(self): + + while self.next_drop <= monotonic(): + self.hits -= 1 + self.next_drop += self.drop_interval + + if self.hits <= 0: + # out of drop interval -> reset stats + if self.next_drop <= monotonic(): + self.next_drop = monotonic() + self.drop_interval + self.skips = 0 + + self.hits = 0 + break + + do_allow = do_test_allow + do_deny = None + + def info(self) -> LeakyBucketLimitInfo: + self.do_test_allow() + + remaining = self.next_drop - monotonic() + if remaining <= 0: + remaining = 0 + + return LeakyBucketLimitInfo( + time_remaining=remaining, hits=self.hits, skips=self.skips, limit=self.allowed + ) diff --git a/src/HABApp/util/rate_limiter/parser.py b/src/HABApp/util/rate_limiter/parser.py new file mode 100644 index 00000000..44e8b4a1 --- /dev/null +++ b/src/HABApp/util/rate_limiter/parser.py @@ -0,0 +1,27 @@ +import re +from typing import Tuple + +LIMIT_REGEX = re.compile( + r""" + \s* ([1-9][0-9]*) + \s* (/|per|in) + \s* ([1-9][0-9]*)? + \s* (s|sec|second|m|min|minute|h|hour|day|month|year)s? + \s*""", + re.IGNORECASE | re.VERBOSE, +) + + +def parse_limit(text: str) -> Tuple[int, int]: + if not isinstance(text, str) or not (m := LIMIT_REGEX.fullmatch(text)): + msg = f'Invalid limit string: "{text:s}"' + raise ValueError(msg) + + count, per, factor, interval = m.groups() + + interval_secs = { + 's': 1, 'sec': 1, 'second': 1, 'm': 60, 'min': 60, 'minute': 60, 'hour': 3600, 'h': 3600, + 'day': 24 * 3600, 'month': 30 * 24 * 3600, 'year': 365 * 24 * 3600 + }[interval] + + return int(count), int(1 if factor is None else factor) * interval_secs diff --git a/src/HABApp/util/rate_limiter/rate_limit.py b/src/HABApp/util/rate_limiter/rate_limit.py deleted file mode 100644 index 500b559b..00000000 --- a/src/HABApp/util/rate_limiter/rate_limit.py +++ /dev/null @@ -1,73 +0,0 @@ -from dataclasses import dataclass -from time import monotonic -from typing import Final - - -@dataclass -class RateLimitInfo: - time_remaining: float #: Time remaining until this window will reset - hits: int #: Hits - skips: int #: Skips - limit: int #: Boundary - - @property - def hits_remaining(self) -> int: - return self.limit - self.hits - - -class RateLimit: - def __init__(self, allowed: int, expiry: int): - super().__init__() - assert allowed > 0, allowed - assert expiry > 0, expiry - - self.expiry: Final = expiry - self.allowed: Final = allowed - - self.start: float = -1.0 - self.stop: float = -1.0 - self.hits: int = 0 - self.skips: int = 0 - - def __repr__(self): - return (f'<{self.__class__.__name__} hits={self.hits:d}/{self.allowed:d} ' - f'expiry={self.expiry:d}s window={self.stop - self.start:.0f}s>') - - def allow(self) -> bool: - now = monotonic() - - if self.stop < now: - self.hits = 0 - self.skips = 0 - self.start = now - self.stop = now + self.expiry - - self.hits += 1 - if self.hits <= self.allowed: - return True - - self.skips += 1 - self.hits = self.allowed - self.stop = now + self.expiry - return False - - def test_allow(self) -> bool: - now = monotonic() - - if self.hits and self.stop < now: - self.hits = 0 - self.skips = 0 - - return self.hits < self.allowed - - def window_info(self) -> RateLimitInfo: - if self.hits <= 0: - remaining = self.expiry - else: - remaining = self.stop - monotonic() - if remaining <= 0: - remaining = 0 - - return RateLimitInfo( - time_remaining=remaining, hits=self.hits, skips=self.skips, limit=self.allowed - ) diff --git a/tests/test_utils/test_rate_limiter.py b/tests/test_utils/test_rate_limiter.py index 820ec820..87c68de0 100644 --- a/tests/test_utils/test_rate_limiter.py +++ b/tests/test_utils/test_rate_limiter.py @@ -1,18 +1,48 @@ +import re + import pytest -import HABApp.util.rate_limiter.limiter as limiter_module -import HABApp.util.rate_limiter.rate_limit as rate_limit_module +import HABApp.util.rate_limiter.limits.fixed_window as fixed_window_module +import HABApp.util.rate_limiter.limits.leaky_bucket as leaky_bucket_module import HABApp.util.rate_limiter.registry as registry_module -from HABApp.util.rate_limiter.limiter import Limiter, RateLimit, parse_limit +from HABApp.util.rate_limiter.limiter import ( + FixedWindowElasticExpiryLimit, + FixedWindowElasticExpiryLimitInfo, + LeakyBucketLimit, + LeakyBucketLimitInfo, + Limiter, + parse_limit, +) +from HABApp.util.rate_limiter.parser import LIMIT_REGEX + + +class MockedMonotonic: + def __init__(self): + self.time = 0 + + def get_time(self): + return self.time + + def __iadd__(self, other): + self.time += other + return self + + +@pytest.fixture() +def time(monkeypatch) -> MockedMonotonic: + m = MockedMonotonic() + monkeypatch.setattr(fixed_window_module, 'monotonic', m.get_time) + monkeypatch.setattr(leaky_bucket_module, 'monotonic', m.get_time) + return m @pytest.mark.parametrize( - 'unit,factor', ( + 'unit,factor', [ ('s', 1), ('sec', 1), ('second', 1), ('m', 60), ('min', 60), ('minute', 60), ('h', 3600), ('hour', 3600), ('day', 24 * 3600), ('month', 30 * 24 * 3600), ('year', 365 * 24 * 3600) - ) + ] ) def test_parse(unit: str, factor: int): assert parse_limit(f' 1 per {unit} ') == (1, factor) @@ -29,138 +59,176 @@ def test_parse(unit: str, factor: int): assert str(e.value) == 'Invalid limit string: "asdf"' -def test_window(monkeypatch): - time = 0 - monkeypatch.setattr(rate_limit_module, 'monotonic', lambda: time) +def test_regex_all_units(): + m = re.search(r'\(([^)]+)\)s\?', LIMIT_REGEX.pattern) + values = m.group(1) - limit = RateLimit(5, 3) - assert str(limit) == '' - assert limit.test_allow() + for unit in values.split('|'): + parse_limit(f'1 in 3 {unit}') + parse_limit(f'1 in 3 {unit}s') + + +def test_fixed_window(time): + + limit = FixedWindowElasticExpiryLimit(5, 3) + assert str(limit) == '' + for _ in range(10): + assert limit.test_allow() assert limit.allow() - assert str(limit) == '' + assert str(limit) == '' for _ in range(4): assert limit.allow() - assert str(limit) == '' + assert str(limit) == '' # Limit is full, stop gets moved further - time = 1 + time += 1 assert not limit.allow() - assert str(limit) == '' + assert str(limit) == '' # move out of interval - time = 4.1 + time += 3.1 + assert limit.allow() + assert limit.hits == 1 + assert str(limit) == '' + + +def test_leaky_bucket(time): + limit = LeakyBucketLimit(4, 2) + assert str(limit) == '' + for _ in range(10): + assert limit.test_allow() + + assert limit.allow() + assert limit.hits == 1 + + assert limit.allow() + assert limit.hits == 2 + + assert limit.allow() + assert limit.allow() + assert not limit.allow() + assert not limit.allow() + assert limit.hits == 4 + + time += 0.5 + assert limit.test_allow() + assert limit.hits == 3 + + time += 1.7 + assert limit.test_allow() + assert limit.hits == 0 + + assert limit.allow() + assert limit.hits == 1 + + time += 0.3 + assert limit.test_allow() + assert limit.hits == 0 + + time += 1 assert limit.allow() + time += 0.4999 + + limit.test_allow() assert limit.hits == 1 - assert str(limit) == '' + time += 0.0001 + limit.test_allow() + assert limit.hits == 0 -def test_window_test_allow(monkeypatch): - time = 0 - monkeypatch.setattr(rate_limit_module, 'monotonic', lambda: time) - limit = RateLimit(5, 3) +def test_window_test_allow(time): + + limit = FixedWindowElasticExpiryLimit(5, 3) limit.hits = 5 limit.stop = 2.99999 assert not limit.test_allow() # expiry when out of window - time = 3 + time += 3 assert limit.test_allow() assert not limit.hits -def test_limiter_add(monkeypatch): +def test_limiter_add(time): limiter = Limiter('test') limiter.add_limit(3, 5).add_limit(3, 5).parse_limits('3 in 5s') assert len(limiter._limits) == 1 -def test_limiter_info(monkeypatch): - time = 0 - monkeypatch.setattr(rate_limit_module, 'monotonic', lambda: time) - monkeypatch.setattr(limiter_module, 'monotonic', lambda: time) +def test_fixed_window_info(time): + limit = FixedWindowElasticExpiryLimit(5, 3) + Info = FixedWindowElasticExpiryLimitInfo - limiter = Limiter('test') + assert limit.info() == Info(hits=0, skips=0, limit=5, time_remaining=3) - info = limiter.info() - assert info.time_remaining == 0 - assert info.skipped == 0 + limit.allow() + assert limit.info() == Info(hits=1, skips=0, limit=5, time_remaining=3) + limit.allow(4) + assert limit.info() == Info(hits=5, skips=0, limit=5, time_remaining=3) + limit.allow() + assert limit.info() == Info(hits=5, skips=1, limit=5, time_remaining=3) - with pytest.raises(ValueError): - limiter.allow() + time += 1 + assert limit.info() == Info(hits=5, skips=1, limit=5, time_remaining=2) - with pytest.raises(ValueError): - limiter.test_allow() + time += 3 + assert limit.info() == Info(hits=0, skips=0, limit=5, time_remaining=3) - limiter.add_limit(3, 3) + assert not limit.test_allow(6) + assert limit.info() == Info(hits=0, skips=0, limit=5, time_remaining=3) - info = limiter.info() - assert info.time_remaining == 0 - assert info.skipped == 0 - w_info = info.limits[0] - assert w_info.limit == 3 - assert w_info.skips == 0 - assert w_info.time_remaining == 3 - assert w_info.hits == 0 +def test_leaky_bucket_info(time): + limit = LeakyBucketLimit(2, 2) + Info = LeakyBucketLimitInfo - limiter.allow() - time = 2 - limiter.allow() + assert limit.info() == Info(hits=0, skips=0, limit=2, time_remaining=1) - info = limiter.info() - assert info.time_remaining == 1 - assert info.skipped == 0 - - w_info = info.limits[0] - assert w_info.limit == 3 - assert w_info.skips == 0 - assert w_info.time_remaining == 1 - assert w_info.hits == 2 - # add a longer limiter - this one should now define the time_remaining - limiter.add_limit(4, 5) - limiter.allow() +def test_registry(monkeypatch): + monkeypatch.setattr(registry_module, '_LIMITERS', {}) - info = limiter.info() - assert info.time_remaining == 5 - assert info.skipped == 0 + obj = registry_module.RateLimiter('Test') + assert obj is registry_module.RateLimiter('TEST') + assert obj is registry_module.RateLimiter('test') - w_info = info.limits[0] - assert w_info.limit == 3 - assert w_info.skips == 0 - assert w_info.time_remaining == 1 - assert w_info.hits == 3 - w_info = info.limits[1] - assert w_info.limit == 4 - assert w_info.skips == 0 - assert w_info.time_remaining == 5 - assert w_info.hits == 1 +def test_limiter(time): - time += 5.0001 + limiter = Limiter('Test') + assert limiter.__repr__() == '' info = limiter.info() - assert info.time_remaining == 0 + assert info.skips == 0 + + with pytest.raises(ValueError): + limiter.allow() - w_info = info.limits[0] - assert w_info.limit == 3 - assert w_info.skips == 0 - assert w_info.time_remaining == 0 - assert w_info.hits == 3 + limiter.add_limit(2, 1).add_limit(2, 2) - w_info = info.limits[1] - assert w_info.limit == 4 - assert w_info.skips == 0 - assert w_info.time_remaining == 0 - assert w_info.hits == 1 + assert limiter.allow() + assert limiter.allow() + time += 0.5 + assert not limiter.allow() + time += 1 + assert not limiter.allow() -def test_registry(monkeypatch): - monkeypatch.setattr(registry_module, '_LIMITERS', {}) + assert limiter.info().skips == 2 + time += 2 - obj = registry_module.RateLimiter('test') - assert obj is registry_module.RateLimiter('TEST') + assert limiter.test_allow() + assert limiter.info().skips == 0 + + assert limiter.allow() + assert limiter.allow() + assert not limiter.allow() + assert limiter.info().skips == 1 + + time += 2 + assert limiter.allow() + assert limiter.info().skips == 0