def test_write_raises(self):
self.hook_inst_mock.load_string.side_effect = Exception('error')
handler = S3TaskHandler()
with mock.patch.object(handler.logger, 'error') as mock_error:
handler.write('text', self.remote_log_location)
msg = 'Could not write logs to %s' % self.remote_log_location
mock_error.assert_called_once_with(msg)
def setUp(self):
super().setUp()
self.remote_log_base = 's3://bucket/remote/log/location'
self.remote_log_location = 's3://bucket/remote/log/location/1.log'
self.remote_log_key = 'remote/log/location/1.log'
self.local_log_location = 'local/log/location'
self.filename_template = '{try_number}.log'
self.s3_task_handler = S3TaskHandler(self.local_log_location,
self.remote_log_base,
self.filename_template)
date = datetime(2016, 1, 1)
self.dag = DAG('dag_for_testing_file_task_handler', start_date=date)
task = DummyOperator(task_id='task_for_testing_file_log_handler',
dag=self.dag)
self.ti = TaskInstance(task=task, execution_date=date)
self.ti.try_number = 1
self.ti.state = State.RUNNING
self.addCleanup(self.dag.clear)
self.conn = boto3.client('s3')
# We need to create the bucket since this is all in Moto's 'virtual'
# AWS account
moto.core.moto_api_backend.reset()
self.conn.create_bucket(Bucket="bucket")
def test_write(self):
S3TaskHandler().write('text', self.remote_log_location)
self.hook_inst_mock.load_string.assert_called_once_with(
'content\ntext',
key=self.remote_log_location,
replace=True,
encrypt=False,
)
def test_read_raises_return_error(self):
self.hook_inst_mock.get_key.side_effect = Exception('error')
handler = S3TaskHandler()
with mock.patch.object(handler.logger, 'error') as mock_error:
result = handler.s3_log_read(self.remote_log_location,
return_error=True)
msg = 'Could not read logs from %s' % self.remote_log_location
self.assertEqual(result, msg)
mock_error.assert_called_once_with(msg)
def test_init_raises(self):
self.hook_mock.side_effect = Exception('Failed to connect')
handler = S3TaskHandler()
with mock.patch.object(handler.logger, 'error') as mock_error:
# Initialize the hook
handler.hook()
mock_error.assert_called_once_with(
'Could not create an S3Hook with connection id "". Please make '
'sure that airflow[s3] is installed and the S3 connection exists.'
)
def test_hook_raises(self):
handler = S3TaskHandler(self.local_log_location, self.remote_log_base,
self.filename_template)
with mock.patch.object(handler.log, 'error') as mock_error:
with mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook"
) as mock_hook:
mock_hook.side_effect = Exception('Failed to connect')
# Initialize the hook
handler.hook
mock_error.assert_called_once_with(
'Could not create an S3Hook with connection id "%s". Please make '
'sure that airflow[aws] is installed and the S3 connection exists.',
'aws_default')
def setUp(self):
super(TestS3TaskHandler, self).setUp()
self.remote_log_location = 'remote/log/location'
self.local_log_location = 'local/log/location'
self.s3_log_location = 's3/log/location'
self.filename_template = ''
self.hook_patcher = mock.patch("airflow.hooks.S3_hook.S3Hook")
self.hook_mock = self.hook_patcher.start()
self.hook_inst_mock = self.hook_mock.return_value
self.hook_key_mock = self.hook_inst_mock.get_key.return_value
self.hook_key_mock.get_contents_as_string.return_value.decode.\
return_value = 'content'
self.s3_task_handler = S3TaskHandler(self.local_log_location,
self.s3_log_location,
self.filename_template)
def test_read_raises(self):
self.hook_inst_mock.get_key.side_effect = Exception('error')
self.assertEqual(S3TaskHandler().read(self.remote_log_location), '')
def test_read_key_empty(self):
self.hook_inst_mock.get_key.return_value = None
self.assertEqual(S3TaskHandler().read(self.remote_log_location), '')
def test_read(self):
self.assertEqual(S3TaskHandler().read(self.remote_log_location),
'content')
def test_log_exists_no_hook(self):
self.hook_mock.side_effect = Exception('Failed to connect')
self.assertFalse(S3TaskHandler().log_exists(self.remote_log_location))
def test_log_exists_raises(self):
self.hook_inst_mock.get_key.side_effect = Exception('error')
self.assertFalse(S3TaskHandler().log_exists(self.remote_log_location))
def test_log_exists_none(self):
self.hook_inst_mock.get_key.return_value = None
self.assertFalse(S3TaskHandler().log_exists(self.remote_log_location))
def test_log_exists(self):
self.assertTrue(S3TaskHandler().log_exists(self.remote_log_location))
def test_init(self):
S3TaskHandler()
self.hook_mock.assert_called_once_with('')
def test_log_exists_false(self):
self.hook_inst_mock.check_for_key.return_value = False
self.assertFalse(S3TaskHandler().log_exists(self.remote_log_location))
