diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 9cd0d082a4..17383cc7d1 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1143,8 +1143,20 @@ async def check_health(self): ) async def _send_packed_command(self, command: Iterable[bytes]) -> None: - self._writer.writelines(command) - await self._writer.drain() + writer = self._writer + if writer is None: + raise ConnectionError("Connection closed while writing") + try: + writer.writelines(command) + await writer.drain() + except AttributeError as error: + if str(error) != "'NoneType' object has no attribute 'writelines'": + raise + raise ConnectionError("Connection closed while writing") from error + except TypeError as error: + if str(error) != "'NoneType' object is not callable": + raise + raise ConnectionError("Connection closed while writing") from error async def send_packed_command( self, command: Union[bytes, str, Iterable[bytes]], check_health: bool = True @@ -1164,8 +1176,7 @@ async def send_packed_command( self._send_packed_command(command), self.socket_timeout ) else: - self._writer.writelines(command) - await self._writer.drain() + await self._send_packed_command(command) except asyncio.TimeoutError: await self.disconnect(nowait=True) raise TimeoutError("Timeout writing to socket") from None @@ -1393,7 +1404,13 @@ def pack_commands(self, commands: Iterable[Iterable[EncodableT]]) -> List[bytes] def _socket_is_empty(self): """Check if the socket is empty""" - return len(self._reader._buffer) == 0 + reader = self._reader + if reader is None: + raise ConnectionError("Connection closed while reading") + try: + return len(reader._buffer) == 0 + except AttributeError as error: + raise ConnectionError("Connection closed while reading") from error async def process_invalidation_messages(self): while not self._socket_is_empty(): diff --git a/tests/test_asyncio/test_connection.py b/tests/test_asyncio/test_connection.py index c56b5d9dd4..834654d08e 100644 --- a/tests/test_asyncio/test_connection.py +++ b/tests/test_asyncio/test_connection.py @@ -83,6 +83,81 @@ def test_connection_default_parser_matches_default_protocol(): assert conn.protocol == 3 +@pytest.mark.parametrize("socket_timeout", [None, 1]) +async def test_closed_writer_during_write_is_connection_error(socket_timeout): + conn = Connection(socket_timeout=socket_timeout) + conn._reader = mock.Mock() + conn._writer = mock.Mock() + conn._writer.writelines.side_effect = AttributeError( + "'NoneType' object has no attribute 'writelines'" + ) + conn.disconnect = mock.AsyncMock() + + with pytest.raises(ConnectionError, match="Connection closed while writing"): + await conn.send_packed_command([b"PING\r\n"], check_health=False) + + conn.disconnect.assert_awaited_once_with(nowait=True) + conn._reader = None + conn._writer = None + + +async def test_invalid_packed_command_type_error_is_not_connection_error(): + conn = Connection() + conn._reader = mock.Mock() + conn._writer = mock.Mock() + conn._writer.writelines.side_effect = TypeError("invalid packed command") + conn.disconnect = mock.AsyncMock() + + with pytest.raises(TypeError, match="invalid packed command"): + await conn.send_packed_command([b"PING\r\n"], check_health=False) + + conn.disconnect.assert_awaited_once_with(nowait=True) + conn._reader = None + conn._writer = None + + +async def test_invalid_packed_command_attribute_error_is_not_connection_error(): + class InvalidCommand: + def __iter__(self): + raise AttributeError("invalid packed command") + + conn = Connection() + conn._reader = mock.Mock() + conn._writer = mock.Mock() + conn._writer.writelines.side_effect = lambda command: list(command) + conn.disconnect = mock.AsyncMock() + + with pytest.raises(AttributeError, match="invalid packed command"): + await conn.send_packed_command(InvalidCommand(), check_health=False) + + conn.disconnect.assert_awaited_once_with(nowait=True) + conn._reader = None + conn._writer = None + + +async def test_closed_writer_type_error_is_connection_error(): + conn = Connection() + conn._reader = mock.Mock() + conn._writer = mock.Mock() + conn._writer.writelines.side_effect = TypeError("'NoneType' object is not callable") + conn.disconnect = mock.AsyncMock() + + with pytest.raises(ConnectionError, match="Connection closed while writing"): + await conn.send_packed_command([b"PING\r\n"], check_health=False) + + conn.disconnect.assert_awaited_once_with(nowait=True) + conn._reader = None + conn._writer = None + + +async def test_closed_reader_during_invalidation_processing_is_connection_error(): + conn = Connection() + conn._reader = None + + with pytest.raises(ConnectionError, match="Connection closed while reading"): + await conn.process_invalidation_messages() + + @pytest.mark.parametrize( ("buffer", "eof", "expected"), [