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

Signed-off-by: Aanand Prasad <aanand.prasad@gmail.com>
This commit is contained in:
Aanand Prasad 2016-07-12 16:23:13 -04:00
parent e64ba8f2b9
commit 73f06e3335
4 changed files with 58 additions and 50 deletions

49
docker/utils/socket.py Normal file
View File

@ -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

View File

@ -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 = []

View File

@ -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)

View File

@ -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):