Merge pull request release-engineering#38 from JAVGan/multitags

AWS: Allow setting tags to snapshots or AMIs only

Author: Rohan McGovern

cloudimg/aws.py

@@ -83,6 +83,12 @@ class AWSPublishingMetadata(PublishingMetadata):
         billing_products (list, optional): Billing product identifiers
 
         boot_mode (str, optional): The boot mode for booting up the AMI on EC2.
+
+        snapshot_tags (dict, optional): Tags to be applied to the snapshot
+        import only.
+
+        ami_tags (dict, optional): Tags to be applied to the registered AMI
+        only.
     """
 
     def __init__(self, *args, **kwargs):
@@ -91,6 +97,8 @@ def __init__(self, *args, **kwargs):
         self.billing_products = kwargs.pop('billing_products', None)
         bmode_str = kwargs.pop('boot_mode', None) or "not_set"
         self.boot_mode = AWSBootMode[bmode_str]
+        self.snapshot_tags = kwargs.pop('snapshot_tags', None)
+        self.ami_tags = kwargs.pop('ami_tags', None)
 
         super(AWSPublishingMetadata, self).__init__(*args, **kwargs)
 
@@ -564,6 +572,13 @@ def publish(self, metadata):
         Returns:
             An EC2 Image
         """
+        def add_tags(tag_parameter_name, extra_kwargs):
+            new_tags = getattr(metadata, tag_parameter_name, None)
+            if new_tags:
+                tags = extra_kwargs.get("tags") or {}
+                new_tags.update(tags)
+                extra_kwargs.update({"tags": new_tags})
+
         log.info('Searching for image: %s', metadata.image_name)
         image = (
             self.get_image_by_name(metadata.image_name) or
@@ -584,8 +599,7 @@ def publish(self, metadata):
 
                 # Set tags when they're provided
                 extra_kwargs = {}
-                if metadata.tags:
-                    extra_kwargs.update({"tags": metadata.tags})
+                add_tags("tags", extra_kwargs)
 
                 if not obj:
                     log.info('Object does not exist: %s', metadata.object_name)
@@ -596,6 +610,7 @@ def publish(self, metadata):
                 else:
                     log.info('Object already exists')
 
+                add_tags("snapshot_tags", extra_kwargs)
                 snapshot = self.import_snapshot(obj,
                                                 metadata.snapshot_name,
                                                 **extra_kwargs)
@@ -771,8 +786,10 @@ def register_image(self, snapshot, metadata):
             BillingProducts=metadata.billing_products,
             **optional_kwargs,
         )
-        if metadata.tags:
-            self.tag_image(image, metadata.tags)
+        if metadata.tags or metadata.ami_tags:
+            tags = metadata.tags or {}
+            tags.update(metadata.ami_tags or {})
+            self.tag_image(image, tags)
         return image
 
     def share_image(self, image, accounts=[], groups=[]):

tests/test_aws.py

@@ -700,6 +700,85 @@ def test_publish_tags(self,
                                                     self.md.object_name,
                                                     tags=tags)
 
+    @patch('cloudimg.aws.AWSService.upload_to_container')
+    @patch('cloudimg.aws.AWSService.import_snapshot')
+    @patch('cloudimg.aws.AWSService.register_image')
+    @patch('cloudimg.aws.AWSService.share_image')
+    @patch('cloudimg.aws.AWSService.get_image_by_tags')
+    @patch('cloudimg.aws.AWSService.get_image_by_name')
+    @patch('cloudimg.aws.AWSService.get_snapshot_by_name')
+    @patch('cloudimg.aws.AWSService.get_object_by_name')
+    def test_publish_snapshot_tags(self,
+                                   get_object_by_name,
+                                   get_snapshot_by_name,
+                                   get_image_by_name,
+                                   get_image_by_tags,
+                                   share_image,
+                                   register_image,
+                                   import_snapshot,
+                                   upload_to_container):
+        get_image_by_name.return_value = None
+        get_image_by_tags.return_value = None
+        get_snapshot_by_name.return_value = None
+        get_object_by_name.return_value = None
+        self.md.tags = None
+        self.md.snapshot_tags = {"snapshot": "tag"}
+        published = self.svc.publish(self.md)
+
+        share_image.assert_called_once_with(published,
+                                            accounts=[],
+                                            groups=[])
+        self.assertEqual(register_image.call_count, 1)
+        import_snapshot.assert_called_once_with(
+            ANY,
+            self.md.snapshot_name,
+            tags={"snapshot": "tag"},
+        )
+        upload_to_container.assert_called_once_with(self.md.image_path,
+                                                    self.md.container,
+                                                    self.md.object_name)
+
+    @patch('cloudimg.aws.AWSService.upload_to_container')
+    @patch('cloudimg.aws.AWSService.import_snapshot')
+    @patch('cloudimg.aws.AWSService.register_image')
+    @patch('cloudimg.aws.AWSService.share_image')
+    @patch('cloudimg.aws.AWSService.get_image_by_tags')
+    @patch('cloudimg.aws.AWSService.get_image_by_name')
+    @patch('cloudimg.aws.AWSService.get_snapshot_by_name')
+    @patch('cloudimg.aws.AWSService.get_object_by_name')
+    def test_publish_merged_snapshot_tags(
+        self,
+        get_object_by_name,
+        get_snapshot_by_name,
+        get_image_by_name,
+        get_image_by_tags,
+        share_image,
+        register_image,
+        import_snapshot,
+        upload_to_container,
+    ):
+        get_image_by_name.return_value = None
+        get_image_by_tags.return_value = None
+        get_snapshot_by_name.return_value = None
+        get_object_by_name.return_value = None
+        self.md.tags = {"foo": "bar"}
+        self.md.snapshot_tags = {"snapshot": "tag"}
+        published = self.svc.publish(self.md)
+
+        share_image.assert_called_once_with(published,
+                                            accounts=[],
+                                            groups=[])
+        self.assertEqual(register_image.call_count, 1)
+        import_snapshot.assert_called_once_with(
+            ANY,
+            self.md.snapshot_name,
+            tags={"foo": "bar", "snapshot": "tag"},
+        )
+        upload_to_container.assert_called_once_with(self.md.image_path,
+                                                    self.md.container,
+                                                    self.md.object_name,
+                                                    tags={"foo": "bar"})
+
     @patch('cloudimg.aws.AWSService.tag_image')
     def test_register_image_no_tags(self, tag_image):
         self.mock_register_image.return_value = "fakeimg"
@@ -741,6 +820,33 @@ def test_register_image_tags(self, tag_image):
         tag_image.assert_called_once_with("fakeimg", self.md.tags)
         self.assertEqual(res, "fakeimg")
 
+    @patch('cloudimg.aws.AWSService.tag_image')
+    def test_register_image_ami_tags(self, tag_image):
+        self.md.tags = None
+        self.md.ami_tags = {"ami": "tag"}
+        self.mock_register_image.return_value = "fakeimg"
+
+        res = self.svc.register_image(MagicMock(), self.md)
+
+        self.mock_register_image.assert_called_once()
+        tag_image.assert_called_once_with("fakeimg", {"ami": "tag"})
+        self.assertEqual(res, "fakeimg")
+
+    @patch('cloudimg.aws.AWSService.tag_image')
+    def test_register_image_merged_ami_tags(self, tag_image):
+        self.md.tags = {"tag": "tag"}
+        self.md.ami_tags = {"ami": "tag"}
+        self.mock_register_image.return_value = "fakeimg"
+
+        res = self.svc.register_image(MagicMock(), self.md)
+
+        self.mock_register_image.assert_called_once()
+        tag_image.assert_called_once_with(
+            "fakeimg",
+            {"tag": "tag", "ami": "tag"},
+        )
+        self.assertEqual(res, "fakeimg")
+
     @patch('cloudimg.aws.AWSService.tag_image')
     def test_register_image_boot_mode(self, tag_image):
         self.mock_register_image.return_value = "fakeimg"