Move storage tests to their own files

tests/storage/test_gae_memcached.py

@@ -0,0 +1,52 @@
+import time
+import unittest
+
+import pytest
+
+from limits import RateLimitItemPerSecond
+from limits.storage import GAEMemcachedStorage
+from limits.strategies import (
+    FixedWindowRateLimiter,
+    FixedWindowElasticExpiryRateLimiter
+)
+from tests import RUN_GAE, fixed_start
+
+
+@pytest.mark.unit
+@unittest.skipUnless(RUN_GAE, reason='Only for GAE')
+class GAEMemcachedStorageTests(unittest.TestCase):
+    def setUp(self):
+        from google.appengine.ext import testbed
+        tb = testbed.Testbed()
+        tb.activate()
+        tb.init_memcache_stub()
+
+    @fixed_start
+    def test_fixed_window(self):
+        storage = GAEMemcachedStorage("gaememcached://")
+        limiter = FixedWindowRateLimiter(storage)
+        per_min = RateLimitItemPerSecond(10)
+        start = time.time()
+        count = 0
+        while time.time() - start < 0.5 and count < 10:
+            self.assertTrue(limiter.hit(per_min))
+            count += 1
+        self.assertFalse(limiter.hit(per_min))
+        while time.time() - start <= 1:
+            time.sleep(0.1)
+        self.assertTrue(limiter.hit(per_min))
+
+    @fixed_start
+    def test_fixed_window_with_elastic_expiry_cluster(self):
+        storage = GAEMemcachedStorage("gaememcached://")
+        limiter = FixedWindowElasticExpiryRateLimiter(storage)
+        per_sec = RateLimitItemPerSecond(2, 2)
+
+        self.assertTrue(limiter.hit(per_sec))
+        time.sleep(1)
+        self.assertTrue(limiter.hit(per_sec))
+        self.assertFalse(limiter.test(per_sec))
+        time.sleep(1)
+        self.assertFalse(limiter.test(per_sec))
+        time.sleep(1)
+        self.assertTrue(limiter.test(per_sec))

tests/storage/test_memcached.py

@@ -0,0 +1,101 @@
+import time
+import unittest
+
+import mock
+import pytest
+import pymemcache.client
+
+from limits import RateLimitItemPerSecond, RateLimitItemPerMinute
+from limits.storage import storage_from_string, MemcachedStorage
+from limits.strategies import FixedWindowRateLimiter, \
+    FixedWindowElasticExpiryRateLimiter
+from tests import fixed_start
+
+
+@pytest.mark.unit
+class MemcachedStorageTests(unittest.TestCase):
+    def setUp(self):
+        pymemcache.client.Client(('localhost', 22122)).flush_all()
+        self.storage_url = 'memcached://localhost:22122'
+
+    def test_options(self):
+        with mock.patch(
+            "limits.storage.memcached.get_dependency"
+        ) as get_dependency:
+            storage_from_string(self.storage_url, connect_timeout=1).check()
+            self.assertEqual(
+                get_dependency().Client.call_args[1]['connect_timeout'], 1
+            )
+
+    @fixed_start
+    def test_fixed_window(self):
+        storage = MemcachedStorage("memcached://localhost:22122")
+        limiter = FixedWindowRateLimiter(storage)
+        per_min = RateLimitItemPerSecond(10)
+        start = time.time()
+        count = 0
+        while time.time() - start < 0.5 and count < 10:
+            self.assertTrue(limiter.hit(per_min))
+            count += 1
+        self.assertFalse(limiter.hit(per_min))
+        while time.time() - start <= 1:
+            time.sleep(0.1)
+        self.assertTrue(limiter.hit(per_min))
+
+    @fixed_start
+    def test_fixed_window_cluster(self):
+        storage = MemcachedStorage(
+            "memcached://localhost:22122,localhost:22123"
+        )
+        limiter = FixedWindowRateLimiter(storage)
+        per_min = RateLimitItemPerSecond(10)
+        start = time.time()
+        count = 0
+        while time.time() - start < 0.5 and count < 10:
+            self.assertTrue(limiter.hit(per_min))
+            count += 1
+        self.assertFalse(limiter.hit(per_min))
+        while time.time() - start <= 1:
+            time.sleep(0.1)
+        self.assertTrue(limiter.hit(per_min))
+
+    @fixed_start
+    def test_fixed_window_with_elastic_expiry(self):
+        storage = MemcachedStorage("memcached://localhost:22122")
+        limiter = FixedWindowElasticExpiryRateLimiter(storage)
+        per_sec = RateLimitItemPerSecond(2, 2)
+
+        self.assertTrue(limiter.hit(per_sec))
+        time.sleep(1)
+        self.assertTrue(limiter.hit(per_sec))
+        self.assertFalse(limiter.test(per_sec))
+        time.sleep(1)
+        self.assertFalse(limiter.test(per_sec))
+        time.sleep(1)
+        self.assertTrue(limiter.test(per_sec))
+
+    @fixed_start
+    def test_fixed_window_with_elastic_expiry_cluster(self):
+        storage = MemcachedStorage(
+            "memcached://localhost:22122,localhost:22123"
+        )
+        limiter = FixedWindowElasticExpiryRateLimiter(storage)
+        per_sec = RateLimitItemPerSecond(2, 2)
+
+        self.assertTrue(limiter.hit(per_sec))
+        time.sleep(1)
+        self.assertTrue(limiter.hit(per_sec))
+        self.assertFalse(limiter.test(per_sec))
+        time.sleep(1)
+        self.assertFalse(limiter.test(per_sec))
+        time.sleep(1)
+        self.assertTrue(limiter.test(per_sec))
+
+    def test_clear(self):
+        storage = MemcachedStorage("memcached://localhost:22122")
+        limiter = FixedWindowRateLimiter(storage)
+        per_min = RateLimitItemPerMinute(1)
+        limiter.hit(per_min)
+        self.assertFalse(limiter.hit(per_min))
+        limiter.clear(per_min)
+        self.assertTrue(limiter.hit(per_min))

tests/storage/test_memory.py

@@ -0,0 +1,78 @@
+import time
+import unittest
+
+import hiro
+import pytest
+
+from limits import RateLimitItemPerMinute, RateLimitItemPerSecond
+from limits.storage import MemoryStorage
+from limits.strategies import FixedWindowRateLimiter, MovingWindowRateLimiter
+
+
+@pytest.mark.unit
+class MemoryStorageTests(unittest.TestCase):
+    def setUp(self):
+        self.storage = MemoryStorage()
+
+    def test_in_memory(self):
+        with hiro.Timeline().freeze() as timeline:
+            limiter = FixedWindowRateLimiter(self.storage)
+            per_min = RateLimitItemPerMinute(10)
+            for i in range(0, 10):
+                self.assertTrue(limiter.hit(per_min))
+            self.assertFalse(limiter.hit(per_min))
+            timeline.forward(61)
+            self.assertTrue(limiter.hit(per_min))
+
+    def test_fixed_window_clear(self):
+        limiter = FixedWindowRateLimiter(self.storage)
+        per_min = RateLimitItemPerMinute(1)
+        limiter.hit(per_min)
+        self.assertFalse(limiter.hit(per_min))
+        limiter.clear(per_min)
+        self.assertTrue(limiter.hit(per_min))
+
+    def test_moving_window_clear(self):
+        limiter = MovingWindowRateLimiter(self.storage)
+        per_min = RateLimitItemPerMinute(1)
+        limiter.hit(per_min)
+        self.assertFalse(limiter.hit(per_min))
+        limiter.clear(per_min)
+        self.assertTrue(limiter.hit(per_min))
+
+    def test_reset(self):
+        limiter = FixedWindowRateLimiter(self.storage)
+        per_min = RateLimitItemPerMinute(10)
+        for i in range(0, 10):
+            self.assertTrue(limiter.hit(per_min))
+        self.assertFalse(limiter.hit(per_min))
+        self.storage.reset()
+        for i in range(0, 10):
+            self.assertTrue(limiter.hit(per_min))
+        self.assertFalse(limiter.hit(per_min))
+
+    def test_expiry(self):
+        with hiro.Timeline().freeze() as timeline:
+            limiter = FixedWindowRateLimiter(self.storage)
+            per_min = RateLimitItemPerMinute(10)
+            for i in range(0, 10):
+                self.assertTrue(limiter.hit(per_min))
+            timeline.forward(60)
+            # touch another key and yield
+            limiter.hit(RateLimitItemPerSecond(1))
+            time.sleep(0.1)
+            self.assertTrue(per_min.key_for() not in self.storage.storage)
+
+    def test_expiry_moving_window(self):
+        with hiro.Timeline().freeze() as timeline:
+            limiter = MovingWindowRateLimiter(self.storage)
+            per_min = RateLimitItemPerMinute(10)
+            per_sec = RateLimitItemPerSecond(1)
+            for _ in range(0, 2):
+                for _ in range(0, 10):
+                    self.assertTrue(limiter.hit(per_min))
+                timeline.forward(60)
+                self.assertTrue(limiter.hit(per_sec))
+                timeline.forward(1)
+                time.sleep(0.1)
+                self.assertEqual([], self.storage.events[per_min.key_for()])

tests/storage/test_redis.py

@@ -0,0 +1,99 @@
+import time
+import unittest
+
+import mock
+import pytest
+import redis
+
+from limits import RateLimitItemPerSecond, RateLimitItemPerMinute
+from limits.storage import RedisStorage, storage_from_string
+from limits.strategies import FixedWindowRateLimiter, MovingWindowRateLimiter
+
+
+@pytest.mark.unit
+class SharedRedisTests(object):
+    def test_fixed_window(self):
+        limiter = FixedWindowRateLimiter(self.storage)
+        per_second = RateLimitItemPerSecond(10)
+        start = time.time()
+        count = 0
+        while time.time() - start < 0.5 and count < 10:
+            self.assertTrue(limiter.hit(per_second))
+            count += 1
+        self.assertFalse(limiter.hit(per_second))
+        while time.time() - start <= 1:
+            time.sleep(0.1)
+        [self.assertTrue(limiter.hit(per_second)) for _ in range(10)]
+
+    def test_reset(self):
+        limiter = FixedWindowRateLimiter(self.storage)
+        for i in range(0, 10):
+            rate = RateLimitItemPerMinute(i)
+            limiter.hit(rate)
+        self.assertEqual(self.storage.reset(), 10)
+
+    def test_fixed_window_clear(self):
+        limiter = FixedWindowRateLimiter(self.storage)
+        per_min = RateLimitItemPerMinute(1)
+        limiter.hit(per_min)
+        self.assertFalse(limiter.hit(per_min))
+        limiter.clear(per_min)
+        self.assertTrue(limiter.hit(per_min))
+
+    def test_moving_window_clear(self):
+        limiter = MovingWindowRateLimiter(self.storage)
+        per_min = RateLimitItemPerMinute(1)
+        limiter.hit(per_min)
+        self.assertFalse(limiter.hit(per_min))
+        limiter.clear(per_min)
+        self.assertTrue(limiter.hit(per_min))
+
+    def test_moving_window_expiry(self):
+        limiter = MovingWindowRateLimiter(self.storage)
+        limit = RateLimitItemPerSecond(2)
+        self.assertTrue(limiter.hit(limit))
+        time.sleep(0.9)
+        self.assertTrue(limiter.hit(limit))
+        self.assertFalse(limiter.hit(limit))
+        time.sleep(0.1)
+        self.assertTrue(limiter.hit(limit))
+        last = time.time()
+        while time.time() - last <= 1:
+            time.sleep(0.05)
+        self.assertTrue(
+            self.storage.storage.keys("%s/*" % limit.namespace) == []
+        )
+
+
+@pytest.mark.unit
+class RedisStorageTests(SharedRedisTests, unittest.TestCase):
+    def setUp(self):
+        self.storage_url = "redis://localhost:7379"
+        self.storage = RedisStorage(self.storage_url)
+        redis.from_url(self.storage_url).flushall()
+
+    def test_init_options(self):
+        with mock.patch(
+            "limits.storage.redis.get_dependency"
+        ) as get_dependency:
+            storage_from_string(self.storage_url, connection_timeout=1)
+            self.assertEqual(
+                get_dependency().from_url.call_args[1]['connection_timeout'], 1
+            )
+
+
+@pytest.mark.unit
+class RedisUnixSocketStorageTests(SharedRedisTests, unittest.TestCase):
+    def setUp(self):
+        self.storage_url = "redis+unix:///tmp/limits.redis.sock"
+        self.storage = RedisStorage(self.storage_url)
+        redis.from_url('unix:///tmp/limits.redis.sock').flushall()
+
+    def test_init_options(self):
+        with mock.patch(
+            "limits.storage.redis.get_dependency"
+        ) as get_dependency:
+            storage_from_string(self.storage_url, connection_timeout=1)
+            self.assertEqual(
+                get_dependency().from_url.call_args[1]['connection_timeout'], 1
+            )

tests/storage/test_redis_cluster.py

@@ -0,0 +1,27 @@
+import unittest
+
+import mock
+import pytest
+import rediscluster
+
+from limits.storage import RedisClusterStorage, storage_from_string
+from tests.storage.test_redis import SharedRedisTests
+
+
+@pytest.mark.unit
+class RedisClusterStorageTests(SharedRedisTests, unittest.TestCase):
+    def setUp(self):
+        rediscluster.RedisCluster("localhost", 7000).flushall()
+        self.storage_url = "redis+cluster://localhost:7000"
+        self.storage = RedisClusterStorage("redis+cluster://localhost:7000")
+
+    def test_init_options(self):
+        with mock.patch(
+            "limits.storage.redis_cluster.get_dependency"
+        ) as get_dependency:
+            storage_from_string(self.storage_url, connection_timeout=1)
+            call_args = get_dependency().RedisCluster.call_args
+            self.assertEqual(
+                call_args[1]['connection_timeout'],
+                1
+            )