def test_ratelimit_exceeded(self, limiter):
"""Verify that if we're inside our cell, we reduce remaining."""
rate = helpers.new_quota(
period=datetime.timedelta(seconds=60), count=50
)
mockstore = limiter.store.recording_store
mockstore.get.side_effect = lambda key: (
limit_data.LimitData(
used=49,
remaining=1,
created_at=datetime.datetime.now(datetime.timezone.utc),
time=(
datetime.datetime.now(datetime.timezone.utc)
+ datetime.timedelta(seconds=60)
),
)
)
limitresult = limiter.rate_limit(key="key", rate=rate, quantity=1)
assert limitresult.limited is True
assert limitresult.remaining == 0
assert (
datetime.timedelta(seconds=60)
<= limitresult.reset_after
< (2 * datetime.timedelta(seconds=60))
)
assert (
datetime.timedelta(seconds=0)
< limitresult.retry_after
<= datetime.timedelta(seconds=3)
)
def test_copy_with_arguments(self):
"""Verify the behaviour of LimitData.copy_with."""
ld = limit_data.LimitData(used=5, remaining=5)
new_ld = ld.copy_with(used=10, remaining=0)
assert new_ld.remaining == 0
assert new_ld.used == 10
assert new_ld.created_at == new_ld.created_at
def test_set_with_time_uses_now(self):
"""Verify we can add data with the current time."""
store = dictstore.DictionaryStore()
new_data = limit_data.LimitData(used=9999, remaining=1)
set_data = store.set_with_time(key="mykey", data=new_data)
assert isinstance(set_data.time, datetime.datetime)
assert store.get("mykey") != {}
def test_get_with_time_defaults_to_now(self):
"""Verify we can retrieve data with a default time."""
data = limit_data.LimitData(used=9999, remaining=1)
store = dictstore.DictionaryStore(store={"mykey": data})
dt, retrieved_data = store.get_with_time("mykey")
assert dt.replace(second=0, microsecond=0) == datetime.datetime.now(
datetime.timezone.utc).replace(second=0, microsecond=0)
assert retrieved_data == data.copy_with(time=dt)
def test_rate_limit_reset_after_period(self, limiter):
"""Verify we allow the last request."""
rate = helpers.new_quota()
mockstore = limiter.store.recording_store
now = datetime.datetime.now(datetime.timezone.utc)
original_created_at = now - datetime.timedelta(seconds=2)
mockstore.get.return_value = limit_data.LimitData(
remaining=0, used=5, created_at=original_created_at)
mockstore.set.return_value = limit_data.LimitData(remaining=4,
used=1,
created_at=now)
limitresult = limiter.rate_limit(key="key", quantity=1, rate=rate)
assert limitresult.remaining == 4
assert limitresult.limit == 5
assert limitresult.limited is False
mockstore.get.assert_called_once_with("key")
mockstore.set.assert_called_once_with(key="key", data=mock.ANY)
def test_last_rate_limit_in_period(self, limiter):
"""Verify we allow the last request."""
rate = helpers.new_quota()
mockstore = limiter.store.recording_store
original_created_at = datetime.datetime.now(
datetime.timezone.utc) - datetime.timedelta(microseconds=1)
mockstore.get.return_value = limit_data.LimitData(
remaining=1, used=4, created_at=original_created_at)
limitresult = limiter.rate_limit(key="key", quantity=1, rate=rate)
assert limitresult.remaining == 0
assert limitresult.limit == 5
assert limitresult.limited is False
mockstore.get.assert_called_once_with("key")
mockstore.set.assert_called_once_with(
key="key",
data=limit_data.LimitData(used=5,
remaining=0,
created_at=original_created_at),
)
def test_preexisting_limitdata(self, limiter):
"""Verify we record the limit data appropriately."""
rate = helpers.new_quota()
mockstore = limiter.store.recording_store
original_created_at = datetime.datetime.now(
datetime.timezone.utc) - datetime.timedelta(microseconds=1)
mockstore.get.return_value = limit_data.LimitData(
remaining=4, used=1, created_at=original_created_at)
limitresult = limiter.rate_limit(key="key", quantity=1, rate=rate)
assert limitresult.remaining == 3
assert limitresult.limit == 5
assert limitresult.limited is False
mockstore.get.assert_called_once_with("key")
mockstore.set.assert_called_once_with(
key="key",
data=limit_data.LimitData(used=2,
remaining=3,
created_at=original_created_at),
)
def test_reset(self, limiter):
"""Verify reset works appropriately."""
rate = helpers.new_quota()
mockstore = limiter.store.recording_store
mockstore.set.return_value = limit_data.LimitData(remaining=5, used=0)
limitresult = limiter.reset(key="key", rate=rate)
assert limitresult.remaining == 5
assert limitresult.limit == 5
assert limitresult.limited is False
mockstore.set.assert_called_once_with(key="key", data=mock.ANY)
def test_result_from_quota(self, limiter):
"""Verify the behaviour of result_from_quota."""
quota = mock.Mock(count=5, period=datetime.timedelta(hours=1))
limitresult = limiter.result_from_quota(
rate=quota,
limited=False,
limitdata=limit_data.LimitData(used=0, remaining=6),
elapsed_since_period_start=datetime.timedelta(seconds=1),
)
assert isinstance(limitresult, result.RateLimitResult)
assert limitresult.limited is False
assert limitresult.limit == 5
assert limitresult.remaining == 6
assert limitresult.retry_after == datetime.timedelta(seconds=-1)
def test_get_with_time_uses_existing_time(self):
"""Verify we can retrieve data from our datastore with its time."""
data = limit_data.LimitData(
used=9999,
remaining=1,
time=datetime.datetime(
year=2018,
month=12,
day=4,
hour=9,
minute=0,
second=0,
tzinfo=datetime.timezone.utc,
),
)
store = dictstore.DictionaryStore(store={"mykey": data})
dt, retrieved_data = store.get_with_time("mykey")
assert dt == data.time
assert retrieved_data == data
def test_set_with_time_uses_provided_value(self):
"""Verify we can add data with a specific time."""
store = dictstore.DictionaryStore()
new_data = limit_data.LimitData(
used=9999,
remaining=1,
time=datetime.datetime(
year=2018,
month=12,
day=4,
hour=9,
minute=0,
second=0,
tzinfo=datetime.timezone.utc,
),
)
set_data = store.set_with_time(key="mykey", data=new_data)
assert set_data == new_data
assert store.get("mykey") == new_data
def test_asdict_without_time(self):
"""Verify the behaviour of LimitData.asdict."""
created_at = datetime.datetime(
year=2018,
month=12,
day=11,
hour=12,
minute=12,
second=15,
microsecond=123_456,
tzinfo=datetime.timezone.utc,
)
ld = limit_data.LimitData(used=0, remaining=5, created_at=created_at)
assert ld.asdict() == {
"time": "",
"created_at": "2018-12-11T12:12:15.123456+0000",
"used": "0",
"remaining": "5",
}
def test_conversion(self):
"""Test attrs conversion just to be safe."""
kwargs = {
"time": "",
"created_at": "2018-12-11T12:12:15.123456+0000",
"used": "0",
"remaining": "5",
}
created_at = datetime.datetime(
year=2018,
month=12,
day=11,
hour=12,
minute=12,
second=15,
microsecond=123_456,
tzinfo=datetime.timezone.utc,
)
ld = limit_data.LimitData(**kwargs)
assert ld.used == 0
assert ld.remaining == 5
assert ld.created_at == created_at
assert ld.time is None
def test_copy_with_no_arguments(self):
"""Verify the behaviour of LimitData.copy_with."""
ld = limit_data.LimitData(used=5, remaining=5)
assert ld.copy_with() == ld
def test_get(self):
"""Verify we can retrieve data from our datastore."""
data = limit_data.LimitData(used=9999, remaining=1)
store = dictstore.DictionaryStore(store={"mykey": data})
assert store.get("mykey") == data
def test_set(self):
"""Verify we can add data."""
store = dictstore.DictionaryStore()
new_data = limit_data.LimitData(used=9999, remaining=1)
assert store.set(key="mykey", data=new_data) == new_data
def test_may_begin_life_with_data(self):
"""Verify that we can give it initial data."""
data = {"a_key": limit_data.LimitData(used=9999, remaining=1)}
store = dictstore.DictionaryStore(store=data)
assert store.store == data
