import socket import pytest from redis.asyncio.retry import Retry from redis.asyncio.sentinel import SentinelManagedConnection from redis.backoff import NoBackoff from .compat import mock pytestmark = pytest.mark.asyncio async def test_connect_retry_on_timeout_error(): """Test that the _connect function is retried in case of a timeout""" connection_pool = mock.AsyncMock() connection_pool.get_master_address = mock.AsyncMock(return_value=("localhost", 6379)) conn = SentinelManagedConnection(retry_on_timeout=True, retry=Retry(NoBackoff(), 3), connection_pool=connection_pool) origin_connect = conn._connect conn._connect = mock.AsyncMock() async def mock_connect(): # connect only on the last retry if conn._connect.call_count <= 2: raise socket.timeout else: return await origin_connect() conn._connect.side_effect = mock_connect await conn.connect() assert conn._connect.call_count == 3