Apply exarkun's fix to stdout pass-through

2 files changed, 29 insertions, 19 deletions

diff --git a/python/subunit/__init__.py b/python/subunit/__init__.py
index 7a04eb9..f0b0987 100644
--- a/python/subunit/__init__.py
+++ b/python/subunit/__init__.py
@@ -50,7 +50,7 @@ class TestProtocolServer(object):
     READING_FAILURE = 2
     READING_ERROR = 3
 
-    def __init__(self, client):
+    def __init__(self, client, stream=sys.stdout):
         """Create a TestProtocol server instance.
 
         client should be an object that provides
@@ -63,6 +63,7 @@ class TestProtocolServer(object):
         """
         self.state = TestProtocolServer.OUTSIDE_TEST
         self.client = client
+        self._stream = stream
 
     def _addError(self, offset, line):
         if (self.state == TestProtocolServer.TEST_STARTED and
@@ -135,17 +136,21 @@ class TestProtocolServer(object):
               self.state == TestProtocolServer.READING_ERROR):
             self._appendMessage(line)
         else:
-            cmd, rest = line.split(None, 1)
-            offset = len(cmd) + 1
-            cmd = cmd.strip(':')
-            if cmd in ('test', 'testing'):
-                self._startTest(offset, line)
-            elif cmd == 'error':
-                self._addError(offset, line)
-            elif cmd == 'failure':
-                self._addFailure(offset, line)
-            elif cmd in ('success', 'successful'):
-                self._addSuccess(offset, line)
+            parts = line.split(None, 1)
+            if len(parts) == 2:
+                cmd, rest = parts
+                offset = len(cmd) + 1
+                cmd = cmd.strip(':')
+                if cmd in ('test', 'testing'):
+                    self._startTest(offset, line)
+                elif cmd == 'error':
+                    self._addError(offset, line)
+                elif cmd == 'failure':
+                    self._addFailure(offset, line)
+                elif cmd in ('success', 'successful'):
+                    self._addSuccess(offset, line)
+                else:
+                    self.stdOutLineReceived(line)
             else:
                 self.stdOutLineReceived(line)
 
@@ -187,7 +192,7 @@ class TestProtocolServer(object):
             self.stdOutLineReceived(line)
 
     def stdOutLineReceived(self, line):
-        sys.stdout.write(line)
+        self._stream.write(line)
 
 
 class RemoteException(Exception):
diff --git a/python/subunit/tests/test_test_protocol.py b/python/subunit/tests/test_test_protocol.py
index 9950f8a..ba33665 100644
--- a/python/subunit/tests/test_test_protocol.py
+++ b/python/subunit/tests/test_test_protocol.py
@@ -110,6 +110,7 @@ class TestMockTestProtocolServer(unittest.TestCase):
         self.assertEqual(protocol.success_calls, [])
         self.assertEqual(protocol.start_calls, [])
 
+
 class TestTestImports(unittest.TestCase):
 
     def test_imports(self):
@@ -177,15 +178,10 @@ class TestTestProtocolServerPassThrough(unittest.TestCase):
 
     def setUp(self):
         from StringIO import StringIO
-        self.real_stdout = sys.stdout
         self.stdout = StringIO()
-        sys.stdout = self.stdout
         self.test = subunit.RemotedTestCase("old mcdonald")
         self.client = MockTestProtocolServerClient()
-        self.protocol = subunit.TestProtocolServer(self.client)
-
-    def tearDown(self):
-        sys.stdout = self.real_stdout
+        self.protocol = subunit.TestProtocolServer(self.client, self.stdout)
 
     def keywords_before_test(self):
         self.protocol.lineReceived("failure a\n")
@@ -307,6 +303,15 @@ class TestTestProtocolServerPassThrough(unittest.TestCase):
         self.assertEqual(self.client.error_calls, [])
         self.assertEqual(self.client.success_calls, [])
 
+    def test_stdout_passthrough(self):
+        """
+        Verify that lines received which cannot be interpreted as any
+        protocol action are passed through to sys.stdout.
+        """
+        bytes = "randombytes\n"
+        self.protocol.lineReceived(bytes)
+        self.assertEqual(self.stdout.getvalue(), bytes)
+
 
 class TestTestProtocolServerLostConnection(unittest.TestCase):