Improve coverages for _DuplicateWriter

Author: ikedam

tests/testsuite.py

@@ -9,15 +9,18 @@
 
 from xmlrunner.unittest import unittest
 import xmlrunner
+from xmlrunner.result import _DuplicateWriter
 from xmlrunner.result import _XMLTestResult
 import doctest
 import tests.doctest_example
 from six import StringIO, BytesIO
 from tempfile import mkdtemp
+from tempfile import mkstemp
 from shutil import rmtree
 from glob import glob
 from xml.dom import minidom
 from lxml import etree
+import os
 import os.path
 
 
@@ -573,3 +576,48 @@ def test_xmlrunner_hold_traceback(self):
         runner = self._test_xmlrunner(suite)
         countAfterTest = sys.getrefcount(self.DummyRefCountTest.dummy)
         self.assertEqual(countBeforeTest, countAfterTest)
+
+
+class DuplicateWriterTestCase(unittest.TestCase):
+    def setUp(self):
+        fd, self.file = mkstemp()
+        self.fh = os.fdopen(fd, 'w')
+        self.buffer = StringIO()
+        self.writer = _DuplicateWriter(self.fh, self.buffer)
+
+    def tearDown(self):
+        self.buffer.close()
+        self.fh.close()
+        os.unlink(self.file)
+
+    def getFirstContent(self):
+        with open(self.file, 'r') as f:
+            return f.read()
+
+    def getSecondContent(self):
+        return self.buffer.getvalue()
+
+    def test_flush(self):
+        self.writer.write('foobarbaz')
+        self.writer.flush()
+        self.assertEqual(self.getFirstContent(), self.getSecondContent())
+
+    def test_writable(self):
+        self.assertTrue(self.writer.writable())
+
+    def test_writelines(self):
+        self.writer.writelines([
+            'foo\n',
+            'bar\n',
+            'baz\n',
+        ])
+        self.writer.flush()
+        self.assertEqual(self.getFirstContent(), self.getSecondContent())
+
+    def test_write(self):
+        # try long buffer (1M)
+        buffer = 'x' * (1024 * 1024)
+        wrote = self.writer.write(buffer)
+        self.writer.flush()
+        self.assertEqual(self.getFirstContent(), self.getSecondContent())
+        self.assertEqual(wrote, len(self.getSecondContent()))

xmlrunner/result.py

@@ -83,7 +83,7 @@ def testcase_name(test_method):
     return result
 
 
-class _DuplicateWriter(io.RawIOBase):
+class _DuplicateWriter(io.TextIOBase):
     """
     Duplicate output from the first handle to the second handle
 
@@ -107,11 +107,20 @@ def writelines(self, lines):
         self._second.writelines(lines)
 
     def write(self, b):
-        len = self._first.write(b)
+        if isinstance(self._first, io.TextIOBase):
+            wrote = self._first.write(b)
 
-        # expected to always succeed to write
-        self._second.write(b[:len])
-        return len
+            if wrote is not None:
+                # expected to always succeed to write
+                self._second.write(b[:wrote])
+
+            return wrote
+        else:
+            # file-like object in Python2
+            # It doesn't return wrote bytes.
+            self._first.write(b)
+            self._second.write(b)
+            return len(b)
 
 
 class _TestInfo(object):