Added a health check response counter

Author: barshaul

redis/client.py

@@ -1282,15 +1282,11 @@ def __init__(
         self.encoder = encoder
         if self.encoder is None:
             self.encoder = self.connection_pool.get_encoder()
-        self.health_check_response_b = self.encoder.encode(
-            self.HEALTH_CHECK_MESSAGE)
+        self.health_check_response_b = self.encoder.encode(self.HEALTH_CHECK_MESSAGE)
         if self.encoder.decode_responses:
             self.health_check_response = ["pong", self.HEALTH_CHECK_MESSAGE]
         else:
-            self.health_check_response = [
-                b"pong",
-                self.health_check_response_b
-            ]
+            self.health_check_response = [b"pong", self.health_check_response_b]
         self.reset()
 
     def __enter__(self):
@@ -1315,6 +1311,7 @@ def reset(self):
             self.connection_pool.release(self.connection)
             self.connection = None
         self.channels = {}
+        self.health_check_response_counter = 0
         self.pending_unsubscribe_channels = set()
         self.patterns = {}
         self.pending_unsubscribe_patterns = set()
@@ -1370,12 +1367,19 @@ def clean_health_check_responses(self):
         """
         If any health check responses are present, clean them
         """
+        ttl = 10
         conn = self.connection
-        while self._execute(conn, conn.can_read, timeout=0):
-            response = self._execute(conn, conn.read_response)
-            if not self.is_health_check_response(response):
-                raise PubSubError('A non health check response was cleaned by '
-                                  'execute_command: {0}'.format(response))
+        while self.health_check_response_counter > 0 and ttl > 0:
+            if self._execute(conn, conn.can_read, timeout=1):
+                response = self._execute(conn, conn.read_response)
+                if self.is_health_check_response(response):
+                    self.health_check_response_counter -= 1
+                else:
+                    raise PubSubError(
+                        "A non health check response was cleaned by "
+                        "execute_command: {0}".format(response)
+                    )
+            ttl -= 1
 
     def _disconnect_raise_connect(self, conn, error):
         """
@@ -1418,6 +1422,7 @@ def parse_response(self, block=True, timeout=0):
 
         if self.is_health_check_response(response):
             # ignore the health check message as user might not expect it
+            self.health_check_response_counter -= 1
             return None
         return response
 
@@ -1428,9 +1433,9 @@ def is_health_check_response(self, response):
         bulk response, instead of a multi-bulk with "pong" and the response.
         """
         return response in [
-                self.health_check_response,  # If there was a subscription
-                self.health_check_response_b  # If there wasn't
-                ]
+            self.health_check_response,  # If there was a subscription
+            self.health_check_response_b,  # If there wasn't
+        ]
 
     def check_health(self):
         conn = self.connection
@@ -1442,6 +1447,7 @@ def check_health(self):
 
         if conn.health_check_interval and time.time() > conn.next_health_check:
             conn.send_command("PING", self.HEALTH_CHECK_MESSAGE, check_health=False)
+            self.health_check_response_counter += 1
 
     def _normalize_keys(self, data):
         """

tests/test_pubsub.py

@@ -2,25 +2,23 @@
 import threading
 import time
 from unittest import mock
+from unittest.mock import patch
 
 import pytest
 
 import redis
 from redis.exceptions import ConnectionError
 
-from .conftest import (
-    _get_client,
-    skip_if_redis_enterprise,
-    skip_if_server_version_lt
-)
+from .conftest import _get_client, skip_if_redis_enterprise, skip_if_server_version_lt
 
 
 def wait_for_message(pubsub, timeout=0.1, ignore_subscribe_messages=False):
     now = time.time()
     timeout = now + timeout
     while now < timeout:
         message = pubsub.get_message(
-            ignore_subscribe_messages=ignore_subscribe_messages)
+            ignore_subscribe_messages=ignore_subscribe_messages
+        )
         if message is not None:
             return message
         time.sleep(0.01)
@@ -351,15 +349,6 @@ def test_unicode_pattern_message_handler(self, r):
             "pmessage", channel, "test message", pattern=pattern
         )
 
-    def test_get_message_without_subscribe(self, r):
-        p = r.pubsub()
-        with pytest.raises(RuntimeError) as info:
-            p.get_message()
-        expect = (
-            "connection not set: " "did you forget to call subscribe() or psubscribe()?"
-        )
-        assert expect in info.exconly()
-
 
 class TestPubSubAutoDecoding:
     "These tests only validate that we get unicode values back"
@@ -557,7 +546,7 @@ def test_get_message_not_subscribed_return_none(self, r):
         assert p.subscribed is False
         assert p.get_message() is None
         assert p.get_message(timeout=0.1) is None
-        with patch.object(threading.Event, 'wait') as mock:
+        with patch.object(threading.Event, "wait") as mock:
             mock.return_value = False
             assert p.get_message(timeout=0.01) is None
             assert mock.called
@@ -570,19 +559,19 @@ def poll(ps, expected_res):
             message = ps.get_message(timeout=1)
             assert message == expected_res
 
-        subscribe_response = make_message('subscribe', 'foo', 1)
+        subscribe_response = make_message("subscribe", "foo", 1)
         poller = threading.Thread(target=poll, args=(p, subscribe_response))
         poller.start()
         time.sleep(0.2)
-        p.subscribe('foo')
+        p.subscribe("foo")
         poller.join()
 
     def test_get_message_wait_for_subscription_not_being_called(self, r):
         p = r.pubsub()
-        p.subscribe('foo')
-        with patch.object(threading.Event, 'wait') as mock:
+        p.subscribe("foo")
+        with patch.object(threading.Event, "wait") as mock:
             assert p.subscribed is True
-            assert wait_for_message(p) == make_message('subscribe', 'foo', 1)
+            assert wait_for_message(p) == make_message("subscribe", "foo", 1)
             assert mock.called is False