jpt-archive/rrl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

raise ValueError on bad tier

Browse Source
This commit is contained in:
James Turk 2020-10-13 12:53:42 -04:00
parent 591539a427
commit f65c2c2c70
3 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -1,2 +1,3 @@
*.pyc *.pyc
dist/ dist/
*.egg-info/

5
rrl.py
View File

@ -52,12 +52,15 @@ class RateLimiter:
self.track_daily_usage = track_daily_usage self.track_daily_usage = track_daily_usage
def check_limit(self, zone: str, key: str, tier_name: str) -> bool: def check_limit(self, zone: str, key: str, tier_name: str) -> bool:
try:
tier = self.tiers[tier_name]
except KeyError:
raise ValueError(f"unknown tier: {tier_name}")
if self.use_redis_time: if self.use_redis_time:
timestamp = self.redis.time()[0] timestamp = self.redis.time()[0]
now = datetime.datetime.fromtimestamp(timestamp) now = datetime.datetime.fromtimestamp(timestamp)
else: else:
now = datetime.datetime.utcnow() now = datetime.datetime.utcnow()
tier = self.tiers[tier_name]
pipe = self.redis.pipeline() pipe = self.redis.pipeline()
if tier.per_minute: if tier.per_minute:

8
test_ratelimit.py
View File

@ -58,6 +58,14 @@ def test_using_redis_time():
assert count == 10 assert count == 10
def test_invalid_tier():
redis.flushall()
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)
with pytest.raises(ValueError):
rl.check_limit("test-zone", "test-key", "non-existent-tier")
def test_multiple_zones(): def test_multiple_zones():
redis.flushall() redis.flushall()
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True) rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)