def test_add_step(self):
schema = JsonObject(TABLE_NAME, Property('id', 'VARCHAR(36)'))
bucket = Mock()
database = create_autospec(Database)
expected = BulkCopyFromS3JsonStep(
metadata='',
source='',
schema=schema,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
bucket=bucket,
table=TargetTable(schema, database))
pipeline = S3BulkCopyPipeline(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY,
bucket, database)
pipeline.step(metadata='', source='', schema=schema)
step = pipeline.steps()[0]
self.assertEqual(expected.metadata, step.metadata)
self.assertEqual(expected.source, step.source)
self.assertEqual(expected.schema, step.schema)
self.assertEqual(expected.aws_access_key_id, step.aws_access_key_id)
self.assertEqual(expected.aws_secret_access_key,
step.aws_secret_access_key)
self.assertEqual(expected.bucket, step.bucket)
self.assertEqual(expected.table.schema, step.table.schema)
self.assertEqual(expected.table.database, step.table.database)
def bulk_copy(self, metadata, source, schema, max_error_count=1):
bulk_copy_step = BulkCopyFromS3JsonStep(metadata=metadata,
source=source,
schema=schema,
aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key,
bucket=self.bucket,
table=TargetTable(schema,
self.database))
self.__add_copy_step(bulk_copy_step, max_error_count)
return self
def step(self, metadata, source, schema, max_error_count=1):
bulk_copy_step = BulkCopyFromS3JsonStep(metadata=metadata,
source=source,
schema=schema,
aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key,
bucket=self.bucket,
table=TargetTable(schema,
self.database))
bulk_copy_step.max_error_count = max_error_count
self.steps().append(bulk_copy_step)
return self
def setUp(self):
self.bucket = Bucket(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY,
BUCKET_NAME, Mock())
self.bucket.save = Mock()
self.bucket.delete = Mock()
self.schema = JsonObject(TABLE_NAME).property('eventId', 'VARCHAR(36)')
self.table = TargetTable(self.schema, Mock())
self.step = BulkCopyFromS3JsonStep(metadata='', source=SOURCE,
schema=self.schema,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
bucket=self.bucket, table=self.table)
self.step.sql = Mock()
def test_validate_with_drop_on_run(self):
with patch.object(TargetTable, 'exists') as exists:
exists.return_value = True
target_table = TargetTable(self.schema, Mock())
target_table.stage_update = Mock()
target_table.drop = Mock()
target_table.promote_update = Mock()
step = BulkCopyFromS3JsonStep(metadata='',
source=SOURCE,
schema=self.schema,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
bucket=self.bucket,
table=target_table)
step.sql = Mock()
step.validate()
self.assert_migration_with_drop(target_table)
target_table.database.rollback.assert_called_once_with()