Add exit/quit dot commands

awsshell/app.py

@@ -30,6 +30,7 @@
 
 
 LOG = logging.getLogger(__name__)
+EXIT_REQUESTED = object()
 
 
 def create_aws_shell(completer, model_completer, docs):
@@ -82,9 +83,16 @@ def run(self, command, application):
             p.communicate()
 
 
+class ExitHandler(object):
+    def run(self, command, application):
+        return EXIT_REQUESTED
+
+
 class DotCommandHandler(object):
     HANDLER_CLASSES = {
         'edit': EditHandler,
+        'exit': ExitHandler,
+        'quit': ExitHandler,
     }
 
     def __init__(self, output=sys.stdout, err=sys.stderr):
@@ -109,7 +117,7 @@ def handle_cmd(self, command, application):
         else:
             # Note we expect the class to support no-arg
             # instantiation.
-            self.HANDLER_CLASSES[cmd_name]().run(parts, application)
+            return self.HANDLER_CLASSES[cmd_name]().run(parts, application)
 
     def _unknown_cmd(self, cmd_parts, application):
         self._err.write("Unknown dot command: %s\n" % cmd_parts[0])
@@ -204,12 +212,12 @@ def run(self):
                 self.save_config()
                 break
             else:
-                if text.strip() in ['quit', 'exit']:
-                    break
                 if text.startswith('.'):
                     # These are special commands.  The only one supported for
                     # now is .edit.
-                    self._dot_cmd.handle_cmd(text, application=self)
+                    result = self._dot_cmd.handle_cmd(text, application=self)
+                    if result is EXIT_REQUESTED:
+                        break
                 else:
                     if text.startswith('!'):
                         # Then run the rest as a normally shell command.

tests/unit/test_app.py

@@ -42,3 +42,19 @@ def test_prints_error_message_on_unknown_dot_command(errstream):
     handler = app.DotCommandHandler(err=errstream)
     handler.handle_cmd(".unknown foo bar", None)
     assert errstream.getvalue() == "Unknown dot command: .unknown\n"
+
+
+def test_exit_dot_command_exits_shell():
+    mock_prompter = mock.Mock()
+    # Simulate the user entering '.quit'
+    fake_document = mock.Mock()
+    fake_document.text = '.quit'
+    mock_prompter.run.return_value = fake_document
+    shell = app.AWSShell(mock.Mock(), mock.Mock(), mock.Mock())
+    shell.create_cli_interface = mock.Mock(return_value=mock_prompter)
+    shell.run()
+
+    # Should have only called run() once.  As soon as we
+    # see the .quit command, we immediately exit and stop prompting
+    # for more shell commands.
+    assert mock_prompter.run.call_count == 1