From f23f7068255a0a3816393f49c5f6a61288f76050 Mon Sep 17 00:00:00 2001
From: Rifqi Al Fatih <rifqialfatih@gmail.com>
Date: Sat, 2 May 2020 18:47:59 +0200
Subject: [PATCH] Implement placement constraints model

---
 moto/ecs/models.py               | 13 ++++++++++++-
 moto/ecs/responses.py            |  2 ++
 tests/test_ecs/test_ecs_boto3.py | 33 ++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/moto/ecs/models.py b/moto/ecs/models.py
index 33d4dcf72..1a385226b 100644
--- a/moto/ecs/models.py
+++ b/moto/ecs/models.py
@@ -121,6 +121,7 @@ class TaskDefinition(BaseObject):
         network_mode=None,
         volumes=None,
         tags=None,
+        placement_constraints=None,
     ):
         self.family = family
         self.revision = revision
@@ -137,6 +138,9 @@ class TaskDefinition(BaseObject):
             self.network_mode = "bridge"
         else:
             self.network_mode = network_mode
+        self.placement_constraints = (
+            placement_constraints if placement_constraints is not None else []
+        )
 
     @property
     def response_object(self):
@@ -558,7 +562,13 @@ class EC2ContainerServiceBackend(BaseBackend):
             raise Exception("{0} is not a cluster".format(cluster_name))
 
     def register_task_definition(
-        self, family, container_definitions, volumes=None, network_mode=None, tags=None
+        self,
+        family,
+        container_definitions,
+        volumes=None,
+        network_mode=None,
+        tags=None,
+        placement_constraints=None,
     ):
         if family in self.task_definitions:
             last_id = self._get_last_task_definition_revision_id(family)
@@ -574,6 +584,7 @@ class EC2ContainerServiceBackend(BaseBackend):
             volumes=volumes,
             network_mode=network_mode,
             tags=tags,
+            placement_constraints=placement_constraints,
         )
         self.task_definitions[family][revision] = task_definition
 
diff --git a/moto/ecs/responses.py b/moto/ecs/responses.py
index 49bf022b4..c8f1e06ce 100644
--- a/moto/ecs/responses.py
+++ b/moto/ecs/responses.py
@@ -63,12 +63,14 @@ class EC2ContainerServiceResponse(BaseResponse):
         volumes = self._get_param("volumes")
         tags = self._get_param("tags")
         network_mode = self._get_param("networkMode")
+        placement_constraints = self._get_param("placementConstraints")
         task_definition = self.ecs_backend.register_task_definition(
             family,
             container_definitions,
             volumes=volumes,
             network_mode=network_mode,
             tags=tags,
+            placement_constraints=placement_constraints,
         )
         return json.dumps({"taskDefinition": task_definition.response_object})
 
diff --git a/tests/test_ecs/test_ecs_boto3.py b/tests/test_ecs/test_ecs_boto3.py
index 7fd90b412..f6de59597 100644
--- a/tests/test_ecs/test_ecs_boto3.py
+++ b/tests/test_ecs/test_ecs_boto3.py
@@ -2604,3 +2604,36 @@ def test_ecs_service_untag_resource_multiple_tags():
         resourceArn=response["service"]["serviceArn"]
     )
     response["tags"].should.equal([{"key": "hello", "value": "world"}])
+
+
+@mock_ecs
+def test_ecs_task_definition_placement_constraints():
+    client = boto3.client("ecs", region_name="us-east-1")
+    response = client.register_task_definition(
+        family="test_ecs_task",
+        containerDefinitions=[
+            {
+                "name": "hello_world",
+                "image": "docker/hello-world:latest",
+                "cpu": 1024,
+                "memory": 400,
+                "essential": True,
+                "environment": [
+                    {"name": "AWS_ACCESS_KEY_ID", "value": "SOME_ACCESS_KEY"}
+                ],
+                "logConfiguration": {"logDriver": "json-file"},
+            }
+        ],
+        networkMode="bridge",
+        tags=[
+            {"key": "createdBy", "value": "moto-unittest"},
+            {"key": "foo", "value": "bar"},
+        ],
+        placementConstraints=[
+            {"type": "memberOf", "expression": "attribute:ecs.instance-type =~ t2.*"}
+        ],
+    )
+    type(response["taskDefinition"]["placementConstraints"]).should.be(list)
+    response["taskDefinition"]["placementConstraints"].should.equal(
+        [{"type": "memberOf", "expression": "attribute:ecs.instance-type =~ t2.*"}]
+    )