From 04c5198a0c9f748273e9cc6fea87e9a79bdded27 Mon Sep 17 00:00:00 2001 From: Jessie Nadler Date: Fri, 17 Nov 2017 13:25:08 -0500 Subject: [PATCH] Add default ecs attributes and format in response obj (#1346) * Initialize EC2ContainerServiceBackend and ContainerInstance objects with region_name * Initialize ContainerInstance with default attributes * These attributes are automatically applied by ECS when a container is registered * Docs: http://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-placement-constraints.html#attributes * Format container_instance.attributes for response_object * Python3 * Only use available ECS regions for ecs_backends * Sort dictionaries on key='name' using lambda * Sort all dicts in tests using lambda --- AUTHORS.md | 1 + moto/ecs/models.py | 38 +++++++--- tests/test_ecs/test_ecs_boto3.py | 119 ++++++++++++++++++++++++++++++- 3 files changed, 148 insertions(+), 10 deletions(-) diff --git a/AUTHORS.md b/AUTHORS.md index f4160146c..55ac102d5 100644 --- a/AUTHORS.md +++ b/AUTHORS.md @@ -47,3 +47,4 @@ Moto is written by Steve Pulec with contributions from: * [Adam Stauffer](https://github.com/adamstauffer) * [Guy Templeton](https://github.com/gjtempleton) * [Michael van Tellingen](https://github.com/mvantellingen) +* [Jessie Nadler](https://github.com/nadlerjessie) diff --git a/moto/ecs/models.py b/moto/ecs/models.py index b44184033..e0b29cb01 100644 --- a/moto/ecs/models.py +++ b/moto/ecs/models.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals import uuid from datetime import datetime from random import random, randint +import boto3 import pytz from moto.core.exceptions import JsonRESTError @@ -261,7 +262,7 @@ class Service(BaseObject): class ContainerInstance(BaseObject): - def __init__(self, ec2_instance_id): + def __init__(self, ec2_instance_id, region_name): self.ec2_instance_id = ec2_instance_id self.agent_connected = True self.status = 'ACTIVE' @@ -321,14 +322,29 @@ class ContainerInstance(BaseObject): 'agentHash': '4023248', 'dockerVersion': 'DockerVersion: 1.5.0' } - - self.attributes = {} + ec2_backend = ec2_backends[region_name] + ec2_instance = ec2_backend.get_instance(ec2_instance_id) + self.attributes = { + 'ecs.ami-id': ec2_instance.image_id, + 'ecs.availability-zone': ec2_instance.placement, + 'ecs.instance-type': ec2_instance.instance_type, + 'ecs.os-type': ec2_instance.platform if ec2_instance.platform == 'windows' else 'linux' # options are windows and linux, linux is default + } @property def response_object(self): response_object = self.gen_response_object() + response_object['attributes'] = [self._format_attribute(name, value) for name, value in response_object['attributes'].items()] return response_object + def _format_attribute(self, name, value): + formatted_attr = { + 'name': name, + } + if value is not None: + formatted_attr['value'] = value + return formatted_attr + class ContainerInstanceFailure(BaseObject): @@ -347,12 +363,19 @@ class ContainerInstanceFailure(BaseObject): class EC2ContainerServiceBackend(BaseBackend): - def __init__(self): + def __init__(self, region_name): + super(EC2ContainerServiceBackend, self).__init__() self.clusters = {} self.task_definitions = {} self.tasks = {} self.services = {} self.container_instances = {} + self.region_name = region_name + + def reset(self): + region_name = self.region_name + self.__dict__ = {} + self.__init__(region_name) def describe_task_definition(self, task_definition_str): task_definition_name = task_definition_str.split('/')[-1] @@ -669,7 +692,7 @@ class EC2ContainerServiceBackend(BaseBackend): cluster_name = cluster_str.split('/')[-1] if cluster_name not in self.clusters: raise Exception("{0} is not a cluster".format(cluster_name)) - container_instance = ContainerInstance(ec2_instance_id) + container_instance = ContainerInstance(ec2_instance_id, self.region_name) if not self.container_instances.get(cluster_name): self.container_instances[cluster_name] = {} container_instance_id = container_instance.container_instance_arn.split( @@ -866,6 +889,5 @@ class EC2ContainerServiceBackend(BaseBackend): yield task_fam -ecs_backends = {} -for region, ec2_backend in ec2_backends.items(): - ecs_backends[region] = EC2ContainerServiceBackend() +available_regions = boto3.session.Session().get_available_regions("ecs") +ecs_backends = {region: EC2ContainerServiceBackend(region) for region in available_regions} diff --git a/tests/test_ecs/test_ecs_boto3.py b/tests/test_ecs/test_ecs_boto3.py index 9e5e4ff08..5fcc297aa 100644 --- a/tests/test_ecs/test_ecs_boto3.py +++ b/tests/test_ecs/test_ecs_boto3.py @@ -1624,11 +1624,13 @@ def test_attributes(): clusterName=test_cluster_name ) + instances = [] test_instance = ec2.create_instances( ImageId="ami-1234abcd", MinCount=1, MaxCount=1, )[0] + instances.append(test_instance) instance_id_document = json.dumps( ec2_utils.generate_instance_identity_document(test_instance) @@ -1648,6 +1650,7 @@ def test_attributes(): MinCount=1, MaxCount=1, )[0] + instances.append(test_instance) instance_id_document = json.dumps( ec2_utils.generate_instance_identity_document(test_instance) @@ -1680,7 +1683,10 @@ def test_attributes(): targetType='container-instance' ) attrs = resp['attributes'] - len(attrs).should.equal(4) + + NUM_CUSTOM_ATTRIBUTES = 4 # 2 specific to individual machines and 1 global, going to both machines (2 + 1*2) + NUM_DEFAULT_ATTRIBUTES = 4 + len(attrs).should.equal(NUM_CUSTOM_ATTRIBUTES + (NUM_DEFAULT_ATTRIBUTES * len(instances))) # Tests that the attrs have been set properly len(list(filter(lambda item: item['name'] == 'env', attrs))).should.equal(2) @@ -1692,13 +1698,14 @@ def test_attributes(): {'name': 'attr1', 'value': 'instance2', 'targetId': partial_arn2, 'targetType': 'container-instance'} ] ) + NUM_CUSTOM_ATTRIBUTES -= 1 resp = ecs_client.list_attributes( cluster=test_cluster_name, targetType='container-instance' ) attrs = resp['attributes'] - len(attrs).should.equal(3) + len(attrs).should.equal(NUM_CUSTOM_ATTRIBUTES + (NUM_DEFAULT_ATTRIBUTES * len(instances))) @mock_ecs @@ -1757,6 +1764,114 @@ def test_list_task_definition_families(): len(resp2['families']).should.equal(1) +@mock_ec2 +@mock_ecs +def test_default_container_instance_attributes(): + ecs_client = boto3.client('ecs', region_name='us-east-1') + ec2 = boto3.resource('ec2', region_name='us-east-1') + + test_cluster_name = 'test_ecs_cluster' + + # Create cluster and EC2 instance + _ = ecs_client.create_cluster( + clusterName=test_cluster_name + ) + + test_instance = ec2.create_instances( + ImageId="ami-1234abcd", + MinCount=1, + MaxCount=1, + )[0] + + instance_id_document = json.dumps( + ec2_utils.generate_instance_identity_document(test_instance) + ) + + # Register container instance + response = ecs_client.register_container_instance( + cluster=test_cluster_name, + instanceIdentityDocument=instance_id_document + ) + + response['containerInstance'][ + 'ec2InstanceId'].should.equal(test_instance.id) + full_arn = response['containerInstance']['containerInstanceArn'] + container_instance_id = full_arn.rsplit('/', 1)[-1] + + default_attributes = response['containerInstance']['attributes'] + assert len(default_attributes) == 4 + expected_result = [ + {'name': 'ecs.availability-zone', 'value': test_instance.placement['AvailabilityZone']}, + {'name': 'ecs.ami-id', 'value': test_instance.image_id}, + {'name': 'ecs.instance-type', 'value': test_instance.instance_type}, + {'name': 'ecs.os-type', 'value': test_instance.platform or 'linux'} + ] + assert sorted(default_attributes, key=lambda item: item['name']) == sorted(expected_result, key=lambda item: item['name']) + + +@mock_ec2 +@mock_ecs +def test_describe_container_instances_with_attributes(): + ecs_client = boto3.client('ecs', region_name='us-east-1') + ec2 = boto3.resource('ec2', region_name='us-east-1') + + test_cluster_name = 'test_ecs_cluster' + + # Create cluster and EC2 instance + _ = ecs_client.create_cluster( + clusterName=test_cluster_name + ) + + test_instance = ec2.create_instances( + ImageId="ami-1234abcd", + MinCount=1, + MaxCount=1, + )[0] + + instance_id_document = json.dumps( + ec2_utils.generate_instance_identity_document(test_instance) + ) + + # Register container instance + response = ecs_client.register_container_instance( + cluster=test_cluster_name, + instanceIdentityDocument=instance_id_document + ) + + response['containerInstance'][ + 'ec2InstanceId'].should.equal(test_instance.id) + full_arn = response['containerInstance']['containerInstanceArn'] + container_instance_id = full_arn.rsplit('/', 1)[-1] + default_attributes = response['containerInstance']['attributes'] + + # Set attributes on container instance, one without a value + attributes = [ + {'name': 'env', 'value': 'prod'}, + {'name': 'attr1', 'value': 'instance1', 'targetId': container_instance_id, 'targetType': 'container-instance'}, + {'name': 'attr_without_value'} + ] + ecs_client.put_attributes( + cluster=test_cluster_name, + attributes=attributes + ) + + # Describe container instance, should have attributes previously set + described_instance = ecs_client.describe_container_instances(cluster=test_cluster_name, containerInstances=[container_instance_id]) + + assert len(described_instance['containerInstances']) == 1 + assert isinstance(described_instance['containerInstances'][0]['attributes'], list) + + # Remove additional info passed to put_attributes + cleaned_attributes = [] + for attribute in attributes: + attribute.pop('targetId', None) + attribute.pop('targetType', None) + cleaned_attributes.append(attribute) + described_attributes = sorted(described_instance['containerInstances'][0]['attributes'], key=lambda item: item['name']) + expected_attributes = sorted(default_attributes + cleaned_attributes, key=lambda item: item['name']) + assert described_attributes == expected_attributes + + def _fetch_container_instance_resources(container_instance_description): remaining_resources = {} registered_resources = {}