Use retry mechanism in async version of Connection objects (#2271)

1 files changed, 51 insertions, 2 deletions

diff --git a/tests/test_asyncio/test_connection.py b/tests/test_asyncio/test_connection.py
index f6259ad..78a3efd 100644
--- a/tests/test_asyncio/test_connection.py
+++ b/tests/test_asyncio/test_connection.py
@@ -1,10 +1,18 @@
 import asyncio
+import socket
 import types
+from unittest.mock import patch
 
 import pytest
 
-from redis.asyncio.connection import PythonParser, UnixDomainSocketConnection
-from redis.exceptions import InvalidResponse
+from redis.asyncio.connection import (
+    Connection,
+    PythonParser,
+    UnixDomainSocketConnection,
+)
+from redis.asyncio.retry import Retry
+from redis.backoff import NoBackoff
+from redis.exceptions import ConnectionError, InvalidResponse, TimeoutError
 from redis.utils import HIREDIS_AVAILABLE
 from tests.conftest import skip_if_server_version_lt
 
@@ -60,3 +68,44 @@ async def test_socket_param_regression(r):
 async def test_can_run_concurrent_commands(r):
     assert await r.ping() is True
     assert all(await asyncio.gather(*(r.ping() for _ in range(10))))
+
+
+async def test_connect_retry_on_timeout_error():
+    """Test that the _connect function is retried in case of a timeout"""
+    conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 3))
+    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
+
+
+async def test_connect_without_retry_on_os_error():
+    """Test that the _connect function is not being retried in case of a OSError"""
+    with patch.object(Connection, "_connect") as _connect:
+        _connect.side_effect = OSError("")
+        conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
+        with pytest.raises(ConnectionError):
+            await conn.connect()
+        assert _connect.call_count == 1
+
+
+async def test_connect_timeout_error_without_retry():
+    """Test that the _connect function is not being retried if retry_on_timeout is
+    set to False"""
+    conn = Connection(retry_on_timeout=False)
+    conn._connect = mock.AsyncMock()
+    conn._connect.side_effect = socket.timeout
+
+    with pytest.raises(TimeoutError) as e:
+        await conn.connect()
+    assert conn._connect.call_count == 1
+    assert str(e.value) == "Timeout connecting to server"