Repository: redis/redis-py
File: tests/test_connection.py (Python)

Code before:
        def mock_connect():
            # connect only on the last retry
            if conn._connect.call_count <= 2:
                raise socket.timeout
            else:
                return origin_connect()

        conn._connect.side_effect = mock_connect
        conn.connect()
        assert conn._connect.call_count == 3
        self.clear(conn)

    def test_connect_without_retry_on_non_retryable_error(self):
        """Test that the _connect function is not being retried in case of a non-retryable error"""
        with patch.object(Connection, "_connect") as _connect:
            _connect.side_effect = RedisError("")
            conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
            with pytest.raises(RedisError):
                conn.connect()
            assert _connect.call_count == 1
            self.clear(conn)

    def test_connect_with_retries(self):
        """
        Test that the _connect function is not being retried in case of a CancelledError -
        error that is not in the list of retry-able errors
        """
        with patch.object(socket.socket, "sendall") as sendall:
            sendall.side_effect = OSError(ECONNREFUSED)
            conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
            with pytest.raises(ConnectionError):
                conn.connect()
            # the handshake commands are the failing ones
            # validate that we don't execute too many commands on each retry
            # 3 retries --> 3 commands
            assert sendall.call_count == 3

    def test_connect_timeout_error_without_retry(self):
        """Test that the _connect function is not being retried if retry_on_timeout is
        set to False"""
        conn = Connection(retry_on_timeout=False)
        conn._connect = mock.Mock()
        conn._connect.side_effect = socket.timeout

        with pytest.raises(TimeoutError, match="Timeout connecting to server"):
            conn.connect()
        assert conn._connect.call_count == 1
        self.clear(conn)


@pytest.mark.onlynoncluster

Review comment (suggestion):
The docstring describes testing for CancelledError (a non-retryable error), but the test actually validates retry behavior with OSError during the handshake phase. The docstring should be updated to accurately describe what the test does: validating that retries occur for the entire connect+handshake flow when OSError happens during handshake.
```suggestion
        Validate that retries occur for the entire connect+handshake flow when OSError
        happens during the handshake phase.
```

Code after:
        def mock_connect():
            # connect only on the last retry
            if conn._connect.call_count <= 2:
                raise socket.timeout
            else:
                return origin_connect()

        conn._connect.side_effect = mock_connect
        conn.connect()
        assert conn._connect.call_count == 3
        self.clear(conn)

    def test_connect_without_retry_on_non_retryable_error(self):
        """Test that the _connect function is not being retried in case of a non-retryable error"""
        with patch.object(Connection, "_connect") as _connect:
            _connect.side_effect = RedisError("")
            conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
            with pytest.raises(RedisError):
                conn.connect()
            assert _connect.call_count == 1
            self.clear(conn)

    def test_connect_with_retries(self):
        """
        Validate that retries occur for the entire connect+handshake flow when OSError
        happens during the handshake phase.
        """
        with patch.object(socket.socket, "sendall") as sendall:
            sendall.side_effect = OSError(ECONNREFUSED)
            conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
            with pytest.raises(ConnectionError):
                conn.connect()
            # the handshake commands are the failing ones
            # validate that we don't execute too many commands on each retry
            # 3 retries --> 3 commands
            assert sendall.call_count == 3

    def test_connect_timeout_error_without_retry(self):
        """Test that the _connect function is not being retried if retry_on_timeout is
        set to False"""
        conn = Connection(retry_on_timeout=False)
        conn._connect = mock.Mock()
        conn._connect.side_effect = socket.timeout

        with pytest.raises(TimeoutError, match="Timeout connecting to server"):
            conn.connect()
        assert conn._connect.call_count == 1
        self.clear(conn)


@pytest.mark.onlynoncluster