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)