Switch datapipeline list-runs to use clients

awscli/customizations/datapipeline/__init__.py

@@ -340,15 +340,14 @@ def __init__(self, session, formatter=None):
         self._formatter = formatter
 
     def _run_main(self, parsed_args, parsed_globals, **kwargs):
-        self._set_session_objects(parsed_globals)
+        self._set_client(parsed_globals)
         self._parse_type_args(parsed_args)
         self._list_runs(parsed_args)
 
-    def _set_session_objects(self, parsed_globals):
+    def _set_client(self, parsed_globals):
         # This is called from _run_main and is used to ensure that we have
         # a service/endpoint object to work with.
-        self.service = self._session.get_service('datapipeline')
-        self.endpoint = self.service.get_endpoint(
+        self.client = self._session.create_client('datapipeline',
             region_name=parsed_globals.region,
             endpoint_url=parsed_globals.endpoint_url,
             verify=parsed_globals.verify_ssl)
@@ -389,15 +388,13 @@ def _list_runs(self, parsed_args):
         self._formatter.display_objects_to_user(converted)
 
     def _describe_objects(self, pipeline_id, object_ids):
-        operation = self.service.get_operation('DescribeObjects')
-        http_parsed, parsed = operation.call(
-            self.endpoint, pipeline_id=pipeline_id, object_ids=object_ids)
+        parsed = self.client.describe_objects(
+            pipelineId=pipeline_id, objectIds=object_ids)
         return parsed
 
     def _query_objects(self, pipeline_id, query):
-        operation = self.service.get_operation('QueryObjects')
-        paginator = operation.paginate(
-            self.endpoint, pipeline_id=pipeline_id,
+        paginator = self.client.get_paginator('query_objects').paginate(
+            pipelineId=pipeline_id,
             sphere='INSTANCE', query=query)
         parsed = paginator.build_full_result()
         return parsed['ids']

tests/unit/customizations/datapipeline/test_commands.py

@@ -191,32 +191,26 @@ def setUp(self):
         super(TestCommandsRunProperly, self).setUp()
         self.query_objects = mock.Mock()
         self.describe_objects = mock.Mock()
-
-    def get_service(self, name):
-        if name == 'QueryObjects':
-            return self.query_objects
-        elif name== 'DescribeObjects':
-            return self.describe_objects
+        self.client = mock.Mock()
+        self.client.get_paginator.return_value = self.query_objects
+        self.client.describe_objects = self.describe_objects
 
     def test_list_runs(self):
         self.driver.session = mock.Mock()
         self.driver.session.emit_first_non_none_response.return_value = None
-        self.driver.session.get_service.return_value.get_endpoint.return_value = \
-                mock.sentinel.endpoint
-        self.driver.session.get_service.return_value.get_operation = self.get_service
-        self.query_objects.paginate.return_value.build_full_result.return_value = {
-            'ids': ['object-ids']}
-        self.describe_objects.call.return_value = (
-            None, {'pipelineObjects': [
-                {'fields': [], 'id': 'id', 'name': 'name'}]})
+        self.driver.session.create_client.return_value = self.client
+        self.query_objects.paginate.return_value.build_full_result.\
+            return_value = {
+                'ids': ['object-ids']}
+        self.describe_objects.return_value = {'pipelineObjects': [
+            {'fields': [], 'id': 'id', 'name': 'name'}]}
 
         command = ListRunsCommand(self.driver.session, formatter=mock.Mock())
         command(['--pipeline-id', 'my-pipeline-id'],
                 parsed_globals=FakeParsedArgs(region='us-east-1'))
         self.assertTrue(self.query_objects.paginate.called)
-        self.describe_objects.call.assert_called_with(
-            mock.sentinel.endpoint, pipeline_id='my-pipeline-id',
-            object_ids=['object-ids'])
+        self.describe_objects.assert_called_with(
+            pipelineId='my-pipeline-id', objectIds=['object-ids'])
 
 
 class TestHelpOutput(BaseAWSHelpOutputTest):