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_remove_snapshot_ownership_unexpected_error(
self, mock_aws, mock_boto3):
"""Assert remove snapshot ownership fails due to unexpected error."""
mock_arn = util_helper.generate_dummy_arn()
mock_customer_snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_snapshot_copy_id = util_helper.generate_dummy_snapshot_id()
zone = settings.HOUNDIGRADE_AWS_AVAILABILITY_ZONE
region = zone[:-1]
client_error = ClientError(
error_response={"Error": {
"Code": "InvalidSnapshot.Unknown"
}},
operation_name=Mock(),
)
resource = mock_boto3.resource.return_value
resource.Snapshot.side_effect = client_error
with self.assertRaises(RuntimeError):
tasks.remove_snapshot_ownership(
mock_arn,
mock_customer_snapshot_id,
region,
mock_snapshot_copy_id,
)
mock_aws.remove_snapshot_ownership.assert_not_called()
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_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_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_not_marketplace(self, mock_aws):
"""Assert that an exception is raised when there is an error."""
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
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_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.side_effect = ClientError(
error_response={
'Error': {
'Code': 'ItIsAMystery',
'Message': 'Mystery Error',
}
},
operation_name=Mock(),
)
with self.assertRaises(RuntimeError) as e:
copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
self.assertIn('ClientError', e.exception.args[0])
self.assertIn('ItIsAMystery', e.exception.args[0])
self.assertIn('Mystery Error', e.exception.args[0])
def test_copy_ami_snapshot_not_marketplace(self, mock_aws):
"""Assert that an exception is raised when there is an error."""
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
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()
account_helper.generate_image(ec2_ami_id=mock_image_id)
mock_image = util_helper.generate_mock_image(mock_image_id)
mock_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.side_effect = ClientError(
error_response={
"Error": {
"Code": "ItIsAMystery",
"Message": "Mystery Error"
}
},
operation_name=Mock(),
)
with self.assertRaises(RuntimeError) as e:
tasks.copy_ami_snapshot(mock_arn, mock_image_id, mock_region)
self.assertIn("ClientError", e.exception.args[0])
self.assertIn("ItIsAMystery", e.exception.args[0])
self.assertIn("Mystery Error", e.exception.args[0])
def test_copy_ami_snapshot_marketplace(self, mock_aws):
"""Assert that a suspected marketplace image is checked."""
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
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_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.side_effect = ClientError(
error_response={"Error": {
"Code": "InvalidSnapshot.NotFound"
}},
operation_name=Mock(),
)
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_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_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_generate_dummy_snapshot_id(self):
"""Assert generated id has the appropriate format."""
snapshot_id = helper.generate_dummy_snapshot_id()
hex_part = snapshot_id.split('snap-')[1]
self.assertIn('snap-', snapshot_id)
self.assertEqual(len(hex_part), 17)
for digit in hex_part:
self.assertIn(digit, string.hexdigits)
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_abort_on_snapshot_error(self, mock_aws):
"""Assert that the volume create task does not retry on error."""
ami_id = util_helper.generate_dummy_image_id()
snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_aws.create_volume.side_effect = AwsSnapshotError()
with patch.object(tasks, 'enqueue_ready_volume') as mock_enqueue,\
patch.object(create_volume, 'retry') as mock_retry:
mock_retry.side_effect = Retry()
with self.assertRaises(AwsSnapshotError):
create_volume(ami_id, snapshot_id)
mock_retry.assert_not_called()
mock_enqueue.delay.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_retry_on_snapshot_not_ready(self, mock_aws):
"""Assert that the volume create task retries."""
ami_id = util_helper.generate_dummy_image_id()
snapshot_id = util_helper.generate_dummy_snapshot_id()
mock_aws.create_volume.side_effect = SnapshotNotReadyException(
snapshot_id)
with patch.object(tasks, 'enqueue_ready_volume') as mock_enqueue,\
patch.object(create_volume, 'retry') as mock_retry:
mock_retry.side_effect = Retry()
with self.assertRaises(Retry):
create_volume(ami_id, snapshot_id)
self.assertTrue(mock_retry.called)
mock_enqueue.delay.assert_not_called()
def test_create_volume_success(self, mock_aws):
"""Assert that the volume create task succeeds."""
ami_id = util_helper.generate_dummy_image_id()
snapshot_id = util_helper.generate_dummy_snapshot_id()
zone = settings.HOUNDIGRADE_AWS_AVAILABILITY_ZONE
region = zone[:-1]
mock_volume = util_helper.generate_mock_volume()
mock_aws.create_volume.return_value = mock_volume.id
mock_aws.get_region_from_availability_zone.return_value = region
with patch.object(tasks, 'enqueue_ready_volume') as mock_enqueue:
create_volume(ami_id, snapshot_id)
mock_enqueue.delay.assert_called_with(ami_id, mock_volume.id,
region)
mock_aws.create_volume.assert_called_with(snapshot_id, zone)
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_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_ami_snapshot_marketplace(self, mock_aws):
"""Assert that a suspected marketplace image is checked."""
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
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_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.side_effect = ClientError(
error_response={'Error': {
'Code': 'InvalidSnapshot.NotFound'
}},
operation_name=Mock(),
)
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, maybe_marketplace=True)
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)
