From b4b63202d9bb4879db151c38a778692f23ab8f05 Mon Sep 17 00:00:00 2001
From: David Morrison <drmorr@yelp.com>
Date: Wed, 23 Jan 2019 15:59:31 -0800
Subject: [PATCH] tag specifications for spot fleet requests

---
 moto/ec2/models.py                | 35 ++++++++++++++++++++----
 moto/ec2/responses/spot_fleets.py | 15 +++++++++++
 tests/test_ec2/test_spot_fleet.py | 44 ++++++++++++++++++++++++++++++-
 3 files changed, 88 insertions(+), 6 deletions(-)
 mode change 100755 => 100644 moto/ec2/models.py

diff --git a/moto/ec2/models.py b/moto/ec2/models.py
old mode 100755
new mode 100644
index a450943b7..cd20e7e0e
--- a/moto/ec2/models.py
+++ b/moto/ec2/models.py
@@ -2879,7 +2879,7 @@ class SpotInstanceRequest(BotoSpotRequest, TaggedEC2Resource):
     def __init__(self, ec2_backend, spot_request_id, price, image_id, type,
                  valid_from, valid_until, launch_group, availability_zone_group,
                  key_name, security_groups, user_data, instance_type, placement,
-                 kernel_id, ramdisk_id, monitoring_enabled, subnet_id, spot_fleet_id,
+                 kernel_id, ramdisk_id, monitoring_enabled, subnet_id, tags, spot_fleet_id,
                  **kwargs):
         super(SpotInstanceRequest, self).__init__(**kwargs)
         ls = LaunchSpecification()
@@ -2903,6 +2903,7 @@ class SpotInstanceRequest(BotoSpotRequest, TaggedEC2Resource):
         ls.monitored = monitoring_enabled
         ls.subnet_id = subnet_id
         self.spot_fleet_id = spot_fleet_id
+        self.tags = tags
 
         if security_groups:
             for group_name in security_groups:
@@ -2936,6 +2937,7 @@ class SpotInstanceRequest(BotoSpotRequest, TaggedEC2Resource):
             security_group_names=[],
             security_group_ids=self.launch_specification.groups,
             spot_fleet_id=self.spot_fleet_id,
+            tags=self.tags,
         )
         instance = reservation.instances[0]
         return instance
@@ -2951,15 +2953,16 @@ class SpotRequestBackend(object):
                                valid_until, launch_group, availability_zone_group,
                                key_name, security_groups, user_data,
                                instance_type, placement, kernel_id, ramdisk_id,
-                               monitoring_enabled, subnet_id, spot_fleet_id=None):
+                               monitoring_enabled, subnet_id, tags=None, spot_fleet_id=None):
         requests = []
+        tags = tags or {}
         for _ in range(count):
             spot_request_id = random_spot_request_id()
             request = SpotInstanceRequest(self,
                                           spot_request_id, price, image_id, type, valid_from, valid_until,
                                           launch_group, availability_zone_group, key_name, security_groups,
                                           user_data, instance_type, placement, kernel_id, ramdisk_id,
-                                          monitoring_enabled, subnet_id, spot_fleet_id)
+                                          monitoring_enabled, subnet_id, tags, spot_fleet_id)
             self.spot_instance_requests[spot_request_id] = request
             requests.append(request)
         return requests
@@ -2979,8 +2982,8 @@ class SpotRequestBackend(object):
 
 class SpotFleetLaunchSpec(object):
     def __init__(self, ebs_optimized, group_set, iam_instance_profile, image_id,
-                 instance_type, key_name, monitoring, spot_price, subnet_id, user_data,
-                 weighted_capacity):
+                 instance_type, key_name, monitoring, spot_price, subnet_id, tag_specifications,
+                 user_data, weighted_capacity):
         self.ebs_optimized = ebs_optimized
         self.group_set = group_set
         self.iam_instance_profile = iam_instance_profile
@@ -2990,6 +2993,7 @@ class SpotFleetLaunchSpec(object):
         self.monitoring = monitoring
         self.spot_price = spot_price
         self.subnet_id = subnet_id
+        self.tag_specifications = tag_specifications
         self.user_data = user_data
         self.weighted_capacity = float(weighted_capacity)
 
@@ -3020,6 +3024,7 @@ class SpotFleetRequest(TaggedEC2Resource):
                 monitoring=spec.get('monitoring._enabled'),
                 spot_price=spec.get('spot_price', self.spot_price),
                 subnet_id=spec['subnet_id'],
+                tag_specifications=self._parse_tag_specifications(spec),
                 user_data=spec.get('user_data'),
                 weighted_capacity=spec['weighted_capacity'],
             )
@@ -3102,6 +3107,7 @@ class SpotFleetRequest(TaggedEC2Resource):
                 monitoring_enabled=launch_spec.monitoring,
                 subnet_id=launch_spec.subnet_id,
                 spot_fleet_id=self.id,
+                tags=launch_spec.tag_specifications,
             )
             self.spot_requests.extend(requests)
         self.fulfilled_capacity += added_weight
@@ -3124,6 +3130,25 @@ class SpotFleetRequest(TaggedEC2Resource):
         self.spot_requests = [req for req in self.spot_requests if req.instance.id not in instance_ids]
         self.ec2_backend.terminate_instances(instance_ids)
 
+    def _parse_tag_specifications(self, spec):
+        try:
+            tag_spec_num = max([int(key.split('.')[1]) for key in spec if key.startswith("tag_specification_set")])
+        except ValueError:  # no tag specifications
+            return {}
+
+        tag_specifications = {}
+        for si in range(1, tag_spec_num + 1):
+            resource_type = spec["tag_specification_set.{si}._resource_type".format(si=si)]
+
+            tags = [key for key in spec if key.startswith("tag_specification_set.{si}._tag".format(si=si))]
+            tag_num = max([int(key.split('.')[3]) for key in tags])
+            tag_specifications[resource_type] = dict((
+                spec["tag_specification_set.{si}._tag.{ti}._key".format(si=si, ti=ti)],
+                spec["tag_specification_set.{si}._tag.{ti}._value".format(si=si, ti=ti)],
+            ) for ti in range(1, tag_num + 1))
+
+        return tag_specifications
+
 
 class SpotFleetBackend(object):
     def __init__(self):
diff --git a/moto/ec2/responses/spot_fleets.py b/moto/ec2/responses/spot_fleets.py
index 0366af9d6..bb9aeb4ca 100644
--- a/moto/ec2/responses/spot_fleets.py
+++ b/moto/ec2/responses/spot_fleets.py
@@ -107,6 +107,21 @@ DESCRIBE_SPOT_FLEET_TEMPLATE = """<DescribeSpotFleetRequestsResponse xmlns="http
                             </item>
                             {% endfor %}
                         </groupSet>
+                        <tagSpecificationSet>
+                            {% for resource_type in launch_spec.tag_specifications %}
+                            <item>
+                                <resourceType>{{ resource_type }}</resourceType>
+                                <tag>
+                                {% for key, value in launch_spec.tag_specifications[resource_type].items() %}
+                                    <item>
+                                        <key>{{ key }}</key>
+                                        <value>{{ value }}</value>
+                                    </item>
+                                {% endfor %}
+                                </tag>
+                            </item>
+                            {% endfor %}
+                        </tagSpecificationSet>
                     </item>
                     {% endfor %}
                 </launchSpecifications>
diff --git a/tests/test_ec2/test_spot_fleet.py b/tests/test_ec2/test_spot_fleet.py
index a2bd1d061..190f3b1f1 100644
--- a/tests/test_ec2/test_spot_fleet.py
+++ b/tests/test_ec2/test_spot_fleet.py
@@ -54,7 +54,7 @@ def spot_config(subnet_id, allocation_strategy="lowestPrice"):
             },
             'EbsOptimized': False,
             'WeightedCapacity': 2.0,
-            'SpotPrice': '0.13'
+            'SpotPrice': '0.13',
         }, {
             'ImageId': 'ami-123',
             'KeyName': 'my-key',
@@ -148,6 +148,48 @@ def test_create_diversified_spot_fleet():
     instances[0]['InstanceId'].should.contain("i-")
 
 
+@mock_ec2
+def test_create_spot_fleet_request_with_tag_spec():
+    conn = boto3.client("ec2", region_name='us-west-2')
+    subnet_id = get_subnet_id(conn)
+
+    tag_spec = [
+        {
+            'ResourceType': 'instance',
+            'Tags': [
+                {
+                    'Key': 'tag-1',
+                    'Value': 'foo',
+                },
+                {
+                    'Key': 'tag-2',
+                    'Value': 'bar',
+                },
+            ]
+        },
+    ]
+    config = spot_config(subnet_id)
+    config['LaunchSpecifications'][0]['TagSpecifications'] = tag_spec
+    spot_fleet_res = conn.request_spot_fleet(
+        SpotFleetRequestConfig=config
+    )
+    spot_fleet_id = spot_fleet_res['SpotFleetRequestId']
+    spot_fleet_requests = conn.describe_spot_fleet_requests(
+        SpotFleetRequestIds=[spot_fleet_id])['SpotFleetRequestConfigs']
+    spot_fleet_config = spot_fleet_requests[0]['SpotFleetRequestConfig']
+    spot_fleet_config['LaunchSpecifications'][0]['TagSpecifications'][0][
+            'ResourceType'].should.equal('instance')
+    for tag in tag_spec[0]['Tags']:
+        spot_fleet_config['LaunchSpecifications'][0]['TagSpecifications'][0]['Tags'].should.contain(tag)
+
+    instance_res = conn.describe_spot_fleet_instances(
+        SpotFleetRequestId=spot_fleet_id)
+    instances = conn.describe_instances(InstanceIds=[i['InstanceId'] for i in instance_res['ActiveInstances']])
+    for instance in instances['Reservations'][0]['Instances']:
+        for tag in tag_spec[0]['Tags']:
+            instance['Tags'].should.contain(tag)
+
+
 @mock_ec2
 def test_cancel_spot_fleet_request():
     conn = boto3.client("ec2", region_name='us-west-2')