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 setUp(self): """Set up fixtures.""" self.ec2_instance_id = util_helper.generate_dummy_instance_id() self.ami_id_a = util_helper.generate_dummy_image_id() self.ami_id_b = util_helper.generate_dummy_image_id() self.ami_id_c = util_helper.generate_dummy_image_id() self.ami_mountpoints = [ (self.ami_id_a, "/dev/sdba"), (self.ami_id_b, "/dev/sdbb"), (self.ami_id_c, "/dev/sdbc"), ]
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_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 = random.choice(util_helper.SOME_AWS_REGIONS) mock_aws.copy_ami.side_effect = ClientError( error_response={ 'Error': { 'Code': 'ItIsAMystery', 'Message': 'Mystery Error', } }, operation_name=Mock(), ) with self.assertRaises(RuntimeError) as e: copy_ami_to_customer_account(arn, reference_ami_id, source_region, True) 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( self, mock_aws, mock_aws_machine_image_objects): """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 = random.choice(util_helper.SOME_AWS_REGIONS) mock_ami = Mock() mock_ami.INSPECTED = 'Inspected' mock_aws_machine_image_objects.get.return_value = mock_ami 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(), ) copy_ami_to_customer_account(arn, reference_ami_id, source_region, True) 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) mock_aws_machine_image_objects.get.assert_called_with( ec2_ami_id=reference_ami_id) self.assertEqual(mock_ami.status, mock_ami.INSPECTED) mock_ami.save.assert_called_once()
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_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_openshift_no_tags(self, mock_aws, mock_check_image_state): """Test case where AMI has no tags.""" mock_session = mock_aws.boto3.Session.return_value ami_id = util_helper.generate_dummy_image_id() ami_region = random.choice(util_helper.SOME_AWS_REGIONS) mock_ami = util_helper.generate_mock_image(ami_id) mock_ami.tags = None mock_resource = mock_session.resource.return_value mock_resource.Image.return_value = mock_ami test_user = util_helper.generate_test_user() test_account = AwsAccount.objects.create( user=test_user, aws_account_id=util_helper.generate_dummy_aws_account_id, account_arn=util_helper.generate_dummy_arn) test_image = AwsMachineImage.objects.create(account=test_account, ec2_ami_id=ami_id) serializer = AwsAccountSerializer() openshift_tag = test_image.tags.filter(description='openshift').first() self.assertEqual(openshift_tag, None) serializer.add_openshift_tag(mock_session, ami_id, ami_region, test_image) openshift_tag = test_image.tags.filter(description='openshift').first() self.assertEqual(openshift_tag, None)
def test_openshift_tag_added(self, mock_aws, mock_check_image_state): """Test openshift tag added to MachineImage.""" mock_session = mock_aws.boto3.Session.return_value ami_id = util_helper.generate_dummy_image_id() ami_region = random.choice(util_helper.SOME_AWS_REGIONS) mock_ami = util_helper.generate_mock_image(ami_id) mock_ami.tags = [{ 'Key': 'cloudigrade-ocp-present', 'Value': 'cloudigrade-ocp-present' }] mock_resource = mock_session.resource.return_value mock_resource.Image.return_value = mock_ami test_user = util_helper.generate_test_user() test_account = AwsAccount.objects.create( user=test_user, aws_account_id=util_helper.generate_dummy_aws_account_id, account_arn=util_helper.generate_dummy_arn) test_image = AwsMachineImage.objects.create(account=test_account, ec2_ami_id=ami_id) serializer = AwsAccountSerializer() openshift_tag = test_image.tags.filter(description='openshift').first() self.assertEqual(openshift_tag, None) serializer.add_openshift_tag(mock_session, ami_id, ami_region, test_image) openshift_tag = test_image.tags.filter(description='openshift').first() self.assertNotEqual(openshift_tag, None)
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_generate_aws_events_with_args_and_some_times(self): """Assert generation of InstanceEvents with all specified args.""" account = helper.generate_aws_account() instance = helper.generate_aws_instance(account) powered_times = ( (None, util_helper.utc_dt(2017, 1, 1)), (util_helper.utc_dt(2017, 1, 2), util_helper.utc_dt(2017, 1, 3)), (util_helper.utc_dt(2017, 1, 4), None), ) ec2_ami_id = util_helper.generate_dummy_image_id() instance_type = random.choice(util_helper.SOME_EC2_INSTANCE_TYPES) subnet = str(uuid.uuid4()) events = helper.generate_aws_instance_events( instance, powered_times, ec2_ami_id=ec2_ami_id, instance_type=instance_type, subnet=subnet, ) self.assertEqual(len(events), 4) # We don't care to check *everything* in the events since that should # already be covered by ``test_generate_events_with_some_times``. # Here we only care that argument values were set correctly. for event in events: self.assertEqual(event.machineimage.ec2_ami_id, ec2_ami_id) self.assertEqual(event.instance_type, instance_type) self.assertEqual(event.subnet, subnet)
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_attach_volumes_to_cluster_instance_not_running( self, mock_boto3, mock_run_inspection_cluster ): """Asserts that an exception is raised if instance exists but is not running.""" ami_id = util_helper.generate_dummy_image_id() volume_id = util_helper.generate_dummy_volume_id() account_helper.generate_image(ec2_ami_id=ami_id, status=MachineImage.PENDING) instance_id = util_helper.generate_dummy_instance_id() mock_list_container_instances = {"containerInstanceArns": [instance_id]} mock_ec2 = Mock() mock_ec2_instance = mock_ec2.Instance.return_value mock_ec2_instance.state = EC2_INSTANCE_STATE_STOPPED mock_ecs = MagicMock() mock_ecs.list_container_instances.return_value = mock_list_container_instances mock_boto3.client.return_value = mock_ecs mock_boto3.resource.return_value = mock_ec2 messages = [{"ami_id": ami_id, "volume_id": volume_id}] with self.assertRaises(AwsECSInstanceNotReady): tasks.attach_volumes_to_cluster(messages) mock_run_inspection_cluster.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_persist_aws_inspection_cluster_results_image_has_errors(self): """Assert that inspection results with image errors are logged.""" ami_id = util_helper.generate_dummy_image_id() api_helper.generate_image( is_encrypted=False, is_windows=False, ec2_ami_id=ami_id ) error_message = _faker.sentence() inspection_results = { "cloud": "aws", "images": { ami_id: { "rhel_found": False, "rhel_version": None, "syspurpose": None, "errors": [error_message], } }, } with self.assertLogs("api.clouds.aws.util", level="INFO") as logging_watcher: util.persist_aws_inspection_cluster_results(inspection_results) self.assertIn( "Error reported in inspection results for image", logging_watcher.output[0], ) self.assertIn(ami_id, logging_watcher.output[0]) self.assertIn(error_message, logging_watcher.output[0]) aws_machine_image = AwsMachineImage.objects.get(ec2_ami_id=ami_id) machine_image = aws_machine_image.machine_image.get() self.assertFalse(machine_image.rhel_detected)
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_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_generate_aws_events_with_args_and_some_times(self): """Assert generation of InstanceEvents with all specified args.""" account = helper.generate_cloud_account() ec2_ami_id = util_helper.generate_dummy_image_id() image = helper.generate_image(ec2_ami_id=ec2_ami_id) instance = helper.generate_instance(account, image=image) powered_times = ( (None, util_helper.utc_dt(2017, 1, 1)), (util_helper.utc_dt(2017, 1, 2), util_helper.utc_dt(2017, 1, 3)), (util_helper.utc_dt(2017, 1, 4), None), ) instance_type = util_helper.get_random_instance_type() subnet = str(uuid.uuid4()) events = helper.generate_instance_events( instance, powered_times, instance_type=instance_type, subnet=subnet, ) self.assertEqual(len(events), 4) # We don't care to check *everything* in the events since that should # already be covered by ``test_generate_events_with_some_times``. # Here we only care that argument values were set correctly. for event in events: if event.event_type != event.TYPE.power_off: self.assertEqual( event.instance.machine_image.content_object.ec2_ami_id, ec2_ami_id ) self.assertEqual(event.content_object.instance_type, instance_type) self.assertEqual(event.content_object.subnet, subnet)
def test_run_inspection_cluster_success(self, mock_boto3): """Asserts successful starting of the houndigrade task.""" mock_list_container_instances = { 'containerInstanceArns': [util_helper.generate_dummy_instance_id()] } mock_ec2 = Mock() mock_ecs = MagicMock() mock_ecs.list_container_instances.return_value = \ mock_list_container_instances mock_boto3.client.return_value = mock_ecs mock_boto3.resource.return_value = mock_ec2 messages = [{ 'ami_id': util_helper.generate_dummy_image_id(), 'volume_id': util_helper.generate_dummy_volume_id() }] tasks.run_inspection_cluster(messages) mock_ecs.list_container_instances.assert_called_once_with( cluster=settings.HOUNDIGRADE_ECS_CLUSTER_NAME) mock_ecs.describe_container_instances.assert_called_once_with( containerInstances=[ mock_list_container_instances['containerInstanceArns'][0] ], cluster=settings.HOUNDIGRADE_ECS_CLUSTER_NAME) mock_ecs.register_task_definition.assert_called_once() mock_ecs.run_task.assert_called_once() mock_ec2.Volume.assert_called_once_with(messages[0]['volume_id']) mock_ec2.Volume.return_value.attach_to_instance.assert_called_once()
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_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_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_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_get_ami_snapshot_id(self): """Assert that an AMI returns a snapshot id.""" mock_image_id = helper.generate_dummy_image_id() mock_image = helper.generate_mock_image(mock_image_id) expected_id = mock_image.block_device_mappings[0]['Ebs']['SnapshotId'] actual_id = ec2.get_ami_snapshot_id(mock_image) self.assertEqual(expected_id, actual_id)
def setUp(self): """Set up basic image objects.""" self.machine_image = helper.generate_image() copy_ec2_ami_id = util_helper.generate_dummy_image_id() api.clouds.aws.util.create_aws_machine_image_copy( copy_ec2_ami_id, self.machine_image.content_object.ec2_ami_id) self.aws_machine_image_copy = aws_models.AwsMachineImageCopy.objects.get( reference_awsmachineimage=self.machine_image.content_object)
def test_objects_exist(self): """Test happy path when both objects exist.""" ec2_ami_id = util_helper.generate_dummy_image_id() expected_machine_image = api_helper.generate_image(ec2_ami_id=ec2_ami_id) expected_aws_machine_image = expected_machine_image.content_object result = util.get_aws_machine_image(ec2_ami_id) self.assertEqual(expected_aws_machine_image, result[0]) self.assertEqual(expected_machine_image, result[1])
def test_attach_volumes_to_cluster_with_no_known_images( self, mock_boto3, mock_scale_down, mock_run_inspection_cluster ): """Assert that inspection is skipped if no known images are given.""" messages = [{"ami_id": util_helper.generate_dummy_image_id()}] tasks.attach_volumes_to_cluster(messages) mock_scale_down.delay.assert_called() mock_boto3.client.assert_not_called() mock_run_inspection_cluster.delay.assert_not_called()
def test_persist_aws_inspection_cluster_results_our_model_is_gone(self): """ Assert that we handle when the AwsMachineImage model has been deleted. This can happen if the customer deletes their cloud account while we are running or waiting on the async inspection. Deleting the cloud account results in the instances and potentially the images being deleted, and when we get the inspection results back for that deleted image, we should just quietly drop the results and move on to other results that may still need processing. """ deleted_ami_id = util_helper.generate_dummy_image_id() ami_id = util_helper.generate_dummy_image_id() api_helper.generate_image( is_encrypted=False, is_windows=False, ec2_ami_id=ami_id ) rhel_version_a = _faker.slug() rhel_version_b = _faker.slug() inspection_results = { "cloud": "aws", "images": { deleted_ami_id: { "rhel_found": True, "rhel_release_files_found": True, "rhel_version": rhel_version_b, }, ami_id: { "rhel_found": True, "rhel_release_files_found": True, "rhel_version": rhel_version_a, }, }, } util.persist_aws_inspection_cluster_results(inspection_results) aws_machine_image = AwsMachineImage.objects.get(ec2_ami_id=ami_id) machine_image = aws_machine_image.machine_image.get() self.assertTrue(machine_image.rhel) self.assertEqual(machine_image.rhel_version, rhel_version_a) with self.assertRaises(AwsMachineImage.DoesNotExist): AwsMachineImage.objects.get(ec2_ami_id=deleted_ami_id)