Merge pull request #669 from dano/fix-tty-streams2

Fix handling output from tty-enabled containers.

Author: shin-

docker/client.py

@@ -40,28 +40,7 @@ def attach(self, container, stdout=True, stderr=True,
         u = self._url("/containers/{0}/attach".format(container))
         response = self._post(u, params=params, stream=stream)
 
-        # Stream multi-plexing was only introduced in API v1.6. Anything before
-        # that needs old-style streaming.
-        if utils.compare_version('1.6', self._version) < 0:
-            def stream_result():
-                self._raise_for_status(response)
-                for line in response.iter_lines(chunk_size=1,
-                                                decode_unicode=True):
-                    # filter out keep-alive new lines
-                    if line:
-                        yield line
-
-            return stream_result() if stream else \
-                self._result(response, binary=True)
-
-        sep = bytes() if six.PY3 else str()
-
-        if stream:
-            return self._multiplexed_response_stream_helper(response)
-        else:
-            return sep.join(
-                [x for x in self._multiplexed_buffer_helper(response)]
-            )
+        return self._get_result(container, stream, response)
 
     @check_resource
     def attach_socket(self, container, params=None, ws=False):
@@ -363,17 +342,7 @@ def exec_start(self, exec_id, detach=False, tty=False, stream=False):
 
         res = self._post_json(self._url('/exec/{0}/start'.format(exec_id)),
                               data=data, stream=stream)
-        self._raise_for_status(res)
-        if stream:
-            return self._multiplexed_response_stream_helper(res)
-        elif six.PY3:
-            return bytes().join(
-                [x for x in self._multiplexed_buffer_helper(res)]
-            )
-        else:
-            return str().join(
-                [x for x in self._multiplexed_buffer_helper(res)]
-            )
+        return self._get_result_tty(stream, res, tty)
 
     @check_resource
     def export(self, container):
@@ -588,16 +557,7 @@ def logs(self, container, stdout=True, stderr=True, stream=False,
                 params['tail'] = tail
             url = self._url("/containers/{0}/logs".format(container))
             res = self._get(url, params=params, stream=stream)
-            if stream:
-                return self._multiplexed_response_stream_helper(res)
-            elif six.PY3:
-                return bytes().join(
-                    [x for x in self._multiplexed_buffer_helper(res)]
-                )
-            else:
-                return str().join(
-                    [x for x in self._multiplexed_buffer_helper(res)]
-                )
+            return self._get_result(container, stream, res)
         return self.attach(
             container,
             stdout=stdout,

docker/clientbase.py

@@ -221,6 +221,46 @@ def _multiplexed_response_stream_helper(self, response):
                 break
             yield data
 
+    def _stream_raw_result_old(self, response):
+        ''' Stream raw output for API versions below 1.6 '''
+        self._raise_for_status(response)
+        for line in response.iter_lines(chunk_size=1,
+                                        decode_unicode=True):
+            # filter out keep-alive new lines
+            if line:
+                yield line
+
+    def _stream_raw_result(self, response):
+        ''' Stream result for TTY-enabled container above API 1.6 '''
+        self._raise_for_status(response)
+        for out in response.iter_content(chunk_size=1, decode_unicode=True):
+            yield out
+
+    def _get_result(self, container, stream, res):
+        cont = self.inspect_container(container)
+        return self._get_result_tty(stream, res, cont['Config']['Tty'])
+
+    def _get_result_tty(self, stream, res, is_tty):
+        # Stream multi-plexing was only introduced in API v1.6. Anything
+        # before that needs old-style streaming.
+        if utils.compare_version('1.6', self._version) < 0:
+            return self._stream_raw_result_old(res)
+
+        # We should also use raw streaming (without keep-alives)
+        # if we're dealing with a tty-enabled container.
+        if is_tty:
+            return self._stream_raw_result(res) if stream else \
+                self._result(res, binary=True)
+
+        self._raise_for_status(res)
+        sep = six.binary_type()
+        if stream:
+            return self._multiplexed_response_stream_helper(res)
+        else:
+            return sep.join(
+                [x for x in self._multiplexed_buffer_helper(res)]
+            )
+
     def get_adapter(self, url):
         try:
             return super(ClientBase, self).get_adapter(url)

tests/fake_api.py

@@ -129,11 +129,11 @@ def post_fake_create_container():
     return status_code, response
 
 
-def get_fake_inspect_container():
+def get_fake_inspect_container(tty=False):
     status_code = 200
     response = {
         'Id': FAKE_CONTAINER_ID,
-        'Config': {'Privileged': True},
+        'Config': {'Privileged': True, 'Tty': tty},
         'ID': FAKE_CONTAINER_ID,
         'Image': 'busybox:latest',
         "State": {

tests/test.py

@@ -69,6 +69,14 @@ def fake_resolve_authconfig(authconfig, registry=None):
     return None
 
 
+def fake_inspect_container(self, container, tty=False):
+    return fake_api.get_fake_inspect_container(tty=tty)[1]
+
+
+def fake_inspect_container_tty(self, container):
+    return fake_inspect_container(self, container, tty=True)
+
+
 def fake_resp(url, data=None, **kwargs):
     status_code, content = fake_api.fake_responses[url]()
     return response(status_code=status_code, content=content)
@@ -1546,7 +1554,9 @@ def test_url_compatibility_tcp(self):
 
     def test_logs(self):
         try:
-            logs = self.client.logs(fake_api.FAKE_CONTAINER_ID)
+            with mock.patch('docker.Client.inspect_container',
+                            fake_inspect_container):
+                logs = self.client.logs(fake_api.FAKE_CONTAINER_ID)
         except Exception as e:
             self.fail('Command should not raise exception: {0}'.format(e))
 
@@ -1565,7 +1575,9 @@ def test_logs(self):
 
     def test_logs_with_dict_instead_of_id(self):
         try:
-            logs = self.client.logs({'Id': fake_api.FAKE_CONTAINER_ID})
+            with mock.patch('docker.Client.inspect_container',
+                            fake_inspect_container):
+                logs = self.client.logs({'Id': fake_api.FAKE_CONTAINER_ID})
         except Exception as e:
             self.fail('Command should not raise exception: {0}'.format(e))
 
@@ -1584,7 +1596,9 @@ def test_logs_with_dict_instead_of_id(self):
 
     def test_log_streaming(self):
         try:
-            self.client.logs(fake_api.FAKE_CONTAINER_ID, stream=True)
+            with mock.patch('docker.Client.inspect_container',
+                            fake_inspect_container):
+                self.client.logs(fake_api.FAKE_CONTAINER_ID, stream=True)
         except Exception as e:
             self.fail('Command should not raise exception: {0}'.format(e))
 
@@ -1598,7 +1612,10 @@ def test_log_streaming(self):
 
     def test_log_tail(self):
         try:
-            self.client.logs(fake_api.FAKE_CONTAINER_ID, stream=False, tail=10)
+            with mock.patch('docker.Client.inspect_container',
+                            fake_inspect_container):
+                self.client.logs(fake_api.FAKE_CONTAINER_ID, stream=False,
+                                 tail=10)
         except Exception as e:
             self.fail('Command should not raise exception: {0}'.format(e))
 
@@ -1610,6 +1627,27 @@ def test_log_tail(self):
             stream=False
         )
 
+    def test_log_tty(self):
+        try:
+            m = mock.Mock()
+            with mock.patch('docker.Client.inspect_container',
+                            fake_inspect_container_tty):
+                with mock.patch('docker.Client._stream_raw_result',
+                                m):
+                    self.client.logs(fake_api.FAKE_CONTAINER_ID,
+                                     stream=True)
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        self.assertTrue(m.called)
+        fake_request.assert_called_with(
+            url_prefix + 'containers/3cc2351ab11b/logs',
+            params={'timestamps': 0, 'follow': 1, 'stderr': 1, 'stdout': 1,
+                    'tail': 'all'},
+            timeout=DEFAULT_TIMEOUT_SECONDS,
+            stream=True
+        )
+
     def test_diff(self):
         try:
             self.client.diff(fake_api.FAKE_CONTAINER_ID)