Touch up test coverage

Fix a typo in the redis store's compare_and_swap as well

src/rush/stores/redis.py

@@ -75,7 +75,7 @@ def compare_and_swap(
                         actual_limit_data=current_data,
                     )
                 p.multi()
-                p.hmset(key, new.asdata())
+                p.hmset(key, new.asdict())
                 p.execute()
             except redis.WatchError as we:
                 raise exceptions.DataChangedInStoreError(

test/unit/test_redis_gcra_limiter.py

@@ -0,0 +1,111 @@
+"""Tests for our fancy Generic Cell Ratelimiter."""
+import collections
+import datetime
+
+import mock
+import pytest
+
+from rush.limiters import redis_gcra as gcra
+from rush.stores import redis
+
+from . import helpers  # noqa: I202
+
+
+LimiterFixture = collections.namedtuple(
+    "LimiterFixture", "client store check_lua apply_lua limiter"
+)
+
+
+@pytest.fixture
+def limiterf():
+    """Provide instantiated periodic limiter."""
+    client = mock.Mock()
+    check_lua = mock.MagicMock()
+    apply_lua = mock.MagicMock()
+    client.register_script.side_effect = [check_lua, apply_lua]
+    store = redis.RedisStore("redis://", client=client)
+    return LimiterFixture(
+        client,
+        store,
+        check_lua,
+        apply_lua,
+        gcra.GenericCellRatelimiter(store=store),
+    )
+
+
+class TestGenericCellRatelimiter:
+    """Tests that exercise our GCRA implementation."""
+
+    def test_reset(self, limiterf):
+        """Verify we reset our GCRA limiter properly."""
+        rate = helpers.new_quota()
+        limiter = limiterf.limiter
+        redis_client = limiterf.client
+
+        limitresult = limiter.reset(key="key", rate=rate)
+
+        redis_client.delete.assert_called_once_with("key")
+        assert limitresult.remaining == 5
+        assert limitresult.limit == 5
+        assert limitresult.limited is False
+        assert limitresult.reset_after == datetime.timedelta(seconds=-1)
+        assert limitresult.retry_after == datetime.timedelta(seconds=-1)
+
+    def test_ratelimit_first_time_seeing_key(self, limiterf):
+        """Verify our behaviour for the first time we see a key."""
+        rate = helpers.new_quota(
+            period=datetime.timedelta(seconds=60), count=50
+        )
+        limiter = limiterf.limiter
+        apply_lua = limiterf.apply_lua
+        apply_lua.reset_mock()
+        apply_lua.return_value = (0, 49, "-1", "-1")
+        check_lua = limiterf.check_lua
+        check_lua.return_value = (0, 49, "-1", "-1")
+
+        _ = limiter.rate_limit(key="key", rate=rate, quantity=1)
+
+        apply_lua.assert_called_once_with(
+            keys=["key"],
+            args=[
+                rate.limit,
+                rate.count / rate.period.total_seconds(),
+                rate.period.total_seconds(),
+                1,
+            ],
+        )
+
+        _ = limiter.rate_limit(key="key", rate=rate, quantity=0)
+
+        check_lua.assert_called_once_with(
+            keys=["key"],
+            args=[
+                rate.limit,
+                rate.count / rate.period.total_seconds(),
+                rate.period.total_seconds(),
+            ],
+        )
+
+    def test_ratelimit_exceeded(self, limiterf):
+        """Verify that if we're inside our cell, we reduce remaining."""
+        rate = helpers.new_quota(
+            period=datetime.timedelta(seconds=60), count=50
+        )
+        limiter = limiterf.limiter
+        apply_lua = limiterf.apply_lua
+        apply_lua.reset_mock()
+        apply_lua.return_value = (1, 0, "10", "15")
+
+        limitresult = limiter.rate_limit(key="key", rate=rate, quantity=1)
+
+        apply_lua.assert_called_once_with(
+            keys=["key"],
+            args=[
+                rate.limit,
+                rate.count / rate.period.total_seconds(),
+                rate.period.total_seconds(),
+                1,
+            ],
+        )
+        assert limitresult.limited is True
+        assert limitresult.remaining == 0

test/unit/test_redis_store.py

@@ -4,6 +4,7 @@
 import mock
 import pytest
 import redis
+import redis.client
 import rfc3986
 
 from rush import exceptions as rexc
@@ -103,3 +104,49 @@ def test_current_time_uses_redis_time(self):
             2018, 12, 25, 12, 27, 7, 608_937, tzinfo=datetime.timezone.utc
         )
         client.time.assert_called_once_with()
+
+    def test_compare_and_set(self):
+        """Verify we call the right things for atomic operation."""
+        pipeline = mock.MagicMock(autospec=redis.client.Pipeline)
+        pipeline.__enter__.return_value = pipeline
+        client = mock.Mock()
+        client.pipeline.return_value = pipeline
+        pipeline.hgetall.return_value = {}
+        data = limit_data.LimitData(used=5, remaining=10)
+        store = redstore.RedisStore(url="redis://", client=client)
+
+        store.compare_and_swap("key", old=None, new=data)
+        assert client.pipeline.call_count == 1
+        pipeline.watch.assert_called_once_with("key")
+        pipeline.hgetall.assert_called_once_with("key")
+        pipeline.multi.assert_called_once_with()
+        pipeline.hmset.assert_called_once_with("key", data.asdict())
+        pipeline.execute.assert_called_once_with()
+
+    def test_compare_and_set_raises_mismatched_data_error(self):
+        """Verify we avoid setting data when we can't verify the old state."""
+        pipeline = mock.MagicMock(autospec=redis.client.Pipeline)
+        pipeline.__enter__.return_value = pipeline
+        client = mock.Mock()
+        client.pipeline.return_value = pipeline
+        pipeline.hgetall.return_value = {"used": 5, "remaining": 10}
+        old = limit_data.LimitData(used=4, remaining=11)
+        new = limit_data.LimitData(used=5, remaining=10)
+        store = redstore.RedisStore(url="redis://", client=client)
+
+        with pytest.raises(rexc.MismatchedDataError):
+            store.compare_and_swap("key", old=old, new=new)
+
+    def test_compare_and_set_raises_data_changed_in_store_error(self):
+        """Verify we avoid setting data when we can't verify the old state."""
+        pipeline = mock.MagicMock(autospec=redis.client.Pipeline)
+        pipeline.__enter__.return_value = pipeline
+        client = mock.Mock()
+        client.pipeline.return_value = pipeline
+        pipeline.hgetall.side_effect = redis.WatchError("'key' changed")
+        old = limit_data.LimitData(used=4, remaining=11)
+        new = limit_data.LimitData(used=5, remaining=10)
+        store = redstore.RedisStore(url="redis://", client=client)
+
+        with pytest.raises(rexc.DataChangedInStoreError):
+            store.compare_and_swap("key", old=old, new=new)