close socket after server disconnect (#1797)

Author: pawl

redis/connection.py

@@ -717,9 +717,14 @@ def disconnect(self):
         self._parser.on_disconnect()
         if self._sock is None:
             return
-        try:
-            if os.getpid() == self.pid:
+
+        if os.getpid() == self.pid:
+            try:
                 self._sock.shutdown(socket.SHUT_RDWR)
+            except OSError:
+                pass
+
+        try:
             self._sock.close()
         except OSError:
             pass

tests/test_connection.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from redis.connection import Connection
 from redis.exceptions import InvalidResponse
 from redis.utils import HIREDIS_AVAILABLE
 
@@ -40,3 +41,36 @@ def inner():
     # mod = j(modclient)
     # mod.set("fookey", ".", d)
     # assert mod.get('fookey') == d
+
+
+class TestConnection:
+    def test_disconnect(self):
+        conn = Connection()
+        mock_sock = mock.Mock()
+        conn._sock = mock_sock
+        conn.disconnect()
+        mock_sock.shutdown.assert_called_once()
+        mock_sock.close.assert_called_once()
+        assert conn._sock is None
+
+    def test_disconnect__shutdown_OSError(self):
+        """An OSError on socket shutdown will still close the socket."""
+        conn = Connection()
+        mock_sock = mock.Mock()
+        conn._sock = mock_sock
+        conn._sock.shutdown.side_effect = OSError
+        conn.disconnect()
+        mock_sock.shutdown.assert_called_once()
+        mock_sock.close.assert_called_once()
+        assert conn._sock is None
+
+    def test_disconnect__close_OSError(self):
+        """An OSError on socket close will still clear out the socket."""
+        conn = Connection()
+        mock_sock = mock.Mock()
+        conn._sock = mock_sock
+        conn._sock.close.side_effect = OSError
+        conn.disconnect()
+        mock_sock.shutdown.assert_called_once()
+        mock_sock.close.assert_called_once()
+        assert conn._sock is None