diff --git a/api/tests/unit_tests/services/account/test_adapters.py b/api/tests/unit_tests/services/account/test_adapters.py index 759747a6e5d..b92a2106da6 100644 --- a/api/tests/unit_tests/services/account/test_adapters.py +++ b/api/tests/unit_tests/services/account/test_adapters.py @@ -1,6 +1,6 @@ from datetime import UTC, datetime from hashlib import sha256 -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest @@ -259,9 +259,14 @@ def test_security_gateway_counts_normal_ip_request(redis_transport: tuple[RedisC assert gateway.is_ip_limited("127.0.0.1") is False - commands.assert_any_call("SETEX", "email_send_ip_limit_minute:127.0.0.1", 60, 1) - commands.assert_any_call("EXPIRE", "email_send_ip_limit_minute:127.0.0.1", 60) - assert commands.call_count == 4 + freeze_key = "email_send_ip_limit_freeze:127.0.0.1" + minute_key = "email_send_ip_limit_minute:127.0.0.1" + assert commands.call_args_list == [ + call("GET", freeze_key, keys=[freeze_key]), + call("GET", minute_key, keys=[minute_key]), + call("SETEX", minute_key, 60, 1), + call("EXPIRE", minute_key, 60), + ] def test_security_gateway_freezes_second_over_limit_ip_strike( @@ -278,8 +283,15 @@ def test_security_gateway_freezes_second_over_limit_ip_strike( assert gateway.is_ip_limited("127.0.0.1") is True - commands.assert_any_call("SETEX", "email_send_ip_limit_freeze:127.0.0.1", 3600, 1) - assert commands.call_count == 4 + freeze_key = "email_send_ip_limit_freeze:127.0.0.1" + minute_key = "email_send_ip_limit_minute:127.0.0.1" + hour_key = "email_send_ip_limit_hour:127.0.0.1" + assert commands.call_args_list == [ + call("GET", freeze_key, keys=[freeze_key]), + call("GET", minute_key, keys=[minute_key]), + call("GET", hour_key, keys=[hour_key]), + call("SETEX", freeze_key, 3600, 1), + ] def test_verification_gateway_binds_token_to_the_target_account() -> None: diff --git a/api/tests/unit_tests/services/account/test_email_registration_adapters.py b/api/tests/unit_tests/services/account/test_email_registration_adapters.py index 3a3c533cabb..c403e032533 100644 --- a/api/tests/unit_tests/services/account/test_email_registration_adapters.py +++ b/api/tests/unit_tests/services/account/test_email_registration_adapters.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, Mock, call, patch import pytest from redis import RedisError @@ -79,12 +79,23 @@ def test_security_gateway_preserves_shared_ip_limit( ) assert gateway.is_ip_limited("127.0.0.1") is limited + freeze_key = "email_send_ip_limit_freeze:127.0.0.1" + minute_key = "email_send_ip_limit_minute:127.0.0.1" + hour_key = "email_send_ip_limit_hour:127.0.0.1" if limited: - commands.assert_any_call("SET", "email_send_ip_limit_hour:127.0.0.1", 1, "NX", "EX", 600) + assert commands.call_args_list == [ + call("GET", freeze_key, keys=[freeze_key]), + call("GET", minute_key, keys=[minute_key]), + call("GET", hour_key, keys=[hour_key]), + call("SET", hour_key, 1, "NX", "EX", 600), + ] else: - commands.assert_any_call("SETEX", "email_send_ip_limit_minute:127.0.0.1", 60, count + 1) - commands.assert_any_call("EXPIRE", "email_send_ip_limit_minute:127.0.0.1", 60) - assert commands.call_count == 4 + assert commands.call_args_list == [ + call("GET", freeze_key, keys=[freeze_key]), + call("GET", minute_key, keys=[minute_key]), + call("SETEX", minute_key, 60, count + 1), + call("EXPIRE", minute_key, 60), + ] def test_security_gateway_uses_registration_and_login_keys( @@ -106,10 +117,13 @@ def test_security_gateway_uses_registration_and_login_keys( gateway.reset_verification_failures("user@example.com") gateway.reset_login_failures("user@example.com") - commands.assert_any_call("SETEX", "email_register_error_rate_limit:user@example.com", 600, 2) - commands.assert_any_call("DEL", "email_register_error_rate_limit:user@example.com") + key = "email_register_error_rate_limit:user@example.com" + assert commands.call_args_list == [ + call("GET", key, keys=[key]), + call("SETEX", key, 600, 2), + call("DEL", key), + ] reset_login_error_rate_limit.assert_called_once_with("user@example.com") - assert commands.call_count == 3 @pytest.mark.parametrize(("count", "limited"), [(5, False), (6, True)]) diff --git a/api/tests/unit_tests/services/account/test_login_adapters.py b/api/tests/unit_tests/services/account/test_login_adapters.py index 3138115cd3e..7484ec4cdd3 100644 --- a/api/tests/unit_tests/services/account/test_login_adapters.py +++ b/api/tests/unit_tests/services/account/test_login_adapters.py @@ -3,8 +3,9 @@ from __future__ import annotations import logging from collections.abc import Callable from dataclasses import dataclass, field +from datetime import timedelta from typing import override -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest from sqlalchemy import select @@ -67,14 +68,13 @@ def test_security_gateway_owns_login_failure_state(redis_transport: tuple[RedisC gateway.record_login_failure("user@example.com") gateway.reset_login_failures("user@example.com") - commands.assert_any_call( - "SETEX", - "login_error_rate_limit:user@example.com", - adapters.dify_config.LOGIN_LOCKOUT_DURATION, - 3, - ) - commands.assert_any_call("DEL", "login_error_rate_limit:user@example.com") - assert commands.call_count == 4 + key = "login_error_rate_limit:user@example.com" + assert commands.call_args_list == [ + call("GET", key, keys=[key]), + call("GET", key, keys=[key]), + call("SETEX", key, adapters.dify_config.LOGIN_LOCKOUT_DURATION, 3), + call("DEL", key), + ] def test_security_gateway_owns_email_send_ip_limit( @@ -86,8 +86,15 @@ def test_security_gateway_owns_email_send_ip_limit( gateway = adapters.RedisConsoleAuthSecurityGateway(redis=redis) assert gateway.is_email_send_ip_limited("127.0.0.1") is True - commands.assert_any_call("SET", "email_send_ip_limit_hour:127.0.0.1", 1, "NX", "EX", 600) - assert commands.call_count == 4 + freeze_key = "email_send_ip_limit_freeze:127.0.0.1" + minute_key = "email_send_ip_limit_minute:127.0.0.1" + hour_key = "email_send_ip_limit_hour:127.0.0.1" + assert commands.call_args_list == [ + call("GET", freeze_key, keys=[freeze_key]), + call("GET", minute_key, keys=[minute_key]), + call("GET", hour_key, keys=[hour_key]), + call("SET", hour_key, 1, "NX", "EX", 600), + ] def test_session_gateway_owns_refresh_token_storage( @@ -110,9 +117,11 @@ def test_session_gateway_owns_refresh_token_storage( assert result == AccountSessionTokens(access_token="access", refresh_token="refresh", csrf_token="csrf") assert issued_payloads[0]["user_id"] == "account-1" - assert commands.call_args_list[0].args[1] == "refresh_token:refresh" - assert commands.call_args_list[1].args[1] == "account_refresh_token:account-1" - assert commands.call_count == 2 + expires_in = int(timedelta(days=adapters.dify_config.REFRESH_TOKEN_EXPIRE_DAYS).total_seconds()) + assert commands.call_args_list == [ + call("SETEX", "refresh_token:refresh", expires_in, "account-1"), + call("SETEX", "account_refresh_token:account-1", expires_in, "refresh"), + ] def test_session_gateway_resolves_rotates_and_revokes_refresh_tokens(