Merge pull request #31 from mogproject/topic-unittest-#30

add withBytesOutput method to TestCase closes #30

Author: mogproject

src/mog_commons/unittest.py

@@ -113,15 +113,15 @@ def assertRaisesMessage(self, expected_exception, expected_message, callable_obj
         self.assertEqual(msg, expected_message)
 
     @contextmanager
-    def withOutput(self):
+    def withOutput(self, buffer_type=StringBuffer):
         """
         Capture and suppress stdout and stderr.
         example:
             with self.withOutput() as (out, err):
                 (do main logic)
             (verify out.getvalue() or err.getvalue())
         """
-        new_out, new_err = StringBuffer(), StringBuffer()
+        new_out, new_err = buffer_type(), buffer_type()
         old_out, old_err = sys.stdout, sys.stderr
 
         try:
@@ -130,6 +130,18 @@ def withOutput(self):
         finally:
             sys.stdout, sys.stderr = old_out, old_err
 
+    @contextmanager
+    def withBytesOutput(self):
+        """
+        Capture and suppress stdout and stderr. The value is represented as bytes.
+        example:
+            with self.withBytesOutput() as (out, err):
+                (do main logic)
+            (verify out.getvalue() or err.getvalue())
+        """
+        with self.withOutput(six.BytesIO) as (out, err):
+            yield out, err
+
     @contextmanager
     def withAssertOutput(self, expected_stdout, expected_stderr, encoding='utf-8'):
         with self.withOutput() as (out, err):

tests/mog_commons/test_unittest.py

@@ -30,6 +30,14 @@ def test_assert_system_exit(self):
         self.assertSystemExit(123, lambda: sys.exit(123))
         self.assertSystemExit(234, lambda x: sys.exit(x), 234)
 
+    def test_with_bytes_output(self):
+        with self.withBytesOutput() as (out, err):
+            out.write(b'\xff\xfe')
+            out.write('あいうえお'.encode('utf-8'))
+            err.write(b'\xfd\xfc')
+        self.assertEqual(out.getvalue(), b'\xff\xfe' + 'あいうえお'.encode('utf-8'))
+        self.assertEqual(err.getvalue(), b'\xfd\xfc')
+
     def test_with_assert_output_file(self):
         def f(text):
             with self.withAssertOutputFile(os.path.join('tests', 'resources', 'utf8_ja.txt')) as out: