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, 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_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_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_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_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_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_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_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_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 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_generate_mock_image(self): """Assert generated image contains given value.""" image_id = helper.generate_dummy_image_id() encrypted = random.choice((True, False)) image = helper.generate_mock_image(image_id, encrypted) self.assertEqual(image.image_id, image_id) self.assertIsNotNone(image.root_device_name) self.assertIsNotNone(image.root_device_type) self.assertIsInstance(image.block_device_mappings, list) self.assertIsInstance(image.block_device_mappings[0], dict)
def test_check_image_state_mystery_load_failure(self): """Assert raised exception when image loading fails mysteriously.""" error_response = { "Error": { "Code": "itisamystery.gif", } } exception = ClientError(error_response, Mock()) mock_image = helper.generate_mock_image() mock_image.load.side_effect = exception with self.assertRaises(ClientError): ec2.check_image_state(mock_image)
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_check_image_state_not_found(self): """Assert raised exception when image is not found.""" # Dummy response inspired by real exception recorded at # https://sentry.io/organizations/cloudigrade/issues/970892568/ error_response = { "Error": { "Code": "InvalidAMIID.NotFound", "Message": "The image id '[POTATO]' does not exist", } } exception = ClientError(error_response, Mock()) mock_image = helper.generate_mock_image() mock_image.load.side_effect = exception with self.assertRaises(AwsImageError): ec2.check_image_state(mock_image)
def test_get_ami(self): """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 = random.choice(helper.SOME_AWS_REGIONS) 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)
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_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_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_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_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 = 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_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, 'create_volume') as mock_create_volume: with patch.object(tasks, 'remove_snapshot_ownership') as \ mock_remove_snapshot_ownership: 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_check_image_state_unhandled(self): """Assert raised exception when image state is unhandled.""" mock_image = helper.generate_mock_image(state='itisamystery.gif') with self.assertRaises(ImageNotReadyException): ec2.check_image_state(mock_image)
def test_check_image_state_failed(self): """Assert raised exception when image state is failed.""" mock_image = helper.generate_mock_image(state='failed') with self.assertRaises(AwsImageError): ec2.check_image_state(mock_image)
def test_check_image_state_available(self): """Assert clean return when image state is available.""" mock_image = helper.generate_mock_image(state='available') ec2.check_image_state(mock_image)
def test_check_image_state_no_meta(self): """Assert raised exception when image has no metadata.""" mock_image = helper.generate_mock_image() mock_image.meta.data = None with self.assertRaises(AwsImageError): ec2.check_image_state(mock_image)