Add disconnect_on_error for redis-py 4.5.5 support

cunla · May 9, 2023 · 21d68bb · 21d68bb
1 parent 620e22e
commit 21d68bb
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -53,7 +53,7 @@ jobs:
       matrix:
         redis-image: [ "redis:6.2.12", "redis:7.0.11" ]
         python-version: [ "3.7", "3.10" ]
-        redis-py: [ "4.3.6", "4.5.4" ]
+        redis-py: [ "4.3.6", "4.5.5" ]
         include:
           - python-version: "3.11"
             redis-image: "redis:6.2.12"

diff --git a/fakeredis/_server.py b/fakeredis/_server.py
@@ -91,11 +91,13 @@ def _decode(self, response):
         else:
             return response
 
-    def read_response(self, disable_decoding=False):
+    def read_response(self, disable_decoding=False, disconnect_on_error=True):
         if not self._server.connected:
             try:
                 response = self._sock.responses.get_nowait()
             except queue.Empty:
+                if disconnect_on_error:
+                    self.disconnect()
                 raise redis.ConnectionError(msgs.CONNECTION_ERROR_MSG)
         else:
             response = self._sock.responses.get()

diff --git a/fakeredis/aioredis.py b/fakeredis/aioredis.py
@@ -150,6 +150,8 @@ async def read_response(self, **kwargs):
             try:
                 response = self._sock.responses.get_nowait()
             except asyncio.QueueEmpty:
+                if kwargs.get('disconnect_on_error', True):
+                    await self.disconnect()
                 raise redis_async.ConnectionError(msgs.CONNECTION_ERROR_MSG)
         else:
             timeout = kwargs.pop('timeout', None)