diff --git a/remoulade/rate_limits/backends/redis.py b/remoulade/rate_limits/backends/redis.py index 80fdb57f5..8348321ea 100644 --- a/remoulade/rate_limits/backends/redis.py +++ b/remoulade/rate_limits/backends/redis.py @@ -25,8 +25,10 @@ class RedisBackend(RateLimitBackend): """Redis backend using ``limits`` RedisStorage. Parameters: - url(str): A redis connection URL. If both a URL and + url(str): An optional connection URL. If both a URL and connection parameters are provided, the URL is used. + client(Redis): An optional client. If this is passed, + then all other parameters are ignored. key_prefix(str): A prefix to prepend to all keys used for rate limiting. strategy(str): The rate limiting strategy to use. One of ``fixed_window``, ``moving_window``, or ``sliding_window``. @@ -41,15 +43,21 @@ class RedisBackend(RateLimitBackend): def __init__( self, *, - url: str, + url: str | None = None, + client=None, key_prefix: str = "remoulade-rate-limit:", strategy: str = "sliding_window", **parameters, ): super().__init__() + if client is None and url is None: + raise ValueError("Either url or client must be provided") + storage: Storage - if "sentinel" in url: + if client: + storage = RedisStorage("redis://localhost", connection_pool=client.connection_pool, key_prefix=key_prefix) + elif "sentinel" in url: storage = RedisSentinelStorage(url, key_prefix=key_prefix, **parameters) else: storage = RedisStorage(url, key_prefix=key_prefix, **parameters) diff --git a/tests/conftest.py b/tests/conftest.py index 1cbd25c11..2f872421a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -174,10 +174,40 @@ def scheduler_thread(): @pytest.fixture -def rate_limit_backend(): +def stub_rate_limit_backend(): return rl_backends.StubBackend() +@pytest.fixture +def redis_rate_limit_backend(): + redis_url = os.getenv("REMOULADE_TEST_REDIS_URL") or "redis://localhost:6481/0" + client = redis.Redis.from_url(redis_url) + check_redis(client) + return rl_backends.RedisBackend(url=redis_url) + + +@pytest.fixture +def redis_rate_limit_backend_from_client(): + redis_url = os.getenv("REMOULADE_TEST_REDIS_URL") or "redis://localhost:6481/0" + client = redis.Redis.from_url(redis_url) + check_redis(client) + return rl_backends.RedisBackend(client=client) + + +@pytest.fixture +def rate_limit_backends(stub_rate_limit_backend, redis_rate_limit_backend, redis_rate_limit_backend_from_client): + return { + "stub": stub_rate_limit_backend, + "redis": redis_rate_limit_backend, + "redis_client": redis_rate_limit_backend_from_client, + } + + +@pytest.fixture(params=["stub", "redis", "redis_client"]) +def rate_limit_backend(request, rate_limit_backends): + return rate_limit_backends[request.param] + + @pytest.fixture def redis_result_backend(): redis_url = os.getenv("REMOULADE_TEST_REDIS_URL") or "redis://localhost:6481/0" diff --git a/tests/middleware/test_rate_limit.py b/tests/middleware/test_rate_limit.py index 566cb1ef3..04701ae12 100644 --- a/tests/middleware/test_rate_limit.py +++ b/tests/middleware/test_rate_limit.py @@ -79,3 +79,10 @@ def do_work(): def test_rate_limit_unknown_strategy(): with pytest.raises(ValueError, match="Unknown rate limit strategy"): StubBackend(strategy="unknown") + + +def test_redis_backend_requires_url_or_client(): + from remoulade.rate_limits.backends import RedisBackend + + with pytest.raises(ValueError, match="Either url or client must be provided"): + RedisBackend()