From ddc51c4ace0caa0787715801b9df42e65c790d46 Mon Sep 17 00:00:00 2001
From: Bar Shaul <88437685+barshaul@users.noreply.github.com>
Date: Thu, 23 Dec 2021 11:42:30 +0200
Subject: Support for specifying error types with retry (#1817)

---
 tests/test_retry.py | 125 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 124 insertions(+), 1 deletion(-)

(limited to 'tests/test_retry.py')

diff --git a/tests/test_retry.py b/tests/test_retry.py
index c4650bc..0094787 100644
--- a/tests/test_retry.py
+++ b/tests/test_retry.py
@@ -1,10 +1,20 @@
+from unittest.mock import patch
+
 import pytest
 
 from redis.backoff import NoBackoff
+from redis.client import Redis
 from redis.connection import Connection, UnixDomainSocketConnection
-from redis.exceptions import ConnectionError
+from redis.exceptions import (
+    BusyLoadingError,
+    ConnectionError,
+    ReadOnlyError,
+    TimeoutError,
+)
 from redis.retry import Retry
 
+from .conftest import _get_client
+
 
 class BackoffMock:
     def __init__(self):
@@ -39,6 +49,37 @@ class TestConnectionConstructorWithRetry:
         assert isinstance(c.retry, Retry)
         assert c.retry._retries == retries
 
+    @pytest.mark.parametrize("Class", [Connection, UnixDomainSocketConnection])
+    def test_retry_on_error(self, Class):
+        c = Class(retry_on_error=[ReadOnlyError])
+        assert c.retry_on_error == [ReadOnlyError]
+        assert isinstance(c.retry, Retry)
+        assert c.retry._retries == 1
+
+    @pytest.mark.parametrize("Class", [Connection, UnixDomainSocketConnection])
+    def test_retry_on_error_empty_value(self, Class):
+        c = Class(retry_on_error=[])
+        assert c.retry_on_error == []
+        assert isinstance(c.retry, Retry)
+        assert c.retry._retries == 0
+
+    @pytest.mark.parametrize("Class", [Connection, UnixDomainSocketConnection])
+    def test_retry_on_error_and_timeout(self, Class):
+        c = Class(
+            retry_on_error=[ReadOnlyError, BusyLoadingError], retry_on_timeout=True
+        )
+        assert c.retry_on_error == [ReadOnlyError, BusyLoadingError, TimeoutError]
+        assert isinstance(c.retry, Retry)
+        assert c.retry._retries == 1
+
+    @pytest.mark.parametrize("retries", range(10))
+    @pytest.mark.parametrize("Class", [Connection, UnixDomainSocketConnection])
+    def test_retry_on_error_retry(self, Class, retries):
+        c = Class(retry_on_error=[ReadOnlyError], retry=Retry(NoBackoff(), retries))
+        assert c.retry_on_error == [ReadOnlyError]
+        assert isinstance(c.retry, Retry)
+        assert c.retry._retries == retries
+
 
 class TestRetry:
     "Test that Retry calls backoff and retries the expected number of times"
@@ -65,3 +106,85 @@ class TestRetry:
         assert self.actual_failures == 1 + retries
         assert backoff.reset_calls == 1
         assert backoff.calls == retries
+
+
+@pytest.mark.onlynoncluster
+class TestRedisClientRetry:
+    "Test the standalone Redis client behavior with retries"
+
+    def test_client_retry_on_error_with_success(self, request):
+        with patch.object(Redis, "parse_response") as parse_response:
+
+            def mock_parse_response(connection, *args, **options):
+                def ok_response(connection, *args, **options):
+                    return "MOCK_OK"
+
+                parse_response.side_effect = ok_response
+                raise ReadOnlyError()
+
+            parse_response.side_effect = mock_parse_response
+            r = _get_client(Redis, request, retry_on_error=[ReadOnlyError])
+            assert r.get("foo") == "MOCK_OK"
+            assert parse_response.call_count == 2
+
+    def test_client_retry_on_error_raise(self, request):
+        with patch.object(Redis, "parse_response") as parse_response:
+            parse_response.side_effect = BusyLoadingError()
+            retries = 3
+            r = _get_client(
+                Redis,
+                request,
+                retry_on_error=[ReadOnlyError, BusyLoadingError],
+                retry=Retry(NoBackoff(), retries),
+            )
+            with pytest.raises(BusyLoadingError):
+                try:
+                    r.get("foo")
+                finally:
+                    assert parse_response.call_count == retries + 1
+
+    def test_client_retry_on_error_different_error_raised(self, request):
+        with patch.object(Redis, "parse_response") as parse_response:
+            parse_response.side_effect = TimeoutError()
+            retries = 3
+            r = _get_client(
+                Redis,
+                request,
+                retry_on_error=[ReadOnlyError],
+                retry=Retry(NoBackoff(), retries),
+            )
+            with pytest.raises(TimeoutError):
+                try:
+                    r.get("foo")
+                finally:
+                    assert parse_response.call_count == 1
+
+    def test_client_retry_on_error_and_timeout(self, request):
+        with patch.object(Redis, "parse_response") as parse_response:
+            parse_response.side_effect = TimeoutError()
+            retries = 3
+            r = _get_client(
+                Redis,
+                request,
+                retry_on_error=[ReadOnlyError],
+                retry_on_timeout=True,
+                retry=Retry(NoBackoff(), retries),
+            )
+            with pytest.raises(TimeoutError):
+                try:
+                    r.get("foo")
+                finally:
+                    assert parse_response.call_count == retries + 1
+
+    def test_client_retry_on_timeout(self, request):
+        with patch.object(Redis, "parse_response") as parse_response:
+            parse_response.side_effect = TimeoutError()
+            retries = 3
+            r = _get_client(
+                Redis, request, retry_on_timeout=True, retry=Retry(NoBackoff(), retries)
+            )
+            with pytest.raises(TimeoutError):
+                try:
+                    r.get("foo")
+                finally:
+                    assert parse_response.call_count == retries + 1
-- 
cgit v1.2.1