Improve WizardLoader interface

Author: joguSD

awsshell/app.py

@@ -149,11 +149,12 @@ def run(self, command, application):
 
 
 class WizardHandler(object):
-    def __init__(self, output=sys.stdout, err=sys.stderr,
-                 loader=WizardLoader()):
+    def __init__(self, output=sys.stdout, err=sys.stderr, loader=None):
         self._output = output
         self._err = err
         self._wizard_loader = loader
+        if self._wizard_loader is None:
+            self._wizard_loader = WizardLoader()
 
     def run(self, command, application):
         """Run the specified wizard.

awsshell/wizard.py

@@ -51,7 +51,7 @@ class WizardLoader(object):
     """
 
     def __init__(self, session=None, interaction_loader=None,
-                 error_handler=None, delegation_loader=None):
+                 error_handler=None):
         """Initialize a wizard factory.
 
         :type session: :class:`botocore.session.Session`
@@ -71,10 +71,6 @@ def __init__(self, session=None, interaction_loader=None,
         self._interaction_loader = interaction_loader
         if interaction_loader is None:
             self._interaction_loader = InteractionLoader()
-        if delegation_loader is None:
-            self._delegation_loader = self
-        else:
-            self._delegation_loader = delegation_loader
         self._error_handler = error_handler
         if error_handler is None:
             self._error_handler = stage_error_handler
@@ -126,8 +122,7 @@ def _load_stage(self, stage, env):
         }
         creator = self._cached_creator
         interaction = self._interaction_loader
-        delegation = self._delegation_loader
-        return Stage(env, creator, interaction, delegation, **stage_attrs)
+        return Stage(env, creator, interaction, self, **stage_attrs)
 
     def _load_stages(self, stages, env):
         return [self._load_stage(stage, env) for stage in stages]

tests/unit/test_wizard.py

@@ -1,5 +1,8 @@
 import mock
 import pytest
+import botocore.session
+
+from botocore.loaders import Loader
 from botocore.session import Session
 from awsshell.utils import FileReadError
 from awsshell.wizard import stage_error_handler
@@ -289,9 +292,6 @@ def test_wizard_basic_interaction(wizard_spec):
 
 
 def test_wizard_basic_delegation(wizard_spec):
-    mock_session = mock.Mock(spec=Session)
-    mock_loader = mock.Mock(spec=WizardLoader)
-    loader = WizardLoader(mock_session, delegation_loader=mock_loader)
     main_spec = {
         "StartStage": "One",
         "Stages": [
@@ -306,7 +306,6 @@ def test_wizard_basic_delegation(wizard_spec):
             }
         ]
     }
-    wizard = loader.create_wizard(main_spec)
     sub_spec = {
         "StartStage": "SubOne",
         "Stages": [
@@ -320,9 +319,19 @@ def test_wizard_basic_delegation(wizard_spec):
             }
         ]
     }
-    mock_loader.load_wizard.return_value = loader.create_wizard(sub_spec)
+
+    mock_loader = mock.Mock(spec=Loader)
+    mock_loader.list_available_services.return_value = ['wizards']
+    mock_load_model = mock_loader.load_service_model
+    mock_load_model.return_value = sub_spec
+
+    session = botocore.session.get_session()
+    session.register_component('data_loader', mock_loader)
+    loader = WizardLoader(session)
+    wizard = loader.create_wizard(main_spec)
+
     result = wizard.execute()
-    mock_loader.load_wizard.assert_called_once_with('SubWizard')
+    mock_load_model.assert_called_once_with('wizards', 'SubWizard')
     assert result == 'Result from sub'