switch to a dataclass for complex resposne
This commit is contained in:
parent
b82c0af81e
commit
c2c0396b20
10
rrl.py
10
rrl.py
@ -13,6 +13,12 @@ class Tier:
|
||||
per_day: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class DailyUsage:
|
||||
date: datetime.date
|
||||
calls: int
|
||||
|
||||
|
||||
class RateLimitExceeded(Exception):
|
||||
pass
|
||||
|
||||
@ -101,7 +107,7 @@ class RateLimiter:
|
||||
key: str,
|
||||
start: datetime.date,
|
||||
end: typing.Optional[datetime.date] = None,
|
||||
) -> typing.List[typing.Dict[str, typing.Union[datetime.date, int]]]:
|
||||
) -> typing.List[DailyUsage]:
|
||||
if not end:
|
||||
end = datetime.date.today()
|
||||
days = []
|
||||
@ -113,6 +119,6 @@ class RateLimiter:
|
||||
f"{self.prefix}:{zone}:{key}:d{day.strftime('%Y%m%d')}" for day in days
|
||||
]
|
||||
return [
|
||||
{"date": d, "calls": int(calls.decode()) if calls else 0}
|
||||
DailyUsage(d, int(calls.decode()) if calls else 0)
|
||||
for d, calls in zip(days, self.redis.mget(day_keys))
|
||||
]
|
||||
|
@ -1,6 +1,6 @@
|
||||
import datetime
|
||||
import pytest
|
||||
from rrl import Tier, RateLimiter, RateLimitExceeded, _get_redis_connection
|
||||
from rrl import Tier, RateLimiter, RateLimitExceeded, _get_redis_connection, DailyUsage
|
||||
from freezegun import freeze_time
|
||||
|
||||
redis = _get_redis_connection()
|
||||
@ -104,9 +104,9 @@ def test_get_daily_usage():
|
||||
|
||||
with freeze_time("2020-01-15"):
|
||||
usage = rl.get_usage_since("zone", "test-key", datetime.date(2020, 1, 1))
|
||||
assert usage[0] == {"date": datetime.date(2020, 1, 1), "calls": 1}
|
||||
assert usage[3] == {"date": datetime.date(2020, 1, 4), "calls": 4}
|
||||
assert usage[8] == {"date": datetime.date(2020, 1, 9), "calls": 9}
|
||||
assert usage[9] == {"date": datetime.date(2020, 1, 10), "calls": 0}
|
||||
assert usage[14] == {"date": datetime.date(2020, 1, 15), "calls": 0}
|
||||
assert usage[0] == DailyUsage(datetime.date(2020, 1, 1), 1)
|
||||
assert usage[3] == DailyUsage(datetime.date(2020, 1, 4), 4)
|
||||
assert usage[8] == DailyUsage(datetime.date(2020, 1, 9), 9)
|
||||
assert usage[9] == DailyUsage(datetime.date(2020, 1, 10), 0)
|
||||
assert usage[14] == DailyUsage(datetime.date(2020, 1, 15), 0)
|
||||
assert len(usage) == 15
|
||||
|
Loading…
Reference in New Issue
Block a user