Update waiter code based on feedback

awscli/clidriver.py

@@ -117,7 +117,8 @@ def _build_command_table(self):
         command_table = self._build_builtin_commands(self.session)
         self.session.emit('building-command-table.main',
                           command_table=command_table,
-                          session=self.session)
+                          session=self.session,
+                          command_object=self)
         return command_table
 
     def _build_builtin_commands(self, session):
@@ -334,6 +335,10 @@ def name(self):
     def name(self, value):
         self._name = value
 
+    @property
+    def service_object(self):
+        return self._service_object
+
     def _get_command_table(self):
         if self._command_table is None:
             self._command_table = self._create_command_table()
@@ -367,7 +372,7 @@ def _create_command_table(self):
         self.session.emit('building-command-table.%s' % self._name,
                           command_table=command_table,
                           session=self.session,
-                          service_object=service_object)
+                          command_object=self)
         return command_table
 
     def create_help_command(self):

awscli/customizations/commands.py

@@ -211,7 +211,8 @@ def _build_subcommand_table(self):
             subcommand_table[subcommand_name] = subcommand_class(self._session)
         self._session.emit('building-command-table.%s' % self.NAME,
                            command_table=subcommand_table,
-                           session=self._session)
+                           session=self._session,
+                           command_object=self)
         return subcommand_table
 
     def _display_help(self, parsed_args, parsed_globals):

awscli/customizations/waiters.py

@@ -21,8 +21,11 @@ def register_add_waiters(cli):
     cli.register('building-command-table', add_waiters)
 
 
-def add_waiters(command_table, session, service_object=None, **kwargs):
-    # If a service object was passed in, try to add a wait command.
+def add_waiters(command_table, session, command_object, **kwargs):
+    # Check if the command object passed in has a ``service_object``. We
+    # only want to add wait commands to top level model-driven services.
+    # These require service objects.
+    service_object = getattr(command_object, 'service_object', None)
     if service_object is not None:
         # Get a client out of the service object.
         client = translate_service_object_to_client(service_object)
@@ -104,7 +107,7 @@ def _build_waiter_state_cmd(self, waiter_name):
         waiter_state_command = WaiterStateCommand(
             name=waiter_cli_name, parent_name='wait',
             operation_object=operation_object,
-            operation_caller=WaiterCaller(self._client, waiter),
+            operation_caller=WaiterCaller(self._client, waiter_name),
             service_object=self._service_object
         )
         # Build the top level description for the waiter state command.
@@ -169,20 +172,20 @@ def _build_operation_description(self, operation):
 
 
 class WaiterCaller(object):
-    def __init__(self, client, waiter):
+    def __init__(self, client, waiter_name):
         self._client = client
-        self._waiter = waiter
+        self._waiter_name = waiter_name
 
     def invoke(self, operation_object, parameters, parsed_globals):
         # Create the endpoint based on the parsed globals
         endpoint = operation_object.service.get_endpoint(
             region_name=parsed_globals.region,
             endpoint_url=parsed_globals.endpoint_url,
             verify=parsed_globals.verify_ssl)
-        # Change the client's endpoint using the newly configured endpoint
-        self._client._endpoint = endpoint
-        # Call the waiter's wait method.
-        self._waiter.wait(**parameters)
+        # Make a clone of the client using the newly configured endpoint
+        client = self._client.clone_client(endpoint=endpoint)
+        # Make the waiter and call its wait method.
+        client.get_waiter(self._waiter_name).wait(**parameters)
         return 0
 
 

tests/unit/customizations/test_waiters.py

@@ -25,6 +25,9 @@ def setUp(self):
         self.session = mock.Mock()
         self.client = mock.Mock()
 
+        self.command_object = mock.Mock()
+        self.command_object.service_object = self.service_object
+
         # Set up the mock service object.
         self.service_object.session = self.session
 
@@ -36,22 +39,23 @@ def setUp(self):
 
     def test_add_waiters(self):
         command_table = {}
-        add_waiters(command_table, self.session, self.service_object)
+        add_waiters(command_table, self.session, self.command_object)
         # Make sure a wait command was added.
         self.assertIn('wait', command_table)
         self.assertIsInstance(command_table['wait'], WaitCommand)
 
     def test_add_waiters_no_waiter_names(self):
         self.client.waiter_names = []
         command_table = {}
-        add_waiters(command_table, self.session, self.service_object)
+        add_waiters(command_table, self.session, self.command_object)
         # Make sure that no wait command was added since the service object
         # has no waiters.
         self.assertEqual(command_table, {})
 
     def test_add_waiters_no_service_object(self):
         command_table = {}
-        add_waiters(command_table, self.session, None)
+        self.command_object.service_object = None
+        add_waiters(command_table, self.session, self.command_object)
         # Make sure that no wait command was added since no service object
         # was passed in.
         self.assertEqual(command_table, {})
@@ -280,22 +284,33 @@ class TestWaiterCaller(unittest.TestCase):
     def test_invoke(self):
         client = mock.Mock()
         waiter = mock.Mock()
+        waiter_name = 'my_waiter'
         operation_object = mock.Mock()
 
+        # Mock the clone of the client
+        cloned_client = mock.Mock()
+        cloned_client.get_waiter.return_value = waiter
+        client.clone_client.return_value = cloned_client
+
         parameters = {'Foo': 'bar', 'Baz': 'biz'}
         parsed_globals = mock.Mock()
         parsed_globals.region = 'us-east-1'
         parsed_globals.endpoint_url = 'myurl'
         parsed_globals.verify_ssl = True
 
-        waiter_caller = WaiterCaller(client, waiter)
+        waiter_caller = WaiterCaller(client, waiter_name)
         waiter_caller.invoke(operation_object, parameters, parsed_globals)
         # Make sure the endpoint was created properly
         operation_object.service.get_endpoint.assert_called_with(
             region_name=parsed_globals.region,
             endpoint_url=parsed_globals.endpoint_url,
             verify=parsed_globals.verify_ssl
         )
+        # Ensure the client was cloned with using the new endpoint.
+        clone_kwargs = client.clone_client.call_args[1]
+        self.assertIn('endpoint', clone_kwargs)
+        # Ensure we get the waiter.
+        cloned_client.get_waiter.assert_called_with(waiter_name)
         # Ensure the wait command was called properly.
         waiter.wait.assert_called_with(
             Foo='bar', Baz='biz')