def generate_aws_account(arn=None, aws_account_id=None, user=None):
"""
Generate an AwsAccount for testing.
Any optional arguments not provided will be randomly generated.
Args:
arn (str): Optional ARN.
aws_account_id (12-digit string): Optional AWS account ID.
user (User): Optional Django auth User to be this account's owner.
Returns:
AwsAccount: The created AwsAccount.
"""
if arn is None:
if aws_account_id is None:
arn = helper.generate_dummy_arn(generate_account_id=True)
else:
arn = helper.generate_dummy_arn(account_id=aws_account_id)
if user is None:
user = helper.generate_test_user()
return AwsAccount.objects.create(
account_arn=arn,
aws_account_id=aws.AwsArn(arn).account_id,
user=user,
)
def test_success(self, mock_primary_id, mock_tasks_aws, mock_create):
"""Assert the task happy path upon normal operation."""
# User that would ultimately own the created objects.
user = User.objects.create()
# Dummy values for the various interactions.
session_account_id = util_helper.generate_dummy_aws_account_id()
auth_id = _faker.pyint()
application_id = _faker.pyint()
source_id = _faker.pyint()
customer_secret_access_key = util_helper.generate_dummy_arn(
account_id=session_account_id)
primary_account_id = util_helper.generate_dummy_aws_account_id()
mock_primary_id.return_value = primary_account_id
mock_tasks_aws.get_session_account_id.return_value = session_account_id
mock_tasks_aws.AwsArn.return_value.account_id = session_account_id
# Fake out the policy verification.
policy_name = _faker.slug()
policy_arn = util_helper.generate_dummy_arn(
account_id=session_account_id)
mock_ensure_policy = mock_tasks_aws.ensure_cloudigrade_policy
mock_ensure_policy.return_value = (policy_name, policy_arn)
# Fake out the role verification.
role_name = _faker.slug()
role_arn = util_helper.generate_dummy_arn(
account_id=session_account_id)
mock_ensure_role = mock_tasks_aws.ensure_cloudigrade_role
mock_ensure_role.return_value = (role_name, role_arn)
tasks.configure_customer_aws_and_create_cloud_account(
user.username,
customer_secret_access_key,
auth_id,
application_id,
source_id,
)
cloud_account_name = api_util.get_standard_cloud_account_name(
"aws", session_account_id)
mock_create.assert_called_with(
user,
role_arn,
cloud_account_name,
auth_id,
application_id,
source_id,
)
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_create_fails_when_arn_access_denied(self):
"""Test that an account is not saved if ARN access is denied."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
validated_data = {
'account_arn': arn,
}
mock_request = Mock()
mock_request.user = util_helper.generate_test_user()
context = {'request': mock_request}
client_error = ClientError(
error_response={'Error': {
'Code': 'AccessDenied'
}},
operation_name=Mock(),
)
with patch.object(aws.sts, 'boto3') as mock_boto3:
mock_assume_role = mock_boto3.client.return_value.assume_role
mock_assume_role.side_effect = client_error
serializer = AwsAccountSerializer(context=context)
with self.assertRaises(ValidationError) as cm:
serializer.create(validated_data)
raised_exception = cm.exception
self.assertIn('account_arn', raised_exception.detail)
self.assertIn(arn, raised_exception.detail['account_arn'][0])
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_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_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_create_fails_when_account_not_verified(self):
"""Test that an account is not saved if verification fails."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
role = util_helper.generate_dummy_role()
validated_data = {
'account_arn': arn,
}
mock_request = Mock()
mock_request.user = util_helper.generate_test_user()
context = {'request': mock_request}
failed_actions = ['foo', 'bar']
with patch.object(aws, 'verify_account_access') as mock_verify, \
patch.object(aws.sts, 'boto3') as mock_boto3:
mock_assume_role = mock_boto3.client.return_value.assume_role
mock_assume_role.return_value = role
mock_verify.return_value = False, failed_actions
serializer = AwsAccountSerializer(context=context)
with self.assertRaises(serializers.ValidationError) as cm:
serializer.create(validated_data)
exception = cm.exception
self.assertIn('account_arn', exception.detail)
for index in range(len(failed_actions)):
self.assertIn(failed_actions[index],
exception.detail['account_arn'][index + 1])
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)
account = AwsAccount(
account_arn=arn,
aws_account_id=aws_account_id,
user=util_helper.generate_test_user(),
)
account.save()
region = random.choice(util_helper.SOME_AWS_REGIONS)
running_instances = {
region: [
util_helper.generate_dummy_describe_instance(
state=aws.InstanceState.running
)
]
}
running_instances[region][0]['Platform'] = 'Windows'
ami_id = running_instances[region][0]['ImageId']
result = util.create_new_machine_images(account, running_instances)
amis = AwsMachineImage.objects.filter(account=account).all()
self.assertEqual(result, [ami_id])
for ami in amis:
self.assertEqual(ami.ec2_ami_id, ami_id)
self.assertEqual(ImageTag.objects.filter(
description='windows').first(),
ami.tags.filter(description='windows').first())
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_create_aws_clount_fails_with_different_user_same_account_id(
self, mock_enable, mock_notify_sources):
"""Test create_aws_cloud_account failure message for different user."""
# The first call just creates the existing objects.
util.create_aws_cloud_account(self.user, self.arn, self.name,
self.auth_id, self.app_id,
self.source_id)
other_arn = util_helper.generate_dummy_arn(
account_id=self.aws_account_id, resource=_faker.word())
other_auth_id = _faker.pyint()
other_app_id = _faker.pyint()
other_source_id = _faker.pyint()
other_user = util_helper.generate_test_user()
with self.assertRaises(ValidationError) as raise_context:
util.create_aws_cloud_account(
other_user,
other_arn,
self.name,
other_auth_id,
other_app_id,
other_source_id,
)
exception_detail = raise_context.exception.detail
self.assertIn("Could not set up cloud metering.",
exception_detail["account_arn"])
self.assertNotIn("CG1002", exception_detail["account_arn"])
def setUp(self):
"""Set up shared values for kafka creation tests."""
self.account_number = str(_faker.pyint())
self.username = _faker.user_name()
self.authentication_id = _faker.pyint()
self.application_id = _faker.pyint()
self.source_id = _faker.pyint()
self.cloudigrade_sources_app_id = _faker.pyint()
(
self.message,
self.headers,
) = util_helper.generate_applicationauthentication_create_message_value(
account_number=self.account_number,
application_id=self.application_id,
authentication_id=self.authentication_id,
)
self.arn = util_helper.generate_dummy_arn()
self.auth_return_value = {
"username": self.arn,
"resource_type": settings.SOURCES_RESOURCE_TYPE,
"resource_id": self.application_id,
"id": self.authentication_id,
"authtype": settings.SOURCES_CLOUDMETER_ARN_AUTHTYPE,
}
def test_create_from_sources_kafka_message_success(
self,
mock_task,
mock_get_app,
mock_get_app_type_id,
mock_get_auth,
):
"""Assert create_from_sources_kafka_message happy path success."""
mock_get_app.return_value = {
"application_type_id": self.cloudigrade_sources_app_id,
"source_id": self.source_id,
}
mock_get_app_type_id.return_value = self.cloudigrade_sources_app_id
arn = util_helper.generate_dummy_arn()
self.auth_return_value["username"] = arn
mock_get_auth.return_value = self.auth_return_value
tasks.create_from_sources_kafka_message(self.message, self.headers)
user = User.objects.get(username=self.account_number)
mock_task.delay.assert_called_with(
user.username,
arn,
self.authentication_id,
self.application_id,
self.source_id,
)
def test_update_from_sources_kafka_message_updates_arn_but_disables_cloud_account(
self, mock_get_auth, mock_notify_sources, mock_get_app):
"""
Assert update_from_source_kafka_message updates the arn and disables clount.
This can happen if we get a well-formed ARN, but the AWS-side verification fails
due to something like badly configured IAM Role or Policy. In this case, we do
want to save the updated ARN, but we need to disable the Cloud Account until it
can verify its permissions (at a later date).
"""
message, headers = util_helper.generate_authentication_create_message_value(
self.account_number, self.username, self.authentication_id)
new_arn = util_helper.generate_dummy_arn(account_id=self.account_id)
self.auth_return_value["username"] = new_arn
mock_get_auth.return_value = self.auth_return_value
mock_get_app.return_value = self.app_return_value
validation_error = ValidationError(
detail={_faker.slug(): _faker.slug()})
with patch("api.clouds.aws.util.verify_permissions"
) as mock_verify_permissions:
mock_verify_permissions.side_effect = validation_error
tasks.update_from_source_kafka_message(message, headers)
mock_verify_permissions.assert_called()
self.clount.refresh_from_db()
self.assertEqual(self.clount.content_object.account_arn, new_arn)
self.assertFalse(self.clount.is_enabled)
def setUp(self):
"""Set up shared variables."""
self.authentication_id = _faker.pyint()
self.account_id = util_helper.generate_dummy_aws_account_id()
self.arn = util_helper.generate_dummy_arn(account_id=self.account_id)
self.clount = api_helper.generate_cloud_account(
arn=self.arn, platform_authentication_id=self.authentication_id)
self.username = _faker.user_name()
self.application_id = _faker.pyint()
self.source_id = _faker.pyint()
self.account_number = str(_faker.pyint())
self.auth_return_value = {
"username": self.arn,
"resource_type": settings.SOURCES_RESOURCE_TYPE,
"resource_id": self.application_id,
"id": self.authentication_id,
"authtype": settings.SOURCES_CLOUDMETER_ARN_AUTHTYPE,
}
self.app_return_value = {
"source_id": self.source_id,
"id": self.application_id,
}
def test_update_from_sources_kafka_message_new_aws_account_id(
self, mock_enable, mock_get_auth, mock_notify_sources,
mock_get_app):
"""
Assert the new cloud account created for new aws_account_id.
A new CloudAccount will get created with the new arn. And the old
CloudAccount will be removed.
"""
message, headers = util_helper.generate_authentication_create_message_value(
self.account_number, self.username, self.authentication_id)
new_account_id = util_helper.generate_dummy_aws_account_id()
new_arn = util_helper.generate_dummy_arn(account_id=new_account_id)
self.auth_return_value["username"] = new_arn
mock_get_auth.return_value = self.auth_return_value
mock_get_app.return_value = self.app_return_value
with patch.object(sts, "boto3") as mock_boto3, patch.object(
aws_models, "_delete_cloudtrail"), patch(
"api.clouds.aws.util.verify_permissions"):
mock_assume_role = mock_boto3.client.return_value.assume_role
mock_assume_role.return_value = util_helper.generate_dummy_role()
tasks.update_from_source_kafka_message(message, headers)
mock_enable.assert_called()
def test_account_not_created_if_enable_fails(self, mock_enable,
mock_notify_sources):
"""
Assert the account is not created if enable fails.
Note:
AwsCloudAccount.enable is responsible for performing AWS permission
verification. This test effectively tests handling of AWS failures there.
"""
user = User.objects.create()
customer_secret_access_key = util_helper.generate_dummy_arn()
auth_id = _faker.pyint()
application_id = _faker.pyint()
source_id = _faker.pyint()
validation_error = ValidationError({"account_arn": "uh oh"})
mock_enable.side_effect = validation_error
tasks.configure_customer_aws_and_create_cloud_account(
user.username,
customer_secret_access_key,
auth_id,
application_id,
source_id,
)
self.assertFalse(CloudAccount.objects.filter(user=user).exists())
def test_create_aws_clount_fails_with_same_user_same_name(
self, mock_enable, mock_notify_sources):
"""Test create_aws_cloud_account fails with same name different platform IDs."""
# The first call just creates the existing objects.
util.create_aws_cloud_account(self.user, self.arn, self.name,
self.auth_id, self.app_id,
self.source_id)
other_arn = util_helper.generate_dummy_arn()
other_auth_id = _faker.pyint()
other_app_id = _faker.pyint()
other_source_id = _faker.pyint()
with self.assertRaises(ValidationError) as raise_context:
util.create_aws_cloud_account(
self.user,
other_arn,
self.name,
other_auth_id,
other_app_id,
other_source_id,
)
exception_detail = raise_context.exception.detail
self.assertIn("name", exception_detail)
self.assertIn("Could not set up cloud metering.",
exception_detail["name"])
self.assertIn("CG1003", exception_detail["name"])
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_create_fails_when_assume_role_fails_unexpectedly(self):
"""Test that account is not saved if assume_role fails unexpectedly."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
validated_data = {
'account_arn': arn,
}
mock_request = Mock()
mock_request.user = util_helper.generate_test_user()
context = {'request': mock_request}
client_error = ClientError(
error_response=MagicMock(),
operation_name=Mock(),
)
with patch.object(aws.sts, 'boto3') as mock_boto3:
mock_assume_role = mock_boto3.client.return_value.assume_role
mock_assume_role.side_effect = client_error
serializer = AwsAccountSerializer(context=context)
with self.assertRaises(ClientError) as cm:
serializer.create(validated_data)
raised_exception = cm.exception
self.assertEqual(raised_exception, client_error)
def test_create_aws_clount_fails_same_aws_account_different_arn(
self, mock_enable, mock_notify_sources):
"""
Test create_aws_cloud_account fails with same ARN and different AWS Account.
If we ever change cloudigrade to support multiple ARNs in the same AWS Account,
this test and its underlying logic must be rewritten.
"""
util.create_aws_cloud_account(self.user, self.arn, self.name,
self.auth_id, self.app_id,
self.source_id)
other_name = self.name = _faker.word()
other_arn = util_helper.generate_dummy_arn(
account_id=self.aws_account_id, resource=_faker.word())
self.assertNotEqual(self.arn, other_arn)
with self.assertRaises(ValidationError) as raise_context:
util.create_aws_cloud_account(
self.user,
other_arn,
other_name,
self.auth_id,
self.app_id,
self.source_id,
)
exception_detail = raise_context.exception.detail
self.assertIn("account_arn", exception_detail)
self.assertIn("Could not set up cloud metering.",
exception_detail["account_arn"])
self.assertIn("CG1002", exception_detail["account_arn"])
def test_get_running_instances(self):
"""Assert we get expected instances in a dict keyed by regions."""
mock_arn = helper.generate_dummy_arn()
mock_regions = [f'region-{uuid.uuid4()}']
mock_role = helper.generate_dummy_role()
mock_running_instance = helper.generate_dummy_describe_instance(
state=aws.InstanceState.running)
mock_stopped_instance = helper.generate_dummy_describe_instance(
state=aws.InstanceState.stopped)
mock_described = {
'Reservations': [
{
'Instances': [
mock_running_instance,
mock_stopped_instance,
],
},
],
}
expected_found = {mock_regions[0]: [mock_running_instance]}
with patch.object(aws, 'get_regions') as mock_get_regions, \
patch.object(aws, 'boto3') as mock_boto3:
mock_assume_role = mock_boto3.client.return_value.assume_role
mock_assume_role.return_value = mock_role
mock_get_regions.return_value = mock_regions
mock_client = mock_boto3.Session.return_value.client.return_value
mock_client.describe_instances.return_value = mock_described
actual_found = aws.get_running_instances(aws.get_session(mock_arn))
self.assertDictEqual(expected_found, actual_found)
def test_create_new_machine_images(self):
"""Test that new machine images are saved to the DB."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(aws_account_id)
account = AwsAccount(
account_arn=arn,
aws_account_id=aws_account_id,
user=util_helper.generate_test_user(),
)
account.save()
region = random.choice(util_helper.SOME_AWS_REGIONS)
running_instances = {
region: [
util_helper.generate_dummy_describe_instance(
state=aws.InstanceState.running
)
]
}
ami_id = running_instances[region][0]['ImageId']
result = util.create_new_machine_images(account, running_instances)
amis = AwsMachineImage.objects.filter(account=account).all()
self.assertEqual(result, [ami_id])
for ami in amis:
self.assertEqual(ami.ec2_ami_id, ami_id)
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_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_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_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_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_parse_arn_with_region_and_account(self):
"""Assert successful account ID parsing from a well-formed ARN."""
mock_account_id = helper.generate_dummy_aws_account_id()
mock_arn = helper.generate_dummy_arn(account_id=mock_account_id,
region="test-region-1")
arn_object = AwsArn(mock_arn)
partition = arn_object.partition
self.assertIsNotNone(partition)
service = arn_object.service
self.assertIsNotNone(service)
region = arn_object.region
self.assertIsNotNone(region)
account_id = arn_object.account_id
self.assertIsNotNone(account_id)
resource_type = arn_object.resource_type
self.assertIsNotNone(resource_type)
resource_separator = arn_object.resource_separator
self.assertIsNotNone(resource_separator)
resource = arn_object.resource
self.assertIsNotNone(resource)
reconstructed_arn = ("arn:" + partition + ":" + service + ":" +
region + ":" + str(account_id) + ":" +
resource_type + resource_separator + resource)
self.assertEqual(mock_account_id, account_id)
self.assertEqual(mock_arn, reconstructed_arn)
def setUp(self):
"""Set up test models."""
aws_account_id = util_helper.generate_dummy_aws_account_id()
arn = util_helper.generate_dummy_arn(account_id=aws_account_id)
self.role = util_helper.generate_dummy_role()
self.account = helper.generate_cloud_account(
arn=arn, aws_account_id=aws_account_id, name="test")
