diff --git a/zerver/lib/rate_limiter.py b/zerver/lib/rate_limiter.py index c086b1fe3e..6381058a72 100644 --- a/zerver/lib/rate_limiter.py +++ b/zerver/lib/rate_limiter.py @@ -1,5 +1,4 @@ import logging -import os import time from abc import ABC, abstractmethod from typing import Dict, List, Optional, Set, Tuple, Type, cast @@ -11,6 +10,7 @@ from django.conf import settings from django.http import HttpRequest from typing_extensions import override +from zerver.lib import redis_utils from zerver.lib.cache import cache_with_key from zerver.lib.exceptions import RateLimitedError from zerver.lib.redis_utils import get_redis_client @@ -22,8 +22,6 @@ from zerver.models import UserProfile client = get_redis_client() rules: Dict[str, List[Tuple[int, int]]] = settings.RATE_LIMITING_RULES -KEY_PREFIX = "" - logger = logging.getLogger(__name__) @@ -159,11 +157,6 @@ class RateLimitedIPAddr(RateLimitedObject): return rules[self.domain] -def bounce_redis_key_prefix_for_testing(test_name: str) -> None: - global KEY_PREFIX - KEY_PREFIX = test_name + ":" + str(os.getpid()) + ":" - - class RateLimiterBackend(ABC): @classmethod @abstractmethod @@ -323,7 +316,8 @@ class RedisRateLimiterBackend(RateLimiterBackend): @classmethod def get_keys(cls, entity_key: str) -> List[str]: return [ - f"{KEY_PREFIX}ratelimit:{entity_key}:{keytype}" for keytype in ["list", "zset", "block"] + f"{redis_utils.REDIS_KEY_PREFIX}ratelimit:{entity_key}:{keytype}" + for keytype in ["list", "zset", "block"] ] @classmethod diff --git a/zerver/lib/redis_utils.py b/zerver/lib/redis_utils.py index 16d9100bb0..deddca2a09 100644 --- a/zerver/lib/redis_utils.py +++ b/zerver/lib/redis_utils.py @@ -1,3 +1,4 @@ +import os import re import secrets from typing import Any, Dict, Mapping, Optional @@ -84,3 +85,11 @@ def validate_key_fits_format(key: str, key_format: str) -> None: if not re.fullmatch(regex, key): raise ZulipRedisKeyOfWrongFormatError(f"{key} does not match format {key_format}") + + +REDIS_KEY_PREFIX = "" + + +def bounce_redis_key_prefix_for_testing(test_name: str) -> None: + global REDIS_KEY_PREFIX + REDIS_KEY_PREFIX = test_name + ":" + str(os.getpid()) + ":" diff --git a/zerver/lib/test_classes.py b/zerver/lib/test_classes.py index 67cd1ad334..25d7f867e0 100644 --- a/zerver/lib/test_classes.py +++ b/zerver/lib/test_classes.py @@ -67,7 +67,7 @@ from zerver.lib.initial_password import initial_password from zerver.lib.message import access_message from zerver.lib.notification_data import UserMessageNotificationsData from zerver.lib.per_request_cache import flush_per_request_caches -from zerver.lib.rate_limiter import bounce_redis_key_prefix_for_testing +from zerver.lib.redis_utils import bounce_redis_key_prefix_for_testing from zerver.lib.sessions import get_session_dict_user from zerver.lib.soft_deactivation import do_soft_deactivate_users from zerver.lib.stream_subscription import get_subscribed_stream_ids_for_user