test: exercise real compliance and deletion rate limiters (#43198)

api/tests/unit_tests/services/account/test_adapters.py

@@ -312,10 +312,12 @@ def test_verification_gateway_creates_six_digit_account_bound_challenge() -> Non
     assert generate_token.call_args.kwargs["additional_data"] == {"code": "123456"}
 
 
-def test_verification_notifier_preserves_rate_limit_before_enqueuing_email() -> None:
-    limiter = MagicMock(spec=RateLimiter)
-    limiter.is_rate_limited.return_value = True
-    limiter.time_window = 60
+def test_verification_notifier_preserves_rate_limit_before_enqueuing_email(
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+) -> None:
+    redis, commands = redis_transport
+    commands.return_value = 10
+    limiter = RateLimiter("deletion", 10, 60, redis_client=redis)
     notifier = CeleryAccountDeletionVerificationNotifier(rate_limiter=limiter)
 
     with (
@@ -324,5 +326,6 @@ def test_verification_notifier_preserves_rate_limit_before_enqueuing_email() ->
     ):
         notifier.send(email="account@example.com", code="123456")
 
+    assert [call.args[0] for call in commands.call_args_list] == ["ZREMRANGEBYSCORE", "ZCARD"]
     assert error.value.retry_after_minutes == 1
     mail_task.delay.assert_not_called()

api/tests/unit_tests/services/test_compliance_download_service.py

@@ -2,8 +2,10 @@ from unittest.mock import MagicMock
 
 import pytest
 
+from extensions.ext_redis import RedisClientWrapper
+from libs.helper import RateLimiter
 from machinery.context import RequestContext
-from services.compliance_download_service import ComplianceDownloadRateLimiter, ComplianceDownloadService
+from services.compliance_download_service import ComplianceDownloadService
 from services.errors.billing import BillingUpstreamUnavailableError, ComplianceRateLimitExceededError
 
 
@@ -13,10 +15,9 @@ def fetch_link() -> MagicMock:
 
 
 @pytest.fixture
-def rate_limiter() -> MagicMock:
-    limiter = MagicMock(spec=ComplianceDownloadRateLimiter)
-    limiter.is_rate_limited.return_value = False
-    return limiter
+def rate_limiter(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> RateLimiter:
+    redis, _commands = redis_transport
+    return RateLimiter("compliance", 10, 60, redis_client=redis)
 
 
 @pytest.fixture
@@ -32,7 +33,7 @@ def request_context() -> RequestContext:
 @pytest.fixture
 def service(
     fetch_link: MagicMock,
-    rate_limiter: MagicMock,
+    rate_limiter: RateLimiter,
 ) -> ComplianceDownloadService:
     return ComplianceDownloadService(
         fetch_link=fetch_link,
@@ -44,12 +45,20 @@ def test_get_link_checks_limit_fetches_and_increments(
     service: ComplianceDownloadService,
     request_context: RequestContext,
     fetch_link: MagicMock,
-    rate_limiter: MagicMock,
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
 ) -> None:
     events: list[str] = []
-    rate_limiter.is_rate_limited.side_effect = lambda _key: events.append("check") or False
+    _, commands = redis_transport
+
+    def execute(command: str, *_args: object, **_kwargs: object) -> int:
+        if command == "ZCARD":
+            events.append("check")
+        elif command == "ZADD":
+            events.append("increment")
+        return 0
+
+    commands.side_effect = execute
     fetch_link.side_effect = lambda *_args: events.append("fetch") or {"url": "https://example.com/report"}
-    rate_limiter.increment_rate_limit.side_effect = lambda _key: events.append("increment")
 
     result = service.get_link(
         request_context=request_context,
@@ -60,7 +69,8 @@ def test_get_link_checks_limit_fetches_and_increments(
 
     assert result == {"url": "https://example.com/report"}
     assert events == ["check", "fetch", "increment"]
-    rate_limiter.is_rate_limited.assert_called_once_with("account-1:workspace-1")
+    assert [call.args[0] for call in commands.call_args_list] == ["ZREMRANGEBYSCORE", "ZCARD", "ZADD", "EXPIRE"]
+    assert all(call.args[1] == "compliance:account-1:workspace-1" for call in commands.call_args_list)
     fetch_link.assert_called_once_with(
         "SOC2_Type_II",
         "account-1",
@@ -68,16 +78,17 @@ def test_get_link_checks_limit_fetches_and_increments(
         "127.0.0.1",
         "test-agent",
     )
-    rate_limiter.increment_rate_limit.assert_called_once_with("account-1:workspace-1")
+    assert commands.call_args.args == ("EXPIRE", "compliance:account-1:workspace-1", 120)
 
 
 def test_get_link_rejects_rate_limited_request(
     service: ComplianceDownloadService,
     request_context: RequestContext,
     fetch_link: MagicMock,
-    rate_limiter: MagicMock,
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
 ) -> None:
-    rate_limiter.is_rate_limited.return_value = True
+    _, commands = redis_transport
+    commands.return_value = 10
 
     with pytest.raises(ComplianceRateLimitExceededError):
         service.get_link(
@@ -88,15 +99,16 @@ def test_get_link_rejects_rate_limited_request(
         )
 
     fetch_link.assert_not_called()
-    rate_limiter.increment_rate_limit.assert_not_called()
+    assert [call.args[0] for call in commands.call_args_list] == ["ZREMRANGEBYSCORE", "ZCARD"]
 
 
 def test_get_link_does_not_increment_after_fetch_failure(
     service: ComplianceDownloadService,
     request_context: RequestContext,
     fetch_link: MagicMock,
-    rate_limiter: MagicMock,
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
 ) -> None:
+    _, commands = redis_transport
     fetch_link.side_effect = BillingUpstreamUnavailableError
 
     with pytest.raises(BillingUpstreamUnavailableError):
@@ -107,4 +119,4 @@ def test_get_link_does_not_increment_after_fetch_failure(
             device_info="test-agent",
         )
 
-    rate_limiter.increment_rate_limit.assert_not_called()
+    assert [call.args[0] for call in commands.call_args_list] == ["ZREMRANGEBYSCORE", "ZCARD"]