diff --git a/awscli/customizations/datapipeline/__init__.py b/awscli/customizations/datapipeline/__init__.py index 437c37d09b78..3ffc70bab083 100644 --- a/awscli/customizations/datapipeline/__init__.py +++ b/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'] diff --git a/tests/unit/customizations/datapipeline/test_commands.py b/tests/unit/customizations/datapipeline/test_commands.py index ae502d15175d..85a11fd509a9 100644 --- a/tests/unit/customizations/datapipeline/test_commands.py +++ b/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):