Fix socket reading function for TCP (non-HTTPS) connections on Windows

Signed-off-by: Joffrey F <joffrey@docker.com>

2 files changed, 59 insertions, 2 deletions

diff --git a/docker/utils/socket.py b/docker/utils/socket.py
index 0945f0a..7b96d4f 100644
--- a/docker/utils/socket.py
+++ b/docker/utils/socket.py
@@ -1,6 +1,7 @@
 import errno
 import os
 import select
+import socket as pysocket
 import struct
 
 import six
@@ -28,6 +29,8 @@ def read(socket, n=4096):
     try:
         if hasattr(socket, 'recv'):
             return socket.recv(n)
+        if six.PY3 and isinstance(socket, getattr(pysocket, 'SocketIO')):
+            return socket.read(n)
         return os.read(socket.fileno(), n)
     except EnvironmentError as e:
         if e.errno not in recoverable_errors:
diff --git a/tests/unit/api_test.py b/tests/unit/api_test.py
index 46cbd68..ba80840 100644
--- a/tests/unit/api_test.py
+++ b/tests/unit/api_test.py
@@ -365,7 +365,7 @@ class DockerApiTest(BaseAPIClientTest):
         assert result == content
 
 
-class StreamTest(unittest.TestCase):
+class UnixSocketStreamTest(unittest.TestCase):
     def setUp(self):
         socket_dir = tempfile.mkdtemp()
         self.build_context = tempfile.mkdtemp()
@@ -462,7 +462,61 @@ class StreamTest(unittest.TestCase):
                         raise e
 
             assert list(stream) == [
-                str(i).encode() for i in range(50)]
+                str(i).encode() for i in range(50)
+            ]
+
+
+class TCPSocketStreamTest(unittest.TestCase):
+    text_data = b'''
+    Now, those children out there, they're jumping through the
+    flames in the hope that the god of the fire will make them fruitful.
+    Really, you can't blame them. After all, what girl would not prefer the
+    child of a god to that of some acne-scarred artisan?
+    '''
+
+    def setUp(self):
+
+        self.server = six.moves.socketserver.ThreadingTCPServer(
+            ('', 0), self.get_handler_class()
+        )
+        self.thread = threading.Thread(target=self.server.serve_forever)
+        self.thread.setDaemon(True)
+        self.thread.start()
+        self.address = 'http://{}:{}'.format(
+            socket.gethostname(), self.server.server_address[1]
+        )
+
+    def tearDown(self):
+        self.server.shutdown()
+        self.server.server_close()
+        self.thread.join()
+
+    def get_handler_class(self):
+        text_data = self.text_data
+
+        class Handler(six.moves.BaseHTTPServer.BaseHTTPRequestHandler, object):
+            def do_POST(self):
+                self.send_response(101)
+                self.send_header(
+                    'Content-Type', 'application/vnd.docker.raw-stream'
+                )
+                self.send_header('Connection', 'Upgrade')
+                self.send_header('Upgrade', 'tcp')
+                self.end_headers()
+                self.wfile.flush()
+                time.sleep(0.2)
+                self.wfile.write(text_data)
+                self.wfile.flush()
+
+        return Handler
+
+    def test_read_from_socket(self):
+        with APIClient(base_url=self.address) as client:
+            resp = client._post(client._url('/dummy'), stream=True)
+            data = client._read_from_socket(resp, stream=True, tty=True)
+            results = b''.join(data)
+
+        assert results == self.text_data
 
 
 class UserAgentTest(unittest.TestCase):