Added additional call count check for read_response

Author: vladvildanov

tests/test_asyncio/conftest.py

@@ -1,6 +1,6 @@
 import random
 from contextlib import asynccontextmanager as _asynccontextmanager
-from typing import Union
+from typing import Union, Iterator, Any
 from unittest.mock import AsyncMock
 
 import pytest
@@ -19,8 +19,9 @@
 
 
 class AwaitableMock(AsyncMock):
-    def __await__(self):
-        pass
+    def __await__(self) -> Iterator[Any]:
+        self.await_count += 1
+        return iter([])
 
 
 async def _get_info(redis_url):

tests/test_asyncio/test_connection.py

@@ -337,6 +337,7 @@ async def test_client_do_not_retry_write_on_read_failure(mock_connection, mock_p
 
     # If read from socket fails, writes won't be executed.
     mock_connection.send_command.assert_called_once_with("SET", "key", "value")
+    mock_connection.read_response.call_count = 3
 
 
 @pytest.mark.onlynoncluster
@@ -364,6 +365,7 @@ async def test_pipeline_immediate_do_not_retry_write_on_read_failure(
 
     # If read from socket fails, writes won't be executed.
     mock_connection.send_command.assert_called_once_with("SET", "key", "value")
+    mock_connection.read_response.call_count = 3
 
 
 async def test_close_is_aclose(request):

tests/test_connection.py

@@ -268,6 +268,7 @@ def test_client_do_not_retry_write_on_read_failure(mock_connection, mock_pool):
 
     # If read from socket fails, writes won't be executed.
     mock_connection.send_command.assert_called_once_with("SET", "key", "value")
+    mock_connection.read_response.call_count = 3
 
 
 @pytest.mark.onlynoncluster
@@ -291,6 +292,7 @@ def test_pipeline_immediate_do_not_retry_write_on_read_failure(
 
     # If read from socket fails, writes won't be executed.
     mock_connection.send_command.assert_called_once_with("SET", "key", "value")
+    mock_connection.read_response.call_count = 3
 
 
 @pytest.mark.skipif(sys.version_info == (3, 9), reason="Flacky test on Python 3.9")