diff --git a/api/tests/unit_tests/services/account/test_adapters.py b/api/tests/unit_tests/services/account/test_adapters.py index 553d7b0f1a5..759747a6e5d 100644 --- a/api/tests/unit_tests/services/account/test_adapters.py +++ b/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() diff --git a/api/tests/unit_tests/services/test_compliance_download_service.py b/api/tests/unit_tests/services/test_compliance_download_service.py index 43071564778..653e196d07d 100644 --- a/api/tests/unit_tests/services/test_compliance_download_service.py +++ b/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"]