2 files changed, 48 insertions, 1 deletions
diff --git a/cmd2/utils.py b/cmd2/utils.py
index 83185603..bdb488cc 100644
--- a/cmd2/utils.py
+++ b/cmd2/utils.py
@@ -262,6 +262,8 @@ class StdSim(object):
 
         def write(self, b: bytes) -> None:
             """Add bytes to internal bytes buffer and if echo is True, echo contents to inner stream."""
+            if not isinstance(b, bytes):
+                raise TypeError('a bytes-like object is required, not {}'.format(type(b)))
             self.byte_buf += b
             if self.echo:
                 self.inner_stream.buffer.write(b)
@@ -272,6 +274,8 @@ class StdSim(object):
 
     def write(self, s: str) -> None:
         """Add str to internal bytes buffer and if echo is True, echo contents to inner stream."""
+        if not isinstance(s, str):
+            raise TypeError('write() argument must be str, not {}'.format(type(s)))
         b = s.encode()
         self.buffer.write(b)
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 8c8daa39..12d87218 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -5,6 +5,10 @@ Unit testing for cmd2/utils.py module.
 Copyright 2018 Todd Leonhardt <todd.leonhardt@gmail.com>
 Released under MIT license, see LICENSE file
 """
+import sys
+
+import pytest
+
 from colorama import Fore
 import cmd2.utils as cu
 
@@ -110,4 +114,43 @@ def test_quot_string_if_needed_no():
     assert cu.quote_string_if_needed(your_str) == your_str
 
 
-
+@pytest.fixture
+def stdout_sim():
+    stdsim = cu.StdSim(sys.stdout)
+    return stdsim
+
+def test_stdsim_write_str(stdout_sim):
+    my_str = 'Hello World'
+    stdout_sim.write(my_str)
+    assert stdout_sim.getvalue() == my_str
+
+def test_stdsim_write_bytes(stdout_sim):
+    b_str = b'Hello World'
+    with pytest.raises(TypeError):
+        stdout_sim.write(b_str)
+
+def test_stdsim_buffer_write_bytes(stdout_sim):
+    b_str = b'Hello World'
+    stdout_sim.buffer.write(b_str)
+    assert stdout_sim.getvalue() == b_str.decode()
+
+def test_stdsim_buffer_write_str(stdout_sim):
+    my_str = 'Hello World'
+    with pytest.raises(TypeError):
+        stdout_sim.buffer.write(my_str)
+
+def test_stdsim_read(stdout_sim):
+    my_str = 'Hello World'
+    stdout_sim.write(my_str)
+    # getvalue() returns the value and leaves it unaffected internally
+    assert stdout_sim.getvalue() == my_str
+    # read() returns the value and then clears the internal buffer
+    assert stdout_sim.read() == my_str
+    assert stdout_sim.getvalue() == ''
+
+def test_stdsim_clear(stdout_sim):
+    my_str = 'Hello World'
+    stdout_sim.write(my_str)
+    assert stdout_sim.getvalue() == my_str
+    stdout_sim.clear()
+    assert stdout_sim.getvalue() == ''