1 files changed, 20 insertions, 1 deletions
diff --git a/python/google/protobuf/internal/input_stream_test.py b/python/google/protobuf/internal/input_stream_test.py
index 8cc1d126..ecec7f7d 100755
--- a/python/google/protobuf/internal/input_stream_test.py
+++ b/python/google/protobuf/internal/input_stream_test.py
@@ -40,7 +40,14 @@ from google.protobuf.internal import wire_format
 from google.protobuf.internal import input_stream
 
 
-class InputStreamTest(unittest.TestCase):
+class InputStreamBufferTest(unittest.TestCase):
+
+  def setUp(self):
+    self.__original_input_stream = input_stream.InputStream
+    input_stream.InputStream = input_stream.InputStreamBuffer
+
+  def tearDown(self):
+    input_stream.InputStream = self.__original_input_stream
 
   def testEndOfStream(self):
     stream = input_stream.InputStream('abcd')
@@ -291,5 +298,17 @@ class InputStreamTest(unittest.TestCase):
     stream = input_stream.InputStream(s)
     self.assertRaises(message.DecodeError, stream.ReadVarUInt64)
 
+
+class InputStreamArrayTest(InputStreamBufferTest):
+
+  def setUp(self):
+    # Test InputStreamArray against the same tests in InputStreamBuffer
+    self.__original_input_stream = input_stream.InputStream
+    input_stream.InputStream = input_stream.InputStreamArray
+
+  def tearDown(self):
+    input_stream.InputStream = self.__original_input_stream
+
+
 if __name__ == '__main__':
   unittest.main()