Keep unix socket alive with python3

2 files changed, 30 insertions, 12 deletions

diff --git a/docker/client.py b/docker/client.py
index 1aa9ed6..1c130ba 100644
--- a/docker/client.py
+++ b/docker/client.py
@@ -274,18 +274,19 @@ class Client(requests.Session):
     def _get_raw_response_socket(self, response):
         self._raise_for_status(response)
         if six.PY3:
-            sock = response.raw._fp.fp.raw._sock
+            sock = response.raw._fp.fp.raw
+            sock._response = response
         else:
             sock = response.raw._fp.fp._sock
-        try:
-            # Keep a reference to the response to stop it being garbage
-            # collected. If the response is garbage collected, it will close
-            # TLS sockets.
-            sock._response = response
-        except AttributeError:
-            # UNIX sockets can't have attributes set on them, but that's fine
-            # because we won't be doing TLS over them
-            pass
+            try:
+                # Keep a reference to the response to stop it being garbage
+                # collected. If the response is garbage collected, it will
+                # close TLS sockets.
+                sock._response = response
+            except AttributeError:
+                # UNIX sockets can't have attributes set on them, but that's
+                # fine because we won't be doing TLS over them
+                pass
 
         return sock
 
@@ -324,7 +325,10 @@ class Client(requests.Session):
         def recvall(socket, size):
             blocks = []
             while size > 0:
-                block = socket.recv(size)
+                if six.PY3:
+                    block = socket._sock.recv(size)
+                else:
+                    block = socket.recv(size)
                 if not block:
                     return None
 
@@ -336,7 +340,10 @@ class Client(requests.Session):
             return data
 
         while True:
-            socket.settimeout(None)
+            if six.PY3:
+                socket._sock.settimeout(None)
+            else:
+                socket.settimeout(None)
             header = recvall(socket, STREAM_HEADER_SIZE_BYTES)
             if not header:
                 break
diff --git a/tests/integration_test.py b/tests/integration_test.py
index 801dd82..f06fbcf 100644
--- a/tests/integration_test.py
+++ b/tests/integration_test.py
@@ -870,6 +870,17 @@ class TestExecuteCommandStreaming(BaseTestCase):
         self.assertEqual(res, expected)
 
 
+class TestRunContainerStreaming(BaseTestCase):
+    def runTest(self):
+        container = self.client.create_container('busybox', '/bin/sh',
+                                                 detach=True, stdin_open=True)
+        id = container['Id']
+        self.client.start(id)
+        self.tmp_containers.append(id)
+        socket = self.client.attach_socket(container, ws=False)
+        self.assertTrue(socket.fileno() > -1)
+
+
 class TestPauseUnpauseContainer(BaseTestCase):
     def runTest(self):
         container = self.client.create_container('busybox', ['sleep', '9999'])