def test_main(self):
# Mock out all of utils except parser
process._utils = MagicMock()
process._utils.add_default_client_arguments = _utils.add_default_client_arguments
# Set some static returns
process._utils.create_processing_job.return_value = 'job-name'
process._utils.get_processing_job_outputs.return_value = mock_outputs = {
'val1': 's3://1',
'val2': 's3://2'
}
process.main(required_args)
# Check if correct requests were created and triggered
process._utils.create_processing_job.assert_called()
process._utils.wait_for_processing_job.assert_called()
process._utils.print_logs_for_job.assert_called()
# Check the file outputs
process._utils.write_output.assert_has_calls([
call('/tmp/job_name_output_path', 'job-name'),
call('/tmp/output_artifacts_output_path',
mock_outputs,
json_encode=True)
])
def test_main(self):
# Mock out all of utils except parser
process._utils = MagicMock()
process._utils.add_default_client_arguments = _utils.add_default_client_arguments
# Set some static returns
process._utils.create_processing_job.return_value = 'job-name'
process._utils.get_processing_job_outputs.return_value = mock_outputs = {'val1': 's3://1', 'val2': 's3://2'}
with patch('builtins.open', mock_open()) as file_open:
process.main(required_args)
# Check if correct requests were created and triggered
process._utils.create_processing_job.assert_called()
process._utils.wait_for_processing_job.assert_called()
# Check the file outputs
file_open.assert_has_calls([
call('/tmp/job_name.txt', 'w'),
call('/tmp/output_artifacts.txt', 'w')
], any_order=True)
file_open().write.assert_has_calls([
call('job-name'),
call(json.dumps(mock_outputs))
], any_order=False) # Must be in the same order as called
def test_main_assumes_role(self):
# Mock out all of utils except parser
process._utils = MagicMock()
process._utils.add_default_client_arguments = _utils.add_default_client_arguments
# Set some static returns
process._utils.create_processing_job.return_value = 'job-name'
process._utils.get_processing_job_outputs.return_value = mock_outputs = {'val1': 's3://1', 'val2': 's3://2'}
assume_role_args = required_args + ['--assume_role', 'my-role']
process.main(assume_role_args)
process._utils.get_sagemaker_client.assert_called_once_with('us-west-2', None, assume_role_arn='my-role')