Fix infinite loop in Client.logs(stream=True) caused by connection close.

This commit is contained in:
Fedor Korotkiy 2013-12-24 09:17:04 +00:00
parent 20a2c23bc2
commit 09f95768f3
2 changed files with 25 additions and 1 deletions

View File

@ -260,6 +260,9 @@ class Client(requests.Session):
data = ''
while size > 0:
block = socket.recv(size)
if not block:
return None
data += block
size -= len(block)
return data
@ -272,7 +275,11 @@ class Client(requests.Session):
_, length = struct.unpack('>BxxxL', header)
if not length:
break
yield recvall(socket, length).rstrip()
data = recvall(socket, length)
if data:
yield data
else:
break
def attach(self, container):
socket = self.attach_socket(container)

View File

@ -274,6 +274,23 @@ class TestLogs(BaseTestCase):
logs = self.client.logs(id)
self.assertEqual(logs, snippet + '\n')
class TestLogsStreaming(BaseTestCase):
def runTest(self):
snippet = 'Flowering Nights (Sakuya Iyazoi)'
container = self.client.create_container(
'busybox', 'echo {0}'.format(snippet)
)
id = container['Id']
self.client.start(id)
self.tmp_containers.append(id)
logs = ''
for chunk in self.client.logs(id, stream=True):
logs += chunk
exitcode = self.client.wait(id)
self.assertEqual(exitcode, 0)
self.assertEqual(logs, snippet + '\n')
class TestLogsWithDictInsteadOfId(BaseTestCase):
def runTest(self):