diff --git a/diskcache/recipes.py b/diskcache/recipes.py index babb68f..7cee914 100644 --- a/diskcache/recipes.py +++ b/diskcache/recipes.py @@ -265,6 +265,7 @@ def throttle( Assumes keys will not be evicted. Set the eviction policy to 'none' on the cache to guarantee the keys are not evicted. + If the bucket expires, the next call starts with a full bucket of tokens. >>> import diskcache, time >>> cache = diskcache.Cache() @@ -291,8 +292,8 @@ def decorator(func): def wrapper(*args, **kwargs): while True: with cache.transact(retry=True): - last, tally = cache.get(key) now = time_func() + last, tally = cache.get(key, default=(now, count)) tally += (now - last) * rate delay = 0 diff --git a/tests/test_recipes.py b/tests/test_recipes.py index ae74459..ef68132 100644 --- a/tests/test_recipes.py +++ b/tests/test_recipes.py @@ -3,6 +3,7 @@ import shutil import threading import time +from unittest import mock import pytest @@ -28,6 +29,41 @@ def test_averager(cache): assert nums.pop() == 9.5 +@pytest.mark.parametrize('expire', [None, 1]) +@pytest.mark.parametrize('before_first_call', [False, True]) +def test_throttle_after_idle(cache, expire, before_first_call): + clock = [time.time()] + delays = [] + + def sleep(seconds): + delays.append(seconds) + clock[0] += seconds + + with mock.patch('time.time', lambda: clock[0]): + + @dc.throttle( + cache, + 2, + 1, + name='limiter', + expire=expire, + time_func=lambda: clock[0], + sleep_func=sleep, + ) + def work(): + return 'ok' + + if not before_first_call: + assert work() == 'ok' + clock[0] += 10 + + assert work() == 'ok' + assert work() == 'ok' + assert delays == [] + assert work() == 'ok' + assert delays == [0.5] + + def test_lock(cache): state = {'num': 0} lock = dc.Lock(cache, 'demo')