From e2af8bc83691b82da3e60ce35c2219f7a0ceff0a Mon Sep 17 00:00:00 2001
From: Tom Elliff <tomelliff@gmail.com>
Date: Wed, 4 Apr 2018 17:21:08 +0100
Subject: [PATCH] Allow tagging snapshots on creation

Largely copying what was done for volume creation in https://github.com/spulec/moto/pull/1432
---
 moto/ec2/responses/elastic_block_store.py | 13 +++++++
 tests/test_ec2/test_tags.py               | 42 +++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/moto/ec2/responses/elastic_block_store.py b/moto/ec2/responses/elastic_block_store.py
index 2d43f8ffb..cdc5b18e9 100644
--- a/moto/ec2/responses/elastic_block_store.py
+++ b/moto/ec2/responses/elastic_block_store.py
@@ -23,8 +23,11 @@ class ElasticBlockStore(BaseResponse):
     def create_snapshot(self):
         volume_id = self._get_param('VolumeId')
         description = self._get_param('Description')
+        tags = self._parse_tag_specification("TagSpecification")
+        snapshot_tags = tags.get('snapshot', {})
         if self.is_not_dryrun('CreateSnapshot'):
             snapshot = self.ec2_backend.create_snapshot(volume_id, description)
+            snapshot.add_tags(snapshot_tags)
             template = self.response_template(CREATE_SNAPSHOT_RESPONSE)
             return template.render(snapshot=snapshot)
 
@@ -233,6 +236,16 @@ CREATE_SNAPSHOT_RESPONSE = """<CreateSnapshotResponse xmlns="http://ec2.amazonaw
   <volumeSize>{{ snapshot.volume.size }}</volumeSize>
   <description>{{ snapshot.description }}</description>
   <encrypted>{{ snapshot.encrypted }}</encrypted>
+  <tagSet>
+    {% for tag in snapshot.get_tags() %}
+      <item>
+      <resourceId>{{ tag.resource_id }}</resourceId>
+      <resourceType>{{ tag.resource_type }}</resourceType>
+      <key>{{ tag.key }}</key>
+      <value>{{ tag.value }}</value>
+      </item>
+    {% endfor %}
+  </tagSet>
 </CreateSnapshotResponse>"""
 
 DESCRIBE_SNAPSHOTS_RESPONSE = """<DescribeSnapshotsResponse xmlns="http://ec2.amazonaws.com/doc/2013-10-15/">
diff --git a/tests/test_ec2/test_tags.py b/tests/test_ec2/test_tags.py
index d78fe24c3..c92a4f81f 100644
--- a/tests/test_ec2/test_tags.py
+++ b/tests/test_ec2/test_tags.py
@@ -409,3 +409,45 @@ def test_create_volume_with_tags():
         )
     
     assert response['Tags'][0]['Key'] == 'TEST_TAG'
+
+
+@mock_ec2
+def test_create_snapshot_with_tags():
+    client = boto3.client('ec2', 'us-west-2')
+    volume_id = client.create_volume(
+        AvailabilityZone='us-west-2',
+        Encrypted=False,
+        Size=40,
+        TagSpecifications=[
+            {
+                'ResourceType': 'volume',
+                'Tags': [
+                    {
+                        'Key': 'TEST_TAG',
+                        'Value': 'TEST_VALUE'
+                    }
+                ],
+            }
+        ]
+    )['VolumeId']
+    snapshot = client.create_snapshot(
+        VolumeId=volume_id,
+        TagSpecifications=[
+            {
+                'ResourceType': 'snapshot',
+                'Tags': [
+                    {
+                        'Key': 'TEST_SNAPSHOT_TAG',
+                        'Value': 'TEST_SNAPSHOT_VALUE'
+                    }
+                ],
+            }
+        ]
+    )
+
+    expected_tags = [{
+        'Key': 'TEST_SNAPSHOT_TAG',
+        'Value': 'TEST_SNAPSHOT_VALUE'
+    }]
+
+    assert snapshot['Tags'] == expected_tags