def test_describe_instances(self):
"""Assert that describe_instances returns a dict of instances data."""
instance_ids = [
helper.generate_dummy_instance_id(),
helper.generate_dummy_instance_id(),
helper.generate_dummy_instance_id(),
helper.generate_dummy_instance_id(),
]
individual_described_instances = [
helper.generate_dummy_describe_instance(instance_id)
for instance_id in instance_ids
]
response = {
"Reservations": [
{
"Instances": individual_described_instances[:2],
},
{
"Instances": individual_described_instances[2:],
},
],
}
mock_session = Mock()
mock_client = mock_session.client.return_value
mock_client.describe_instances.return_value = response
region = helper.get_random_region()
described_instances = ec2.describe_instances(mock_session,
instance_ids, region)
self.assertEqual(set(described_instances.keys()), set(instance_ids))
for described_instance in individual_described_instances:
self.assertIn(described_instance, described_instances.values())
def test_generate_aws_ami_messages_deduplicate(self):
"""Test that messages are also deduplicated."""
region = util_helper.get_random_region()
image_id_1 = util_helper.generate_dummy_image_id()
image_id_2 = util_helper.generate_dummy_image_id()
instance_1 = util_helper.generate_dummy_describe_instance(
image_id=image_id_1)
instance_2 = util_helper.generate_dummy_describe_instance(
image_id=image_id_1)
instance_3 = util_helper.generate_dummy_describe_instance(
image_id=image_id_1)
instance_4 = util_helper.generate_dummy_describe_instance(
image_id=image_id_2)
instances_data = {
region: [instance_1, instance_2, instance_3, instance_4]
}
ami_list = [image_id_1, image_id_2]
expected = [
{
"cloud_provider": AWS_PROVIDER_STRING,
"region": region,
"image_id": image_id_1,
},
{
"cloud_provider": AWS_PROVIDER_STRING,
"region": region,
"image_id": image_id_2,
},
]
result = generate_aws_ami_messages(instances_data, ami_list)
self.assertEqual(len(result), len(expected))
for message in expected:
self.assertIn(message, result)
def test_copy_ami_to_customer_account_missing_image(self, mock_aws):
"""Assert early return if the AwsMachineImage doesn't exist."""
arn = util_helper.generate_dummy_arn()
reference_ami_id = util_helper.generate_dummy_image_id()
region = util_helper.get_random_region()
tasks.copy_ami_to_customer_account(arn, reference_ami_id, region)
mock_aws.get_session.assert_not_called()
def test_save_instance_with_available_image(self):
"""Test that save instance events also writes image on instance."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
account = api_helper.generate_cloud_account(
arn=arn, aws_account_id=aws_account_id)
region = util_helper.get_random_region()
instances_data = {
region: [
util_helper.generate_dummy_describe_instance(
state=aws.InstanceState.running)
]
}
ami_id = instances_data[region][0]["ImageId"]
mock_session = Mock()
described_amis = util_helper.generate_dummy_describe_image(
image_id=ami_id, owner_id=aws_account_id)
with patch.object(util.aws, "describe_images") as mock_describe_images:
mock_describe_images.return_value = [described_amis]
util.create_new_machine_images(mock_session, instances_data)
mock_describe_images.assert_called_with(mock_session, {ami_id},
region)
awsinstance = util.save_instance(account, instances_data[region][0],
region)
instance = awsinstance.instance.get()
self.assertEqual(instance.machine_image.content_object.ec2_ami_id,
ami_id)
self.assertEqual(instance.machine_image.status, MachineImage.PENDING)
def test_copy_ami_to_customer_account_private_no_copy_dot_error(
self, mock_aws):
"""Assert that the task marks private (no copy) image as in error."""
arn = util_helper.generate_dummy_arn()
reference_ami_id = util_helper.generate_dummy_image_id()
source_region = util_helper.get_random_region()
image = api_helper.generate_image(ec2_ami_id=reference_ami_id,
status=MachineImage.INSPECTING)
mock_reference_ami = Mock()
mock_reference_ami.public = False
mock_aws.get_ami.return_value = mock_reference_ami
mock_aws.copy_ami.side_effect = ClientError(
error_response={
"Error": {
"Code":
"InvalidRequest",
"Message":
"You do not have permission to access the "
"storage of this ami.",
}
},
operation_name=Mock(),
)
tasks.copy_ami_to_customer_account(arn, reference_ami_id,
source_region)
mock_aws.get_session.assert_called_with(arn)
mock_aws.get_ami.assert_called_with(mock_aws.get_session.return_value,
reference_ami_id, source_region)
image.refresh_from_db()
self.assertEqual(image.status, MachineImage.ERROR)
def test_create_new_machine_images_with_windows_image(self):
"""Test that new windows machine images are marked appropriately."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
api_helper.generate_cloud_account(arn=arn,
aws_account_id=aws_account_id)
region = util_helper.get_random_region()
instances_data = {
region: [
util_helper.generate_dummy_describe_instance(
state=aws.InstanceState.running)
]
}
instances_data[region][0]["Platform"] = "Windows"
ami_id = instances_data[region][0]["ImageId"]
mock_session = Mock()
described_amis = util_helper.generate_dummy_describe_image(
image_id=ami_id, owner_id=aws_account_id)
with patch.object(util.aws, "describe_images") as mock_describe_images:
mock_describe_images.return_value = [described_amis]
result = util.create_new_machine_images(mock_session,
instances_data)
mock_describe_images.assert_called_with(mock_session, {ami_id},
region)
self.assertEqual(result, [ami_id])
images = list(AwsMachineImage.objects.all())
self.assertEqual(len(images), 1)
self.assertEqual(images[0].ec2_ami_id, ami_id)
self.assertEqual(images[0].platform, AwsMachineImage.WINDOWS)
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_to_customer_account_not_marketplace(self, mock_aws):
"""Assert that the task fails when non-marketplace error occurs."""
arn = util_helper.generate_dummy_arn()
reference_ami_id = util_helper.generate_dummy_image_id()
source_region = util_helper.get_random_region()
api_helper.generate_image(ec2_ami_id=reference_ami_id)
mock_aws.copy_ami.side_effect = ClientError(
error_response={
"Error": {
"Code": "ItIsAMystery",
"Message": "Mystery Error"
}
},
operation_name=Mock(),
)
with self.assertRaises(RuntimeError) as e:
tasks.copy_ami_to_customer_account(arn, reference_ami_id,
source_region)
self.assertIn("ClientError", e.exception.args[0])
self.assertIn("ItIsAMystery", e.exception.args[0])
self.assertIn("Mystery Error", e.exception.args[0])
mock_aws.get_session.assert_called_with(arn)
mock_aws.get_ami.assert_called_with(mock_aws.get_session.return_value,
reference_ami_id, source_region)
def test_copy_ami_to_customer_account_marketplace_with_dot_error(
self, mock_aws):
"""Assert that the task marks marketplace image as inspected."""
arn = util_helper.generate_dummy_arn()
reference_ami_id = util_helper.generate_dummy_image_id()
source_region = util_helper.get_random_region()
image = api_helper.generate_image(ec2_ami_id=reference_ami_id,
status=MachineImage.INSPECTING)
mock_aws.copy_ami.side_effect = ClientError(
error_response={
"Error": {
"Code":
"InvalidRequest",
"Message":
"Images with EC2 BillingProduct codes cannot "
"be copied to another AWS account.",
}
},
operation_name=Mock(),
)
tasks.copy_ami_to_customer_account(arn, reference_ami_id,
source_region)
mock_aws.get_session.assert_called_with(arn)
mock_aws.get_ami.assert_called_with(mock_aws.get_session.return_value,
reference_ami_id, source_region)
image.refresh_from_db()
aws_image = image.content_object
aws_image.refresh_from_db()
self.assertEqual(image.status, MachineImage.INSPECTED)
self.assertTrue(aws_image.aws_marketplace_image)
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_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_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_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_get_volume(self, mock_boto3):
"""Test that a Volume is returned."""
region = helper.get_random_region()
zone = helper.generate_dummy_availability_zone(region)
volume_id = helper.generate_dummy_volume_id()
mock_volume = helper.generate_mock_volume(volume_id=volume_id,
zone=zone)
resource = mock_boto3.resource.return_value
resource.Volume.return_value = mock_volume
actual_volume = ec2.get_volume(volume_id, region)
self.assertEqual(actual_volume, mock_volume)
def test_generate_aws_instance_with_args(self):
"""Assert generation of an AwsInstance with all specified args."""
account = helper.generate_cloud_account()
ec2_instance_id = util_helper.generate_dummy_instance_id()
region = util_helper.get_random_region()
instance = helper.generate_instance(
account,
ec2_instance_id=ec2_instance_id,
region=region,
)
self.assertIsInstance(instance, Instance)
self.assertEqual(instance.cloud_account, account)
self.assertEqual(instance.content_object.ec2_instance_id, ec2_instance_id)
self.assertEqual(instance.content_object.region, region)
def test_copy_snapshot_success(self, mock_boto3):
"""Assert that a snapshot copy operation begins."""
mock_region = helper.get_random_region()
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_copy_ami_abort_when_no_image_loaded(self):
"""Test that image copy aborts when no image is loaded."""
mock_session = Mock()
mock_ec2_client = mock_session.client.return_value
image_id = helper.generate_dummy_image_id()
source_region = helper.get_random_region()
with patch.object(ec2, "get_ami") as mock_get_ami:
mock_get_ami.return_value = None
result = ec2.copy_ami(mock_session, image_id, source_region)
self.assertIsNone(result)
mock_ec2_client.copy_image.assert_not_called()
mock_ec2_client.create_tags.assert_not_called()
def test_get_ami_when_load_fails(self, mock_check_image_state):
"""Assert that get_ami returns None when load fails."""
mock_image_id = helper.generate_dummy_image_id()
mock_image = helper.generate_mock_image(mock_image_id)
mock_session = Mock()
mock_resource = mock_session.resource.return_value
mock_resource.Image.return_value = mock_image
mock_region = helper.get_random_region()
mock_check_image_state.side_effect = AwsImageError
actual_image = ec2.get_ami(mock_session, mock_image_id, mock_region)
self.assertIsNone(actual_image)
def test_get_region_from_availability_zone(self, mock_client):
"""Assert that the proper region is returned for an AZ."""
expected_region = test_helper.get_random_region()
zone = test_helper.generate_dummy_availability_zone(expected_region)
az_response = {
"AvailabilityZones": [
{"State": "available", "RegionName": expected_region, "ZoneName": zone}
]
}
mock_desc = mock_client.return_value.describe_availability_zones
mock_desc.return_value = az_response
actual_region = helper.get_region_from_availability_zone(zone)
self.assertEqual(expected_region, actual_region)
def test_generate_aws_ami_messages(self):
"""Test that messages are formatted correctly."""
region = util_helper.get_random_region()
instance = util_helper.generate_dummy_describe_instance()
instances_data = {region: [instance]}
ami_list = [instance["ImageId"]]
expected = [{
"cloud_provider": AWS_PROVIDER_STRING,
"region": region,
"image_id": instance["ImageId"],
}]
result = generate_aws_ami_messages(instances_data, ami_list)
self.assertEqual(result, expected)
def generate_cloudtrail_record(
aws_account_id,
event_name,
event_time=None,
region=None,
request_parameters=None,
response_elements=None,
):
"""
Generate an example CloudTrail log's "Record" dict.
This function needs something in request_parameters or response_elements to
produce actually meaningful output, but this is not strictly required.
Args:
aws_account_id (int): The AWS account ID.
event_name (str): optional AWS event name.
event_time (datetime.datetime): optional time when the even occurred.
region (str): optional AWS region in which the event occurred.
request_parameters (dict): optional data for 'requestParameters' key.
response_elements (dict): optional data for 'responseElements' key.
Returns:
dict: Data that looks like a CloudTrail log Record.
"""
if not region:
region = helper.get_random_region()
if not event_time:
event_time = helper.utc_dt(2018, 1, 1, 0, 0, 0)
if not request_parameters:
request_parameters = {}
if not response_elements:
response_elements = {}
record = {
"awsRegion": region,
"eventName": event_name,
"eventSource": "ec2.amazonaws.com",
"eventTime": event_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
"requestParameters": request_parameters,
"responseElements": response_elements,
"userIdentity": {
"accountId": int(aws_account_id)
},
}
return record
def test_copy_ami_to_customer_account_save_error_when_image_load_fails(
self, mock_aws):
"""Assert that we save error status if image load fails."""
arn = util_helper.generate_dummy_arn()
reference_ami_id = util_helper.generate_dummy_image_id()
snapshot_region = util_helper.get_random_region()
image = api_helper.generate_image(ec2_ami_id=reference_ami_id)
mock_aws.get_ami.return_value = None
with patch.object(tasks.imageprep, "copy_ami_snapshot") as mock_copy:
tasks.copy_ami_to_customer_account(arn, reference_ami_id,
snapshot_region)
mock_copy.delay.assert_not_called()
mock_aws.copy_ami.assert_not_called()
image.refresh_from_db()
self.assertEquals(image.status, image.ERROR)
def test_copy_ami(self):
"""Test that an image is copied via the boto session successfully."""
mock_session = Mock()
mock_ec2_client = mock_session.client.return_value
mock_original_image = Mock()
mock_copied_image_dict = helper.generate_mock_image_dict()
mock_ec2_client.copy_image.return_value = mock_copied_image_dict
image_id = mock_original_image.id
source_region = helper.get_random_region()
with patch.object(ec2, "get_ami") as mock_get_ami:
mock_get_ami.return_value = mock_original_image
result = ec2.copy_ami(mock_session, image_id, source_region)
mock_ec2_client.copy_image.assert_called_once()
mock_ec2_client.create_tags.assert_called_once()
self.assertEqual(result, mock_copied_image_dict["ImageId"])
def test_get_ami(self, mock_check_image_state):
"""Assert that get_ami returns an Image."""
mock_image_id = helper.generate_dummy_image_id()
mock_image = helper.generate_mock_image(mock_image_id)
mock_session = Mock()
mock_resource = mock_session.resource.return_value
mock_resource.Image.return_value = mock_image
mock_region = helper.get_random_region()
actual_image = ec2.get_ami(mock_session, mock_image_id, mock_region)
self.assertEqual(actual_image, mock_image)
mock_session.resource.assert_called_once_with("ec2",
region_name=mock_region)
mock_resource.Image.assert_called_once_with(mock_image_id)
mock_check_image_state.assert_called_once_with(mock_image)
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 = helper.get_random_region()
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_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_save_instance_with_missing_machineimage(self):
"""
Test that save_instance works around a missing MachineImage.
This is a use case that we *shouldn't* need, but until we identify how
exactly AwsMachineImage objects are being saved without their matching
MachineImage objects, we have this workaround to create the
MachineImage at the time it is needed.
"""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
account = api_helper.generate_cloud_account(
arn=arn, aws_account_id=aws_account_id)
region = util_helper.get_random_region()
instances_data = {
region: [
util_helper.generate_dummy_describe_instance(
state=aws.InstanceState.running)
]
}
ami_id = instances_data[region][0]["ImageId"]
# Create the AwsMachineImage without its paired MachineImage.
aws_machine_image = AwsMachineImage.objects.create(
owner_aws_account_id=util_helper.generate_dummy_aws_account_id(),
ec2_ami_id=ami_id,
platform="none",
)
# Verify that the MachineImage was *not* created before proceeding.
with self.assertRaises(MachineImage.DoesNotExist):
aws_machine_image.machine_image.get()
awsinstance = util.save_instance(account, instances_data[region][0],
region)
instance = awsinstance.instance.get()
self.assertEqual(instance.machine_image.content_object.ec2_ami_id,
ami_id)
self.assertEqual(instance.machine_image.status,
MachineImage.UNAVAILABLE)
def test_copy_ami_snapshot_save_error_when_image_snapshot_id_get_fails(
self, mock_aws):
"""Assert that we save error status if snapshot id is not available."""
arn = util_helper.generate_dummy_arn()
ami_id = util_helper.generate_dummy_image_id()
snapshot_region = util_helper.get_random_region()
image = account_helper.generate_image(ec2_ami_id=ami_id)
mock_aws.get_ami.return_value = image
mock_aws.get_ami_snapshot_id.return_value = None
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:
tasks.copy_ami_snapshot(arn, ami_id, snapshot_region)
mock_create_volume.delay.assert_not_called()
mock_copy_ami_to_customer_account.delay.assert_not_called()
image.refresh_from_db()
self.assertEquals(image.status, image.ERROR)
