From e8472be6d5b21246303f2a78eb1cdfeb62ea202a Mon Sep 17 00:00:00 2001
From: Aanand Prasad <aanand.prasad@gmail.com>
Date: Wed, 22 Jan 2014 17:44:04 +0000
Subject: [PATCH] Fig bug in split_buffer where input was being discarded

Also, write some tests for it.
---
 fig/cli/log_printer.py     | 21 ++-------------------
 fig/cli/utils.py           | 25 +++++++++++++++++++++++++
 tests/split_buffer_test.py | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 19 deletions(-)
 create mode 100644 tests/split_buffer_test.py

diff --git a/fig/cli/log_printer.py b/fig/cli/log_printer.py
index a3ad15eb3..0fe3215e6 100644
--- a/fig/cli/log_printer.py
+++ b/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
diff --git a/fig/cli/utils.py b/fig/cli/utils.py
index 2b0eb42d6..3116df98b 100644
--- a/fig/cli/utils.py
+++ b/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
diff --git a/tests/split_buffer_test.py b/tests/split_buffer_test.py
new file mode 100644
index 000000000..b90463c07
--- /dev/null
+++ b/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"])