From 5710268b2809bb0a8119bb62be894d7a4d01b93d Mon Sep 17 00:00:00 2001 From: Ali-Akber Saifee Date: Tue, 21 Oct 2014 08:27:13 +0800 Subject: [PATCH] Fix #17: check response from redis and use defaults. --- flask_limiter/storage.py | 9 +++++---- tests/test_regressions.py | 12 ++++++++++-- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/flask_limiter/storage.py b/flask_limiter/storage.py index e58cdd05..7cc192fd 100644 --- a/flask_limiter/storage.py +++ b/flask_limiter/storage.py @@ -253,7 +253,7 @@ def get(self, key): """ :param str key: the key to get the counter value for """ - return int(self.storage.get(key)) + return int(self.storage.get(key) or 0) def acquire_entry(self, key, limit, expiry, no_add=False): """ @@ -286,15 +286,16 @@ def get_moving_window(self, key, limit, expiry): :param int expiry: expiry of entry """ timestamp = time.time() - return tuple(self.lua_moving_window( + window = self.lua_moving_window( [key], [int(timestamp - expiry), limit] - )) + ) + return window or (timestamp, 0) def get_expiry(self, key): """ :param str key: the key to get the expiry for """ - return int(self.storage.ttl(key) + time.time()) + return int((self.storage.ttl(key) or 0) + time.time()) class MemcachedStorage(Storage): """ diff --git a/tests/test_regressions.py b/tests/test_regressions.py index fae6ec78..ea9670f5 100644 --- a/tests/test_regressions.py +++ b/tests/test_regressions.py @@ -36,7 +36,11 @@ def t1(): return "t1" with app.test_client() as cli: - self.assertEqual(200, cli.get("/t1").status_code) + resp = cli.get("/t1") + self.assertEqual( + resp.headers["X-RateLimit-Remaining"], + 5 + ) def test_redis_request_slower_than_moving_window(self): app, limiter = self.build_app({ @@ -51,4 +55,8 @@ def t1(): return "t1" with app.test_client() as cli: - self.assertEqual(200, cli.get("/t1").status_code) + resp = cli.get("/t1") + self.assertEqual( + resp.headers["X-RateLimit-Remaining"], + 5 + )