Tests for wrapping Python shell

1 files changed, 27 insertions, 1 deletions

diff --git a/tests/test_replwrap.py b/tests/test_replwrap.py
index febf221..442c613 100644
--- a/tests/test_replwrap.py
+++ b/tests/test_replwrap.py
@@ -1,10 +1,11 @@
+import platform
 import unittest
 
 import pexpect
 from pexpect import replwrap
 
 class REPLWrapTestCase(unittest.TestCase):
-    def test_python(self):
+    def test_bash(self):
         bash = replwrap.bash()
         res = bash.run_command("time")
         assert 'real' in res, res
@@ -34,5 +35,30 @@ class REPLWrapTestCase(unittest.TestCase):
         res = repl.run_command("echo $HOME")
         assert res.startswith('/'), res
 
+    def test_python(self):
+        if platform.python_implementation() == 'PyPy':
+            raise unittest.SkipTest("This test fails on PyPy because of REPL differences")
+
+        p = replwrap.python()
+        res = p.run_command('4+7')
+        assert res.strip() == '11'
+
+        res = p.run_command('for a in range(3): print(a)\n')
+        assert res.strip().splitlines() == ['0', '1', '2']
+
+    def test_no_change_prompt(self):
+        if platform.python_implementation() == 'PyPy':
+            raise unittest.SkipTest("This test fails on PyPy because of REPL differences")
+
+        child = pexpect.spawnu('python')
+        # prompt_change=None should mean no prompt change
+        py = replwrap.REPLWrapper(child, replwrap.u(">>> "), prompt_change=None,
+                                  continuation_prompt=replwrap.u("... "))
+        assert py.prompt == ">>> "
+
+        res = py.run_command("for a in range(3): print(a)\n")
+        assert res.strip().splitlines() == ['0', '1', '2']
+
+
 if __name__ == '__main__':
     unittest.main()