Merge pull request #858 from TomasTomecek/improve-exec-api

Improve exec api

6 files changed, 94 insertions, 54 deletions

diff --git a/Makefile b/Makefile
index 4a887e6..22abf50 100644
--- a/Makefile
+++ b/Makefile
@@ -31,6 +31,7 @@ integration-test-py3: build-py3
 	docker run -v /var/run/docker.sock:/var/run/docker.sock docker-py3 py.test tests/integration
 
 integration-dind: build build-py3
+	docker rm -vf dpy-dind || :
 	docker run -d --name dpy-dind --env="DOCKER_HOST=tcp://localhost:2375" --privileged dockerswarm/dind:1.9.0 docker -d -H tcp://0.0.0.0:2375
 	docker run --env="DOCKER_HOST=tcp://docker:2375" --link=dpy-dind:docker docker-py py.test tests/integration
 	docker run --env="DOCKER_HOST=tcp://docker:2375" --link=dpy-dind:docker docker-py3 py.test tests/integration
diff --git a/docker/api/exec_api.py b/docker/api/exec_api.py
index e64a40c..f0e4afa 100644
--- a/docker/api/exec_api.py
+++ b/docker/api/exec_api.py
@@ -7,8 +7,8 @@ from .. import utils
 class ExecApiMixin(object):
     @utils.minimum_version('1.15')
     @utils.check_resource
-    def exec_create(self, container, cmd, stdout=True, stderr=True, tty=False,
-                    privileged=False, user=''):
+    def exec_create(self, container, cmd, stdout=True, stderr=True,
+                    stdin=False, tty=False, privileged=False, user=''):
         if privileged and utils.compare_version('1.19', self._version) < 0:
             raise errors.InvalidVersion(
                 'Privileged exec is not supported in API < 1.19'
@@ -25,7 +25,7 @@ class ExecApiMixin(object):
             'User': user,
             'Privileged': privileged,
             'Tty': tty,
-            'AttachStdin': False,
+            'AttachStdin': stdin,
             'AttachStdout': stdout,
             'AttachStderr': stderr,
             'Cmd': cmd
@@ -53,7 +53,11 @@ class ExecApiMixin(object):
         self._raise_for_status(res)
 
     @utils.minimum_version('1.15')
-    def exec_start(self, exec_id, detach=False, tty=False, stream=False):
+    def exec_start(self, exec_id, detach=False, tty=False, stream=False,
+                   socket=False):
+        # we want opened socket if socket == True
+        if socket:
+            stream = True
         if isinstance(exec_id, dict):
             exec_id = exec_id.get('Id')
 
@@ -65,4 +69,7 @@ class ExecApiMixin(object):
         res = self._post_json(
             self._url('/exec/{0}/start', exec_id), data=data, stream=stream
         )
+
+        if socket:
+            return self._get_raw_response_socket(res)
         return self._get_result_tty(stream, res, tty)
diff --git a/docker/utils/__init__.py b/docker/utils/__init__.py
index bcbf568..e86b724 100644
--- a/docker/utils/__init__.py
+++ b/docker/utils/__init__.py
@@ -4,7 +4,7 @@ from .utils import (
     kwargs_from_env, convert_filters, datetime_to_timestamp, create_host_config,
     create_container_config, parse_bytes, ping_registry, parse_env_file,
     version_lt, version_gte, decode_json_header, split_command,
-) # flake8: noqa
+)  # flake8: noqa
 
 from .types import Ulimit, LogConfig # flake8: noqa
-from .decorators import check_resource, minimum_version #flake8: noqa
+from .decorators import check_resource, minimum_version  # flake8: noqa
diff --git a/tests/helpers.py b/tests/helpers.py
index 01083e3..eca5035 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -1,6 +1,9 @@
+import errno
 import os
 import os.path
+import select
 import shutil
+import struct
 import tarfile
 import tempfile
 import unittest
@@ -64,6 +67,49 @@ def docker_client_kwargs(**kwargs):
     return client_kwargs
 
 
+def read_socket(socket, n=4096):
+    """ Code stolen from dockerpty to read the socket """
+    recoverable_errors = (errno.EINTR, errno.EDEADLK, errno.EWOULDBLOCK)
+
+    # wait for data to become available
+    select.select([socket], [], [])
+
+    try:
+        if hasattr(socket, 'recv'):
+            return socket.recv(n)
+        return os.read(socket.fileno(), n)
+    except EnvironmentError as e:
+        if e.errno not in recoverable_errors:
+            raise
+
+
+def next_packet_size(socket):
+    """ Code stolen from dockerpty to get the next packet size """
+    data = six.binary_type()
+    while len(data) < 8:
+        next_data = read_socket(socket, 8 - len(data))
+        if not next_data:
+            return 0
+        data = data + next_data
+
+    if data is None:
+        return 0
+
+    if len(data) == 8:
+        _, actual = struct.unpack('>BxxxL', data)
+        return actual
+
+
+def read_data(socket, packet_size):
+    data = six.binary_type()
+    while len(data) < packet_size:
+        next_data = read_socket(socket, packet_size - len(data))
+        if not next_data:
+            assert False, "Failed trying to read in the dataz"
+        data += next_data
+    return data
+
+
 class BaseTestCase(unittest.TestCase):
     tmp_imgs = []
     tmp_containers = []
diff --git a/tests/integration/container_test.py b/tests/integration/container_test.py
index 8cb4e97..0396514 100644
--- a/tests/integration/container_test.py
+++ b/tests/integration/container_test.py
@@ -1,8 +1,6 @@
-import errno
 import os
 import shutil
 import signal
-import struct
 import tempfile
 
 import docker
@@ -950,14 +948,15 @@ class AttachContainerTest(helpers.BaseTestCase):
         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)
+        self.client.start(id)
         sock = self.client.attach_socket(container, ws=False)
         self.assertTrue(sock.fileno() > -1)
 
     def test_run_container_reading_socket(self):
         line = 'hi there and stuff and things, words!'
-        command = "echo '{0}'".format(line)
+        # `echo` appends CRLF, `printf` doesn't
+        command = "printf '{0}'".format(line)
         container = self.client.create_container(BUSYBOX, command,
                                                  detach=True, tty=False)
         ident = container['Id']
@@ -965,51 +964,14 @@ class AttachContainerTest(helpers.BaseTestCase):
 
         opts = {"stdout": 1, "stream": 1, "logs": 1}
         pty_stdout = self.client.attach_socket(ident, opts)
+        self.addCleanup(pty_stdout.close)
+
         self.client.start(ident)
 
-        recoverable_errors = (errno.EINTR, errno.EDEADLK, errno.EWOULDBLOCK)
-
-        def read(n=4096):
-            """Code stolen from dockerpty to read the socket"""
-            try:
-                if hasattr(pty_stdout, 'recv'):
-                    return pty_stdout.recv(n)
-                return os.read(pty_stdout.fileno(), n)
-            except EnvironmentError as e:
-                if e.errno not in recoverable_errors:
-                    raise
-
-        def next_packet_size():
-            """Code stolen from dockerpty to get the next packet size"""
-            data = six.binary_type()
-            while len(data) < 8:
-                next_data = read(8 - len(data))
-                if not next_data:
-                    return 0
-                data = data + next_data
-
-            if data is None:
-                return 0
-
-            if len(data) == 8:
-                _, actual = struct.unpack('>BxxxL', data)
-                return actual
-
-        next_size = next_packet_size()
-        self.assertEqual(next_size, len(line) + 1)
-
-        data = six.binary_type()
-        while len(data) < next_size:
-            next_data = read(next_size - len(data))
-            if not next_data:
-                assert False, "Failed trying to read in the dataz"
-            data += next_data
-        self.assertEqual(data.decode('utf-8'), "{0}\n".format(line))
-        pty_stdout.close()
-
-        # Prevent segfault at the end of the test run
-        if hasattr(pty_stdout, "_response"):
-            del pty_stdout._response
+        next_size = helpers.next_packet_size(pty_stdout)
+        self.assertEqual(next_size, len(line))
+        data = helpers.read_data(pty_stdout, next_size)
+        self.assertEqual(data.decode('utf-8'), line)
 
 
 class PauseTest(helpers.BaseTestCase):
diff --git a/tests/integration/exec_test.py b/tests/integration/exec_test.py
index 6025316..9f54808 100644
--- a/tests/integration/exec_test.py
+++ b/tests/integration/exec_test.py
@@ -77,8 +77,8 @@ class ExecTest(helpers.BaseTestCase):
         container = self.client.create_container(BUSYBOX, 'cat',
                                                  detach=True, stdin_open=True)
         id = container['Id']
-        self.client.start(id)
         self.tmp_containers.append(id)
+        self.client.start(id)
 
         exec_id = self.client.exec_create(id, ['echo', 'hello\nworld'])
         self.assertIn('Id', exec_id)
@@ -88,6 +88,30 @@ class ExecTest(helpers.BaseTestCase):
             res += chunk
         self.assertEqual(res, b'hello\nworld\n')
 
+    def test_exec_start_socket(self):
+        if not helpers.exec_driver_is_native():
+            pytest.skip('Exec driver not native')
+
+        container = self.client.create_container(BUSYBOX, 'cat',
+                                                 detach=True, stdin_open=True)
+        container_id = container['Id']
+        self.client.start(container_id)
+        self.tmp_containers.append(container_id)
+
+        line = 'yay, interactive exec!'
+        # `echo` appends CRLF, `printf` doesn't
+        exec_id = self.client.exec_create(
+            container_id, ['printf', line], tty=True)
+        self.assertIn('Id', exec_id)
+
+        socket = self.client.exec_start(exec_id, socket=True)
+        self.addCleanup(socket.close)
+
+        next_size = helpers.next_packet_size(socket)
+        self.assertEqual(next_size, len(line))
+        data = helpers.read_data(socket, next_size)
+        self.assertEqual(data.decode('utf-8'), line)
+
     def test_exec_inspect(self):
         if not helpers.exec_driver_is_native():
             pytest.skip('Exec driver not native')