Move socket-reading test helpers into docker.utils.socket

Signed-off-by: Aanand Prasad <aanand.prasad@gmail.com>

4 files changed, 58 insertions, 50 deletions

diff --git a/docker/utils/socket.py b/docker/utils/socket.py
new file mode 100644
index 0000000..f81d2f5
--- /dev/null
+++ b/docker/utils/socket.py
@@ -0,0 +1,49 @@
+import errno
+import os
+import select
+import struct
+
+import six
+
+
+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 data"
+        data += next_data
+    return data
diff --git a/tests/helpers.py b/tests/helpers.py
index 70be803..94ea388 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -1,9 +1,6 @@
-import errno
 import os
 import os.path
-import select
 import shutil
-import struct
 import tarfile
 import tempfile
 import unittest
@@ -67,49 +64,6 @@ 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 data"
-        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 56b648a..594aaa3 100644
--- a/tests/integration/container_test.py
+++ b/tests/integration/container_test.py
@@ -3,6 +3,8 @@ import signal
 import tempfile
 
 import docker
+from docker.utils.socket import next_packet_size
+from docker.utils.socket import read_data
 import pytest
 import six
 
@@ -1025,9 +1027,9 @@ class AttachContainerTest(helpers.BaseTestCase):
 
         self.client.start(ident)
 
-        next_size = helpers.next_packet_size(pty_stdout)
+        next_size = next_packet_size(pty_stdout)
         self.assertEqual(next_size, len(line))
-        data = helpers.read_data(pty_stdout, next_size)
+        data = read_data(pty_stdout, next_size)
         self.assertEqual(data.decode('utf-8'), line)
 
 
diff --git a/tests/integration/exec_test.py b/tests/integration/exec_test.py
index 9f54808..d0c8c9b 100644
--- a/tests/integration/exec_test.py
+++ b/tests/integration/exec_test.py
@@ -1,5 +1,8 @@
 import pytest
 
+from docker.utils.socket import next_packet_size
+from docker.utils.socket import read_data
+
 from .. import helpers
 
 BUSYBOX = helpers.BUSYBOX
@@ -107,9 +110,9 @@ class ExecTest(helpers.BaseTestCase):
         socket = self.client.exec_start(exec_id, socket=True)
         self.addCleanup(socket.close)
 
-        next_size = helpers.next_packet_size(socket)
+        next_size = next_packet_size(socket)
         self.assertEqual(next_size, len(line))
-        data = helpers.read_data(socket, next_size)
+        data = read_data(socket, next_size)
         self.assertEqual(data.decode('utf-8'), line)
 
     def test_exec_inspect(self):