Fig bug in split_buffer where input was being discarded

Also, write some tests for it.

fig/cli/log_printer.py

@@ -6,6 +6,7 @@ from itertools import cycle
 
 from .multiplexer import Multiplexer
 from . import colors
+from .utils import split_buffer
 
 
 class LogPrinter(object):
@@ -33,7 +34,7 @@ class LogPrinter(object):
         prefix = color_fn(container.name + " | ")
         # Attach to container before log printer starts running
         line_generator = split_buffer(self._attach(container), '\n')
-        return (prefix + line for line in line_generator)
+        return (prefix + line.decode('utf-8') for line in line_generator)
 
     def _attach(self, container):
         params = {
@@ -44,21 +45,3 @@ class LogPrinter(object):
         params.update(self.attach_params)
         params = dict((name, 1 if value else 0) for (name, value) in list(params.items()))
         return container.attach(**params)
-
-def split_buffer(reader, separator):
-    """
-    Given a generator which yields strings and a separator string,
-    joins all input, splits on the separator and yields each chunk.
-    Requires that each input string is decodable as UTF-8.
-    """
-    buffered = ''
-
-    for data in reader:
-        lines = (buffered + data.decode('utf-8')).split(separator)
-        for line in lines[:-1]:
-            yield line + separator
-        if len(lines) > 1:
-            buffered = lines[-1]
-
-    if len(buffered) > 0:
-        yield buffered

fig/cli/utils.py

@@ -83,3 +83,28 @@ def mkdir(path, permissions=0o700):
 
 def docker_url():
     return os.environ.get('DOCKER_HOST')
+
+
+def split_buffer(reader, separator):
+    """
+    Given a generator which yields strings and a separator string,
+    joins all input, splits on the separator and yields each chunk.
+
+    Unlike string.split(), each chunk includes the trailing
+    separator, except for the last one if none was found on the end
+    of the input.
+    """
+    buffered = str('')
+    separator = str(separator)
+
+    for data in reader:
+        buffered += data
+        while True:
+            index = buffered.find(separator)
+            if index == -1:
+                break
+            yield buffered[:index+1]
+            buffered = buffered[index+1:]
+
+    if len(buffered) > 0:
+        yield buffered

tests/split_buffer_test.py

@@ -0,0 +1,37 @@
+from __future__ import unicode_literals
+from __future__ import absolute_import
+from fig.cli.utils import split_buffer
+from . import unittest
+
+class SplitBufferTest(unittest.TestCase):
+    def test_single_line_chunks(self):
+        def reader():
+            yield "abc\n"
+            yield "def\n"
+            yield "ghi\n"
+
+        self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "def\n", "ghi\n"])
+
+    def test_no_end_separator(self):
+        def reader():
+            yield "abc\n"
+            yield "def\n"
+            yield "ghi"
+
+        self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "def\n", "ghi"])
+
+    def test_multiple_line_chunk(self):
+        def reader():
+            yield "abc\ndef\nghi"
+
+        self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "def\n", "ghi"])
+
+    def test_chunked_line(self):
+        def reader():
+            yield "a"
+            yield "b"
+            yield "c"
+            yield "\n"
+            yield "d"
+
+        self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "d"])