def test_describe_endpoint(self, mock_client):
mock_session = mock.Mock()
attrs = {'describe_endpoint.return_value': 'InProgress'}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.describe_endpoint(endpoint_name)
mock_session.describe_endpoint.assert_called_once_with(EndpointName=endpoint_name)
assert response == 'InProgress'
def get_sagemaker_response(self):
sagemaker = SageMakerHook(aws_conn_id=self.aws_conn_id)
self.log.info('Poking Sagemaker Endpoint %s', self.endpoint_name)
return sagemaker.describe_endpoint(self.endpoint_name)
