Don't set socket timeout if it's already disabled when streaming

Signed-off-by: Kevin Frommelt <kevin.frommelt@gmail.com>

2 files changed, 65 insertions, 6 deletions

diff --git a/docker/client.py b/docker/client.py
index 5f60a32..de3cb3c 100644
--- a/docker/client.py
+++ b/docker/client.py
@@ -291,14 +291,29 @@ class Client(
         """ Depending on the combination of python version and whether we're
         connecting over http or https, we might need to access _sock, which
         may or may not exist; or we may need to just settimeout on socket
-         itself, which also may or may not have settimeout on it.
+        itself, which also may or may not have settimeout on it. To avoid
+        missing the correct one, we try both.
 
-        To avoid missing the correct one, we try both.
+        We also do not want to set the timeout if it is already disabled, as
+        you run the risk of changing a socket that was non-blocking to
+        blocking, for example when using gevent.
         """
-        if hasattr(socket, "settimeout"):
-            socket.settimeout(None)
-        if hasattr(socket, "_sock") and hasattr(socket._sock, "settimeout"):
-            socket._sock.settimeout(None)
+        sockets = [socket, getattr(socket, '_sock', None)]
+
+        for s in sockets:
+            if not hasattr(s, 'settimeout'):
+                continue
+
+            timeout = -1
+
+            if hasattr(s, 'gettimeout'):
+                timeout = s.gettimeout()
+
+            # Don't change the timeout if it is already disabled.
+            if timeout is None or timeout == 0.0:
+                continue
+
+            s.settimeout(None)
 
     def _get_result(self, container, stream, res):
         cont = self.inspect_container(container)
diff --git a/tests/unit/client_test.py b/tests/unit/client_test.py
index 1a173b5..b21f1d6 100644
--- a/tests/unit/client_test.py
+++ b/tests/unit/client_test.py
@@ -24,3 +24,47 @@ class ClientTest(base.BaseTestCase):
                           DOCKER_TLS_VERIFY='1')
         client = Client.from_env()
         self.assertEqual(client.base_url, "https://192.168.59.103:2376")
+
+
+class DisableSocketTest(base.BaseTestCase):
+    class DummySocket(object):
+        def __init__(self, timeout=60):
+            self.timeout = timeout
+
+        def settimeout(self, timeout):
+            self.timeout = timeout
+
+        def gettimeout(self):
+            return self.timeout
+
+    def setUp(self):
+        self.client = Client()
+
+    def test_disable_socket_timeout(self):
+        """Test that the timeout is disabled on a generic socket object."""
+        socket = self.DummySocket()
+
+        self.client._disable_socket_timeout(socket)
+
+        self.assertEqual(socket.timeout, None)
+
+    def test_disable_socket_timeout2(self):
+        """Test that the timeouts are disabled on a generic socket object
+        and it's _sock object if present."""
+        socket = self.DummySocket()
+        socket._sock = self.DummySocket()
+
+        self.client._disable_socket_timeout(socket)
+
+        self.assertEqual(socket.timeout, None)
+        self.assertEqual(socket._sock.timeout, None)
+
+    def test_disable_socket_timout_non_blocking(self):
+        """Test that a non-blocking socket does not get set to blocking."""
+        socket = self.DummySocket()
+        socket._sock = self.DummySocket(0.0)
+
+        self.client._disable_socket_timeout(socket)
+
+        self.assertEqual(socket.timeout, None)
+        self.assertEqual(socket._sock.timeout, 0.0)