Fixed how the six.moves.input function is imported and used.

Also added a unit test for the cmd2.Cmd.select() method.

3 files changed, 37 insertions, 7 deletions

diff --git a/cmd2.py b/cmd2.py
index 1a43a47b..12e3cede 100755
--- a/cmd2.py
+++ b/cmd2.py
@@ -51,8 +51,8 @@ from six import next
 # Possible types for text data. This is basestring() in Python 2 and str in Python 3.
 from six import string_types
 
-# raw_input() for Python 2 or input() for Python 3
-from six.moves import input
+# Used for sm.input: raw_input() for Python 2 or input() for Python 3
+import six.moves as sm
 
 # itertools.zip() for Python 2 or zip() for Python 3 - produces an iterator in both cases
 from six.moves import zip
@@ -1010,7 +1010,7 @@ class Cmd(cmd.Cmd):
 
         if self.use_rawinput:
             try:
-                line = input(prompt)
+                line = sm.input(prompt)
             except EOFError:
                 line = 'EOF'
         else:
@@ -1103,7 +1103,7 @@ class Cmd(cmd.Cmd):
         for (idx, (value, text)) in enumerate(fulloptions):
             self.poutput('  %2d. %s\n' % (idx + 1, text))
         while True:
-            response = input(prompt)
+            response = sm.input(prompt)
             try:
                 response = int(response)
                 result = fulloptions[response - 1][0]
@@ -1164,7 +1164,7 @@ class Cmd(cmd.Cmd):
 
     def do_pause(self, arg):
         'Displays the specified text then waits for the user to press RETURN.'
-        input(arg + '\n')
+        sm.input(arg + '\n')
 
     def do_shell(self, arg):
         'execute a command as if at the OS prompt.'
diff --git a/setup.py b/setup.py
index 14c82219..8270b9f6 100755
--- a/setup.py
+++ b/setup.py
@@ -55,6 +55,7 @@ CLASSIFIERS = filter(None, map(str.strip,
                                """.splitlines()))
 
 INSTALL_REQUIRES = ['pyparsing >= 2.0.1', 'six']
+# unitest.mock was added in Python 3.3.  mock is a backport of unittest.mock to all versions of Python
 TESTS_REQUIRE = ['mock', 'pytest']
 
 setup(
diff --git a/tests/test_transcript.py b/tests/test_transcript.py
index a4102bc8..d11eb02a 100644
--- a/tests/test_transcript.py
+++ b/tests/test_transcript.py
@@ -7,8 +7,11 @@ Released under MIT license, see LICENSE file
 """
 import sys
 
+import mock
 import pytest
-from mock import patch
+
+# Used for sm.input: raw_input() for Python 2 or input() for Python 3
+import six.moves as sm
 
 from cmd2 import Cmd, make_option, options, Cmd2TestCase
 from conftest import run_cmd, StdOut, normalize
@@ -55,6 +58,13 @@ class DemoApp(Cmd):
         else:
             self.stdout.write('Hello Nobody\n')
 
+    def do_eat(self, arg):
+        """Eat something, with a selection of sauces to choose from."""
+        sauce = self.select('sweet salty', 'Sauce? ')
+        result = '{food} with {sauce} sauce, yum!'
+        result = result.format(food=arg, sauce=sauce)
+        self.stdout.write(result + '\n')
+
 
 @pytest.fixture
 def _cmdline_app():
@@ -217,7 +227,7 @@ def test_optarser_options_with_spaces_in_quotes(_demo_app):
 def test_commands_at_invocation(_cmdline_app):
     testargs = ["prog", "say hello", "say Gracie", "quit"]
     expected = "hello\nGracie\n"
-    with patch.object(sys, 'argv', testargs):
+    with mock.patch.object(sys, 'argv', testargs):
         app = CmdLineApp()
         app.stdout = StdOut()
         app.cmdloop()
@@ -225,3 +235,22 @@ def test_commands_at_invocation(_cmdline_app):
         assert out == expected
 
 
+def test_select_options(_demo_app):
+    # Mock out the input call so we don't actually wait for a user's `response on stdin
+    m = mock.MagicMock(name='input', return_value='2')
+    sm.input = m
+
+    food = 'bacon'
+    out = run_cmd(_demo_app, "set debug true")
+    out = run_cmd(_demo_app, "eat {}".format(food))
+    expected = normalize("""
+   1. sweet
+   2. salty
+{} with salty sauce, yum!
+""".format(food))
+
+    # Make sure our mock was called with the expected arguments
+    m.assert_called_once_with('Sauce? ')
+
+    # And verify the expected output to stdout
+    assert out == expected