Fix bug in the ratelimit helper & add a few more tests

Author: sylvinus

mrq/helpers.py

@@ -5,7 +5,7 @@
 
 def ratelimit(key, limit, per=1, redis=None):
     """ Returns an integer with the number of available actions for the
-    current period. If zero, rate was already reached. """
+    current period in seconds. If zero, rate was already reached. """
 
     if redis is None:
         redis = connections.redis
@@ -17,10 +17,10 @@ def ratelimit(key, limit, per=1, redis=None):
 
     with redis.pipeline(transaction=True) as pipeline:
         pipeline.incr(k, 1)
-        pipeline.expire(k, 10)
+        pipeline.expire(k, per + 10)
         value = pipeline.execute()
 
-    current = int(value[0])
+    current = int(value[0]) - 1
 
     if current >= limit:
         return 0

tests/test_ratelimit.py

@@ -6,12 +6,16 @@ def test_helpers_ratelimit(worker):
 
     worker.start_deps()
 
-    for i in range(1, 10):
+    assert ratelimit("k3", 1, per=1) == 1
+    assert ratelimit("k3", 1, per=1) == 0
+    assert ratelimit("k3", 1, per=1) == 0
+
+    for i in range(0, 10):
         r = ratelimit("k", 10, per=1)
         assert r == 10 - i
 
     assert ratelimit("k", 10, per=1) == 0
-    assert ratelimit("k2", 5, per=1) == 4
+    assert ratelimit("k2", 5, per=1) == 5
 
     # We *could* have failures there if we go over a second but we've not seen
     # it much so far.
@@ -21,4 +25,9 @@ def test_helpers_ratelimit(worker):
     # TODO: test the "per" argument a bit better.
     time.sleep(1)
 
+    assert ratelimit("k", 10, per=1) == 10
     assert ratelimit("k", 10, per=1) == 9
+
+    # This is actually another counter.
+    assert ratelimit("k", 10, per=10) == 10
+    assert ratelimit("k", 10, per=10) == 9