Merge pull request #25 from mtreinish/add-value-error-to-write-only-stream-check

Check for ValueError in write only check in _unwrap_text()

2 files changed, 44 insertions, 1 deletions

diff --git a/python/subunit/__init__.py b/python/subunit/__init__.py
index 69acef5..dd3e951 100644
--- a/python/subunit/__init__.py
+++ b/python/subunit/__init__.py
@@ -1287,15 +1287,17 @@ def _make_binary_on_windows(fileno):
 
 def _unwrap_text(stream):
     """Unwrap stream if it is a text stream to get the original buffer."""
+    exceptions = (_UnsupportedOperation, IOError)
     if sys.version_info > (3, 0):
         unicode_type = str
     else:
         unicode_type = unicode
+        exceptions += (ValueError,)
     try:
         # Read streams
         if type(stream.read(0)) is unicode_type:
             return stream.buffer
-    except (_UnsupportedOperation, IOError):
+    except exceptions:
         # Cannot read from the stream: try via writes
         try:
             stream.write(_b(''))
diff --git a/python/subunit/tests/test_test_protocol.py b/python/subunit/tests/test_test_protocol.py
index f3cd0e3..7427b12 100644
--- a/python/subunit/tests/test_test_protocol.py
+++ b/python/subunit/tests/test_test_protocol.py
@@ -15,8 +15,11 @@
 #
 
 import datetime
+import io
 import unittest2 as unittest
 import os
+import sys
+import tempfile
 
 from testtools import PlaceHolder, skipIf, TestCase, TestResult
 from testtools.compat import _b, _u, BytesIO
@@ -52,6 +55,44 @@ def details_to_str(details):
     return TestResult()._err_details_to_string(None, details=details)
 
 
+class TestHelpers(TestCase):
+    def test__unwrap_text_file_read_mode(self):
+        fd, file_path = tempfile.mkstemp()
+        self.addCleanup(os.remove, file_path)
+        fake_file = os.fdopen(fd, 'r')
+        if sys.version_info > (3, 0):
+            self.assertEqual(fake_file.buffer,
+                             subunit._unwrap_text(fake_file))
+        else:
+            self.assertEqual(fake_file, subunit._unwrap_text(fake_file))
+
+    def test__unwrap_text_file_write_mode(self):
+        fd, file_path = tempfile.mkstemp()
+        self.addCleanup(os.remove, file_path)
+        fake_file = os.fdopen(fd, 'w')
+        if sys.version_info > (3, 0):
+            self.assertEqual(fake_file.buffer,
+                             subunit._unwrap_text(fake_file))
+        else:
+            self.assertEqual(fake_file, subunit._unwrap_text(fake_file))
+
+    def test__unwrap_text_fileIO_read_mode(self):
+        fd, file_path = tempfile.mkstemp()
+        self.addCleanup(os.remove, file_path)
+        fake_file = io.FileIO(file_path, 'r')
+        self.assertEqual(fake_file, subunit._unwrap_text(fake_file))
+
+    def test__unwrap_text_fileIO_write_mode(self):
+        fd, file_path = tempfile.mkstemp()
+        self.addCleanup(os.remove, file_path)
+        fake_file = io.FileIO(file_path, 'w')
+        self.assertEqual(fake_file, subunit._unwrap_text(fake_file))
+
+    def test__unwrap_text_BytesIO(self):
+        fake_stream = io.BytesIO()
+        self.assertEqual(fake_stream, subunit._unwrap_text(fake_stream))
+
+
 class TestTestImports(unittest.TestCase):
 
     def test_imports(self):