Refactor handler to preserve stack trace

Author: joguSD

awsshell/wizard.py

@@ -35,7 +35,7 @@ def _select_stage_prompt():
         else:
             raise KeyboardInterrupt()
     else:
-        raise error
+        return None
 
 
 class WizardException(Exception):
@@ -180,7 +180,10 @@ def execute(self):
                 current_stage = stage.get_next_stage()
             except Exception as err:
                 stages = [s.name for (s, _) in self._stage_history]
-                (stage, index) = self._error_handler(err, stages)
+                recovery = self._error_handler(err, stages)
+                if recovery is None:
+                    raise
+                (stage, index) = recovery
                 self._pop_stages(index)
                 current_stage = stage
 

tests/unit/test_wizard.py

@@ -309,5 +309,5 @@ def test_stage_exception_handler_other(error_class):
     prompt = mock.Mock()
     confirm = mock.Mock()
     err = error_class()
-    with pytest.raises(error_class):
-        stage_error_handler(err, ['stage'], confirm=confirm, prompt=prompt)
+    res = stage_error_handler(err, ['stage'], confirm=confirm, prompt=prompt)
+    assert res is None