def test_remove_snapshot_ownership_no_copy_snapshot(
self, mock_aws, mock_boto3):
"""Assert remove snapshot ownership task succeeds with missing copy."""
mock_arn = util_helper.generate_dummy_arn()
mock_customer_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_customer_snapshot = util_helper.generate_mock_snapshot(
mock_customer_snapshot_id)
mock_snapshot_copy_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot_copy = util_helper.generate_mock_snapshot(
mock_snapshot_copy_id)
zone = settings.HOUNDIGRADE_AWS_AVAILABILITY_ZONE
region = zone[:-1]
client_error = ClientError(
error_response={"Error": {
"Code": "InvalidSnapshot.NotFound"
}},
operation_name=Mock(),
)
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot_copy
resource.Snapshot.side_effect = client_error
mock_aws.check_snapshot_state.return_value = None
mock_aws.get_snapshot.return_value = mock_customer_snapshot
mock_aws.get_region_from_availability_zone.return_value = region
tasks.remove_snapshot_ownership(mock_arn, mock_customer_snapshot_id,
region, mock_snapshot_copy_id)
mock_aws.remove_snapshot_ownership.assert_called_with(
mock_customer_snapshot)
def test_remove_snapshot_ownership_success(self, mock_aws, mock_boto3):
"""Assert that the remove snapshot ownership task succeeds."""
mock_arn = util_helper.generate_dummy_arn()
mock_customer_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_customer_snapshot = util_helper.generate_mock_snapshot(
mock_customer_snapshot_id)
mock_snapshot_copy_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot_copy = util_helper.generate_mock_snapshot(
mock_snapshot_copy_id)
zone = settings.HOUNDIGRADE_AWS_AVAILABILITY_ZONE
region = zone[:-1]
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot_copy
mock_aws.check_snapshot_state.return_value = None
mock_aws.get_snapshot.return_value = mock_customer_snapshot
mock_aws.get_region_from_availability_zone.return_value = region
tasks.remove_snapshot_ownership(mock_arn, mock_customer_snapshot_id,
region, mock_snapshot_copy_id)
mock_aws.remove_snapshot_ownership.assert_called_with(
mock_customer_snapshot)
def test_copy_ami_snapshot_retry_on_ownership_not_verified(self, mock_aws):
"""Assert that the snapshot copy task fails."""
mock_session = mock_aws.boto3.Session.return_value
mock_account_id = mock_aws.get_session_account_id.return_value
mock_arn = util_helper.generate_dummy_arn()
mock_region = random.choice(util_helper.SOME_AWS_REGIONS)
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
block_mapping = mock_image.block_device_mappings
mock_snapshot_id = block_mapping[0]['Ebs']['SnapshotId']
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, owner_id=mock_account_id)
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
mock_aws.add_snapshot_ownership.side_effect = \
AwsSnapshotNotOwnedError()
with patch.object(tasks, 'create_volume') as mock_create_volume,\
patch.object(copy_ami_snapshot, 'retry') as mock_retry:
mock_retry.side_effect = Retry()
with self.assertRaises(Retry):
copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
mock_create_volume.delay.assert_not_called()
def test_copy_ami_snapshot_private_shared(self, mock_aws):
"""Assert that the task copies the image when it is private/shared."""
mock_account_id = util_helper.generate_dummy_aws_account_id()
mock_session = mock_aws.boto3.Session.return_value
mock_aws.get_session_account_id.return_value = mock_account_id
# the account id to use as the private shared image owner
other_account_id = util_helper.generate_dummy_aws_account_id()
mock_region = util_helper.get_random_region()
mock_arn = util_helper.generate_dummy_arn(mock_account_id, mock_region)
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
mock_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, encrypted=False, owner_id=other_account_id)
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
account_helper.generate_image(ec2_ami_id=mock_image_id)
with patch.object(tasks.imageprep,
"create_volume") as mock_create_volume, patch.object(
tasks.imageprep, "copy_ami_to_customer_account"
) as mock_copy_ami_to_customer_account:
tasks.copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
mock_create_volume.delay.assert_not_called()
mock_copy_ami_to_customer_account.delay.assert_called_with(
mock_arn, mock_image_id, mock_region)
def test_copy_ami_snapshot_retry_on_ownership_not_verified(self, mock_aws):
"""Assert that the snapshot copy task fails."""
mock_session = mock_aws.boto3.Session.return_value
mock_account_id = mock_aws.get_session_account_id.return_value
mock_arn = util_helper.generate_dummy_arn()
mock_region = util_helper.get_random_region()
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
account_helper.generate_image(ec2_ami_id=mock_image_id)
block_mapping = mock_image.block_device_mappings
mock_snapshot_id = block_mapping[0]["Ebs"]["SnapshotId"]
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, owner_id=mock_account_id)
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
mock_aws.add_snapshot_ownership.side_effect = AwsSnapshotNotOwnedError(
)
with patch.object(tasks,
"create_volume") as mock_create_volume, patch.object(
tasks.copy_ami_snapshot, "retry") as mock_retry:
mock_retry.side_effect = Retry()
with self.assertRaises(Retry):
tasks.copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
mock_create_volume.delay.assert_not_called()
def test_copy_ami_snapshot_success_with_reference(self, mock_aws):
"""Assert the snapshot copy task succeeds using a reference AMI ID."""
mock_session = mock_aws.boto3.Session.return_value
mock_account_id = mock_aws.get_session_account_id.return_value
account = account_helper.generate_aws_account()
arn = account.account_arn
region = random.choice(util_helper.SOME_AWS_REGIONS)
new_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(new_image_id)
block_mapping = mock_image.block_device_mappings
mock_snapshot_id = block_mapping[0]['Ebs']['SnapshotId']
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, owner_id=mock_account_id)
mock_new_snapshot_id = util_helper.generate_dummy_snapshot_id()
# This is the original ID of a private/shared image.
# It would have been saved to our DB upon initial discovery.
reference_image = account_helper.generate_aws_image(account=account)
reference_image_id = reference_image.ec2_ami_id
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
mock_aws.copy_snapshot.return_value = mock_new_snapshot_id
with patch.object(tasks, 'create_volume') as mock_create_volume, \
patch.object(tasks, 'remove_snapshot_ownership') as \
mock_remove_snapshot_ownership:
tasks.copy_ami_snapshot(
arn,
new_image_id,
region,
reference_image_id,
)
# arn, customer_snapshot_id, snapshot_region, snapshot_copy_id
mock_remove_snapshot_ownership.delay.assert_called_with(
arn,
mock_snapshot_id,
region,
mock_new_snapshot_id,
)
mock_create_volume.delay.assert_called_with(
reference_image_id,
mock_new_snapshot_id,
)
mock_aws.get_session.assert_called_with(arn)
mock_aws.get_ami.assert_called_with(mock_session, new_image_id, region)
mock_aws.get_ami_snapshot_id.assert_called_with(mock_image)
mock_aws.add_snapshot_ownership.assert_called_with(mock_snapshot)
mock_aws.copy_snapshot.assert_called_with(mock_snapshot_id, region)
# Verify that the copy object was stored correctly to reference later.
copied_image = AwsMachineImageCopy.objects.get(ec2_ami_id=new_image_id)
self.assertIsNotNone(copied_image)
self.assertEqual(copied_image.reference_awsmachineimage.ec2_ami_id,
reference_image_id)
def test_copy_ami_snapshot_encrypted(self, mock_aws):
"""Assert that the task marks the image as encrypted in the DB."""
mock_account_id = util_helper.generate_dummy_aws_account_id()
mock_region = random.choice(util_helper.SOME_AWS_REGIONS)
mock_arn = util_helper.generate_dummy_arn(mock_account_id, mock_region)
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
mock_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot = util_helper.generate_mock_snapshot(mock_snapshot_id,
encrypted=True)
mock_session = mock_aws.boto3.Session.return_value
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
account = AwsAccount(
aws_account_id=mock_account_id,
account_arn=mock_arn,
user=util_helper.generate_test_user(),
)
account.save()
ami = AwsMachineImage.objects.create(account=account,
is_windows=False,
ec2_ami_id=mock_image_id)
ami.save()
with patch.object(tasks, 'create_volume') as mock_create_volume,\
self.assertRaises(AwsSnapshotEncryptedError):
copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
self.assertTrue(ami.is_encrypted)
mock_create_volume.delay.assert_not_called()
def test_remove_snapshot_ownership_not_verified(self):
"""Assert an error is raised when ownership is not removed."""
mock_user_id = str(uuid.uuid4())
mock_snapshot = helper.generate_mock_snapshot()
mock_snapshot.describe_attribute.return_value = {
'CreateVolumePermissions': [{
'UserId': mock_user_id
}],
}
with patch.object(ec2, '_get_primary_account_id') as mock_get_acct_id:
mock_get_acct_id.return_value = mock_user_id
with self.assertRaises(AwsSnapshotOwnedError):
ec2.remove_snapshot_ownership(mock_snapshot)
expected_permission = {'Remove': [{'UserId': mock_user_id}]}
expected_user_ids = [mock_user_id]
mock_snapshot.modify_attribute.assert_called_once_with(
Attribute='createVolumePermission',
CreateVolumePermission=expected_permission,
OperationType='remove',
UserIds=expected_user_ids)
mock_snapshot.describe_attribute.assert_called_once_with(
Attribute='createVolumePermission')
def test_add_snapshot_ownership_success(self):
"""Assert that snapshot ownership is modified successfully."""
mock_user_id = str(uuid.uuid4())
mock_snapshot = helper.generate_mock_snapshot()
mock_snapshot.describe_attribute.return_value = {
'CreateVolumePermissions': [{
'UserId': mock_user_id
}],
}
with patch.object(ec2, '_get_primary_account_id') as mock_get_acct_id:
mock_get_acct_id.return_value = mock_user_id
actual_modified = ec2.add_snapshot_ownership(mock_snapshot)
self.assertIsNone(actual_modified)
expected_permission = {'Add': [{'UserId': mock_user_id}]}
expected_user_ids = [mock_user_id]
mock_snapshot.modify_attribute.assert_called_once_with(
Attribute='createVolumePermission',
CreateVolumePermission=expected_permission,
OperationType='add',
UserIds=expected_user_ids)
mock_snapshot.describe_attribute.assert_called_once_with(
Attribute='createVolumePermission')
def test_copy_ami_snapshot_encrypted(self, mock_aws):
"""Assert that the task marks the image as encrypted in the DB."""
mock_account_id = util_helper.generate_dummy_aws_account_id()
mock_region = util_helper.get_random_region()
mock_arn = util_helper.generate_dummy_arn(mock_account_id, mock_region)
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
mock_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot = util_helper.generate_mock_snapshot(mock_snapshot_id,
encrypted=True)
mock_session = mock_aws.boto3.Session.return_value
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
ami = account_helper.generate_image(ec2_ami_id=mock_image_id)
with patch.object(tasks, "create_volume") as mock_create_volume:
tasks.copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
ami.refresh_from_db()
self.assertTrue(ami.is_encrypted)
self.assertEqual(ami.status, ami.ERROR)
mock_create_volume.delay.assert_not_called()
def test_remove_snapshot_ownership_other_user(self):
"""Assert an error is not raised when other user has ownership."""
mock_user_id = str(uuid.uuid4())
mock_snapshot = helper.generate_mock_snapshot()
mock_snapshot.describe_attribute.return_value = {
"CreateVolumePermissions": [{
"UserId": "mock_user_id"
}],
}
with patch.object(ec2, "_get_primary_account_id") as mock_get_acct_id:
mock_get_acct_id.return_value = mock_user_id
ec2.remove_snapshot_ownership(mock_snapshot)
expected_permission = {"Remove": [{"UserId": mock_user_id}]}
expected_user_ids = [mock_user_id]
mock_snapshot.modify_attribute.assert_called_once_with(
Attribute="createVolumePermission",
CreateVolumePermission=expected_permission,
OperationType="remove",
UserIds=expected_user_ids,
)
mock_snapshot.describe_attribute.assert_called_once_with(
Attribute="createVolumePermission")
def test_remove_snapshot_ownership_success(self):
"""Assert that snapshot ownership is removed successfully."""
mock_user_id = str(uuid.uuid4())
mock_snapshot = helper.generate_mock_snapshot()
mock_snapshot.describe_attribute.return_value = {
"CreateVolumePermissions": [],
}
with patch.object(ec2, "_get_primary_account_id") as mock_get_acct_id:
mock_get_acct_id.return_value = mock_user_id
actual_modified = ec2.remove_snapshot_ownership(mock_snapshot)
self.assertIsNone(actual_modified)
expected_permission = {"Remove": [{"UserId": mock_user_id}]}
expected_user_ids = [mock_user_id]
mock_snapshot.modify_attribute.assert_called_once_with(
Attribute="createVolumePermission",
CreateVolumePermission=expected_permission,
OperationType="remove",
UserIds=expected_user_ids,
)
mock_snapshot.describe_attribute.assert_called_once_with(
Attribute="createVolumePermission")
def test_copy_ami_snapshot_success(self, mock_aws):
"""Assert that the snapshot copy task succeeds."""
mock_arn = util_helper.generate_dummy_arn()
mock_region = random.choice(util_helper.SOME_AWS_REGIONS)
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
block_mapping = mock_image.block_device_mappings
mock_snapshot_id = block_mapping[0]['Ebs']['SnapshotId']
mock_snapshot = util_helper.generate_mock_snapshot(mock_snapshot_id)
mock_new_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_session = mock_aws.boto3.Session.return_value
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
mock_aws.copy_snapshot.return_value = mock_new_snapshot_id
with patch.object(tasks, 'create_volume') as mock_create_volume:
copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
mock_create_volume.delay.assert_called_with(
mock_image_id, mock_new_snapshot_id)
mock_aws.get_session.assert_called_with(mock_arn)
mock_aws.get_ami.assert_called_with(mock_session, mock_image_id,
mock_region)
mock_aws.get_ami_snapshot_id.assert_called_with(mock_image)
mock_aws.add_snapshot_ownership.assert_called_with(mock_snapshot)
mock_aws.copy_snapshot.assert_called_with(mock_snapshot_id,
mock_region)
def test_copy_ami_snapshot_success_with_reference(self, mock_aws):
"""Assert the snapshot copy task succeeds using a reference AMI ID."""
mock_session = mock_aws.boto3.Session.return_value
mock_account_id = mock_aws.get_session_account_id.return_value
account = account_helper.generate_cloud_account()
arn = account.content_object.account_arn
region = util_helper.get_random_region()
new_image_id = util_helper.generate_dummy_image_id()
# unlike non-reference calls to copy_ami_snapshot, we do NOT want to
# call "account_helper.generate_aws_image(ec2_ami_id=new_image_id)"
# here because cloudigrade has only seen the reference, not the new.
mock_image = util_helper.generate_mock_image(new_image_id)
block_mapping = mock_image.block_device_mappings
mock_snapshot_id = block_mapping[0]["Ebs"]["SnapshotId"]
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, owner_id=mock_account_id)
mock_new_snapshot_id = util_helper.generate_dummy_snapshot_id()
# This is the original ID of a private/shared image.
# It would have been saved to our DB upon initial discovery.
reference_image = account_helper.generate_image()
reference_image_id = reference_image.content_object.ec2_ami_id
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
mock_aws.copy_snapshot.return_value = mock_new_snapshot_id
with patch.object(tasks.imageprep,
"create_volume") as mock_create_volume, patch.object(
tasks.imageprep, "remove_snapshot_ownership"
) as mock_remove_snapshot_ownership:
tasks.copy_ami_snapshot(arn, new_image_id, region,
reference_image_id)
# arn, customer_snapshot_id, snapshot_region, snapshot_copy_id
mock_remove_snapshot_ownership.delay.assert_called_with(
arn, mock_snapshot_id, region, mock_new_snapshot_id)
mock_create_volume.delay.assert_called_with(
reference_image_id, mock_new_snapshot_id)
mock_aws.get_session.assert_called_with(arn)
mock_aws.get_ami.assert_called_with(mock_session, new_image_id, region)
mock_aws.get_ami_snapshot_id.assert_called_with(mock_image)
mock_aws.add_snapshot_ownership.assert_called_with(mock_snapshot)
mock_aws.copy_snapshot.assert_called_with(mock_snapshot_id, region)
# Verify that the copy object was stored correctly to reference later.
copied_image = AwsMachineImageCopy.objects.get(ec2_ami_id=new_image_id)
self.assertIsNotNone(copied_image)
self.assertEqual(
copied_image.reference_awsmachineimage.ec2_ami_id,
reference_image_id,
)
def test_create_volume_snapshot_has_error(self, mock_boto3):
"""Test that volume creation aborts when snapshot has error."""
zone = helper.generate_dummy_availability_zone()
mock_snapshot = helper.generate_mock_snapshot(state='error')
mock_ec2 = mock_boto3.resource.return_value
mock_ec2.Snapshot.return_value = mock_snapshot
with self.assertRaises(AwsSnapshotError):
ec2.create_volume(mock_snapshot.snapshot_id, zone)
mock_boto3.resource.assert_called_once_with('ec2')
mock_ec2.create_volume.assert_not_called()
def test_create_volume_snapshot_not_ready(self, mock_boto3):
"""Test that volume creation aborts when snapshot is not ready."""
zone = helper.generate_dummy_availability_zone()
mock_snapshot = helper.generate_mock_snapshot(state='pending')
mock_ec2 = mock_boto3.resource.return_value
mock_ec2.Snapshot.return_value = mock_snapshot
with self.assertRaises(SnapshotNotReadyException):
ec2.create_volume(mock_snapshot.snapshot_id, zone)
mock_boto3.resource.assert_called_once_with('ec2')
mock_ec2.create_volume.assert_not_called()
def test_copy_snapshot_success(self, mock_boto3):
"""Assert that a snapshot copy operation begins."""
mock_region = random.choice(helper.SOME_AWS_REGIONS)
mock_snapshot = helper.generate_mock_snapshot()
mock_copied_snapshot_id = helper.generate_dummy_snapshot_id()
mock_copy_result = {'SnapshotId': mock_copied_snapshot_id}
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot
mock_snapshot.copy.return_value = mock_copy_result
actual_copied_snapshot_id = ec2.copy_snapshot(
mock_snapshot.snapshot_id, mock_region)
self.assertEqual(actual_copied_snapshot_id, mock_copied_snapshot_id)
def test_create_volume_snapshot_ready(self, mock_boto3):
"""Test that volume creation starts when snapshot is ready."""
zone = helper.generate_dummy_availability_zone()
mock_snapshot = helper.generate_mock_snapshot()
mock_volume = helper.generate_mock_volume()
mock_ec2 = mock_boto3.resource.return_value
mock_ec2.Snapshot.return_value = mock_snapshot
mock_ec2.create_volume.return_value = mock_volume
volume_id = ec2.create_volume(mock_snapshot.snapshot_id, zone)
mock_ec2.create_volume.assert_called_with(
SnapshotId=mock_snapshot.snapshot_id, AvailabilityZone=zone)
mock_boto3.resource.assert_called_once_with('ec2')
self.assertEqual(volume_id, mock_volume.id)
def test_get_snapshot(self):
"""Assert that get_snapshot returns a Snapshot."""
mock_snapshot_id = helper.generate_dummy_snapshot_id()
mock_snapshot = helper.generate_mock_snapshot(mock_snapshot_id)
mock_session = Mock()
mock_resource = mock_session.resource.return_value
mock_resource.Snapshot.return_value = mock_snapshot
mock_region = random.choice(helper.SOME_AWS_REGIONS)
actual_snapshot = ec2.get_snapshot(mock_session, mock_snapshot_id,
mock_region)
self.assertEqual(actual_snapshot, mock_snapshot)
mock_session.resource.assert_called_once_with('ec2',
region_name=mock_region)
mock_resource.Snapshot.assert_called_once_with(mock_snapshot_id)
def test_delete_snapshot_success(self, mock_aws, mock_boto3):
"""Assert that the delete snapshot succeeds."""
mock_snapshot_copy_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot_copy = util_helper.generate_mock_snapshot(
mock_snapshot_copy_id)
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot_copy
volume_id = util_helper.generate_dummy_volume_id()
mock_volume = util_helper.generate_mock_volume(volume_id=volume_id,
state="available")
volume_region = mock_volume.zone[:-1]
mock_aws.get_volume.return_value = mock_volume
mock_aws.check_volume_state.return_value = None
delete_snapshot(mock_snapshot_copy_id, volume_id, volume_region)
def test_copy_ami_snapshot_success(self, mock_aws):
"""Assert that the snapshot copy task succeeds."""
mock_session = mock_aws.boto3.Session.return_value
mock_account_id = mock_aws.get_session_account_id.return_value
mock_arn = util_helper.generate_dummy_arn()
mock_region = util_helper.get_random_region()
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
account_helper.generate_image(ec2_ami_id=mock_image_id)
block_mapping = mock_image.block_device_mappings
mock_snapshot_id = block_mapping[0]["Ebs"]["SnapshotId"]
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, owner_id=mock_account_id)
mock_new_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
mock_aws.copy_snapshot.return_value = mock_new_snapshot_id
with patch.object(tasks.imageprep,
"create_volume") as mock_create_volume:
with patch.object(tasks.imageprep, "remove_snapshot_ownership"
) as mock_remove_snapshot_ownership:
tasks.copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
mock_create_volume.delay.assert_called_with(
mock_image_id, mock_new_snapshot_id)
mock_remove_snapshot_ownership.delay.assert_called_with(
mock_arn,
mock_snapshot_id,
mock_region,
mock_new_snapshot_id,
)
mock_aws.get_session.assert_called_with(mock_arn)
mock_aws.get_ami.assert_called_with(mock_session, mock_image_id,
mock_region)
mock_aws.get_ami_snapshot_id.assert_called_with(mock_image)
mock_aws.add_snapshot_ownership.assert_called_with(mock_snapshot)
mock_aws.copy_snapshot.assert_called_with(mock_snapshot_id,
mock_region)
def test_copy_snapshot_failure(self, mock_boto3):
"""Assert that an error is given when copy fails."""
mock_region = helper.get_random_region()
mock_snapshot = helper.generate_mock_snapshot()
mock_copy_error = {
"Error": {
"Code": "MockError",
"Message": "The operation failed."
}
}
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot
mock_snapshot.copy.side_effect = ClientError(mock_copy_error,
"CopySnapshot")
with self.assertRaises(ClientError):
ec2.copy_snapshot(mock_snapshot.snapshot_id, mock_region)
def test_copy_snapshot_failure(self, mock_boto3):
"""Assert that an error is given when copy fails."""
mock_region = random.choice(helper.SOME_AWS_REGIONS)
mock_snapshot = helper.generate_mock_snapshot()
mock_copy_error = {
'Error': {
'Code': 'MockError',
'Message': 'The operation failed.'
}
}
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot
mock_snapshot.copy.side_effect = ClientError(mock_copy_error,
'CopySnapshot')
with self.assertRaises(ClientError):
ec2.copy_snapshot(mock_snapshot.snapshot_id, mock_region)
def test_copy_ami_snapshot_private_shared(self, mock_aws):
"""Assert that the task copies the image when it is private/shared."""
mock_account_id = util_helper.generate_dummy_aws_account_id()
mock_session = mock_aws.boto3.Session.return_value
mock_aws.get_session_account_id.return_value = mock_account_id
# the account id to use as the private shared image owner
other_account_id = util_helper.generate_dummy_aws_account_id()
mock_region = random.choice(util_helper.SOME_AWS_REGIONS)
mock_arn = util_helper.generate_dummy_arn(mock_account_id, mock_region)
mock_image_id = util_helper.generate_dummy_image_id()
mock_image = util_helper.generate_mock_image(mock_image_id)
mock_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot = util_helper.generate_mock_snapshot(
mock_snapshot_id, encrypted=False, owner_id=other_account_id)
mock_aws.get_session.return_value = mock_session
mock_aws.get_ami.return_value = mock_image
mock_aws.get_ami_snapshot_id.return_value = mock_snapshot_id
mock_aws.get_snapshot.return_value = mock_snapshot
account = AwsAccount(
aws_account_id=mock_account_id,
account_arn=mock_arn,
user=util_helper.generate_test_user(),
)
account.save()
ami = AwsMachineImage.objects.create(account=account,
ec2_ami_id=mock_image_id)
ami.save()
with patch.object(tasks, 'create_volume') as mock_create_volume, \
patch.object(tasks, 'copy_ami_to_customer_account') as \
mock_copy_ami_to_customer_account:
copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
mock_create_volume.delay.assert_not_called()
mock_copy_ami_to_customer_account.delay.assert_called_with(
mock_arn, mock_image_id, mock_region)
def test_copy_snapshot_limit_reached(self, mock_boto3):
"""Assert that an error is returned when the copy limit is reached."""
mock_region = helper.get_random_region()
mock_snapshot = helper.generate_mock_snapshot()
mock_copy_error = {
"Error": {
"Code":
"ResourceLimitExceeded",
"Message":
"You have exceeded an Amazon EC2 resource limit. "
"For example, you might have too many snapshot "
"copies in progress.",
}
}
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot
mock_snapshot.copy.side_effect = ClientError(mock_copy_error,
"CopySnapshot")
with self.assertRaises(AwsSnapshotCopyLimitError):
ec2.copy_snapshot(mock_snapshot.snapshot_id, mock_region)
def test_copy_snapshot_limit_reached(self, mock_boto3):
"""Assert that an error is returned when the copy limit is reached."""
mock_region = random.choice(helper.SOME_AWS_REGIONS)
mock_snapshot = helper.generate_mock_snapshot()
mock_copy_error = {
'Error': {
'Code':
'ResourceLimitExceeded',
'Message':
'You have exceeded an Amazon EC2 resource limit. '
'For example, you might have too many snapshot '
'copies in progress.'
}
}
resource = mock_boto3.resource.return_value
resource.Snapshot.return_value = mock_snapshot
mock_snapshot.copy.side_effect = ClientError(mock_copy_error,
'CopySnapshot')
with self.assertRaises(AwsSnapshotCopyLimitError):
ec2.copy_snapshot(mock_snapshot.snapshot_id, mock_region)
def test_add_snapshot_ownership_not_verified(self):
"""Assert an error is raised when ownership is not verified."""
mock_user_id = str(uuid.uuid4())
mock_snapshot = helper.generate_mock_snapshot()
mock_snapshot.describe_attribute.return_value = {
"CreateVolumePermissions": [],
}
with patch.object(ec2, "_get_primary_account_id") as mock_get_acct_id:
mock_get_acct_id.return_value = mock_user_id
with self.assertRaises(AwsSnapshotNotOwnedError):
ec2.add_snapshot_ownership(mock_snapshot)
expected_permission = {"Add": [{"UserId": mock_user_id}]}
expected_user_ids = [mock_user_id]
mock_snapshot.modify_attribute.assert_called_once_with(
Attribute="createVolumePermission",
CreateVolumePermission=expected_permission,
OperationType="add",
UserIds=expected_user_ids,
)
mock_snapshot.describe_attribute.assert_called_once_with(
Attribute="createVolumePermission")
def test_generate_mock_snapshot_id_encrypted(self):
"""Assert created snapshots obey the encrypted arg."""
snapshot = helper.generate_mock_snapshot(encrypted=True)
self.assertTrue(snapshot.encrypted)
def test_generate_mock_snapshot_id_included(self):
"""Assert snapshots are created with provided id."""
snapshot_id = helper.generate_dummy_snapshot_id()
snapshot = helper.generate_mock_snapshot(snapshot_id)
self.assertEqual(snapshot.snapshot_id, snapshot_id)
def test_generate_mock_snapshot_defaults(self):
"""Assert snapshots are created without provided values."""
snapshot = helper.generate_mock_snapshot()
self.assertIsNotNone(snapshot.snapshot_id)
self.assertIsNotNone(snapshot.encrypted)
