ashb commented on a change in pull request #4091: [AIRFLOW-2524] Update
SageMaker hook, operator and sensor for training, tuning and transform
URL: https://github.com/apache/incubator-airflow/pull/4091#discussion_r228727595
##########
File path: tests/contrib/hooks/test_sagemaker_hook.py
##########
@@ -233,94 +254,82 @@ class TestSageMakerHook(unittest.TestCase):
def setUp(self):
configuration.load_test_config()
- db.merge_conn(
- models.Connection(
- conn_id='sagemaker_test_conn_id',
- conn_type='sagemaker',
- login='access_id',
- password='access_key',
- extra=json.dumps(db_config)
- )
- )
+
+ @mock.patch.object(SageMakerHook, 'get_iam_conn')
+ def test_expand_role(self, mock_iam):
+ mock_session = mock.Mock()
+ some_role = 'some-role'
+ role_description = {
+ 'Role': {
+ 'Arn': 'test-arn'
+ }
+ }
+ attrs = {'get_role.return_value':
+ role_description}
+ mock_session.configure_mock(**attrs)
+ mock_iam.return_value = mock_session
+ hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
+ arn = hook.expand_role(some_role)
+ mock_session.get_role.\
+ assert_called_once_with(RoleName=some_role)
+ self.assertEqual(arn, role_description['Role']['Arn'])
+
+ @mock.patch.object(S3Hook, 'create_bucket')
+ @mock.patch.object(S3Hook, 'load_file')
+ def test_configure_s3_resources(self, mock_load_file, mock_create_bucket):
+ hook = SageMakerHook()
+ evaluation_result = {
+ 'Image': image,
+ 'Role': role
+ }
+ self.assertEqual(evaluation_result,
+ hook.configure_s3_resources(test_evaluation_config))
+ mock_create_bucket.assert_called_once_with(bucket_name=bucket)
+ mock_load_file.assert_called_once_with(path, key, bucket)
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(S3Hook, 'check_for_key')
@mock.patch.object(S3Hook, 'check_for_bucket')
@mock.patch.object(S3Hook, 'check_for_prefix')
Review comment:
These mocks are okay as they are though.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
