def test_bad_storage_type(self):
"""
Test an error is raised when an unknown storage type is configured.
"""
# bad_storage_type, anything, anything, anything
with self.assertRaises(FileflowError):
get_storage_driver('bad_storage_type', '', '', '', '')
def test_s3_bad_environment(self):
"""
Test an error is raised when using an incorrect environment name for
S3 storage.
"""
# s3, bucket_name, bad_environment, key1, key2
with self.assertRaises(FileflowError):
get_storage_driver(
's3', '', 'bad_environment', '', '', 'the_bucket'
)
def test_file_driver(self):
"""
Test the file storage driver is returned when configured.
"""
# file, prefix, anything, anything, anything
driver = get_storage_driver('file', '/the/prefix/', '', '', '')
self.assertIsInstance(driver, FileStorageDriver)
self.assertEqual(driver.prefix, '/the/prefix/')
def storage(self):
"""
Lazy load a storage property on access instead of on class instantiation. Something in the storage attribute
is not deep-copyable which causes errors with airflow clear and airflow backfill which both try to deep copy
a target DAG and all its operators, so we only want this property when we actually use it.
"""
if self._storage is None:
self._storage = get_storage_driver()
return self._storage
def test_s3_test_environment(self):
"""
Test the storage driver and bucket name when using S3 in test.
"""
self.conn.create_bucket('the_buckettest')
# s3, bucket_name, test, key1, key2
driver = get_storage_driver('s3', '', 'test', '', '', 'the_bucket')
self.assertIsInstance(driver, S3StorageDriver)
self.assertEqual(driver.bucket_name, 'the_buckettest')
def test_s3_production_environment(self):
"""
Test the storage driver and bucket name when using S3 in production.
"""
self.conn.create_bucket('the_bucket')
# s3, bucket_name, production, any, any
driver = get_storage_driver(
's3', '', 'production', '', '', 'the_bucket'
)
self.assertIsInstance(driver, S3StorageDriver)
self.assertEqual(driver.bucket_name, 'the_bucket')
def __init__(self, context):
# The upstream dependencies
# These must always be specified
# Dictionary can contain any number of keys which must be redirected in the business logic to their read/parse methods
self.data_dependencies = context.pop('data_dependencies', {})
# The task instance.
self.task_instance = context['ti']
self.date = context['execution_date']
# Picking a storage driver for this task instance.
self.storage = get_storage_driver()