From 9aef6944da6a6f1b84be99cefb34f1724e6ca58f Mon Sep 17 00:00:00 2001
From: Miki Watanabe <105326591+MikiPWata@users.noreply.github.com>
Date: Mon, 5 Feb 2024 16:16:58 -0500
Subject: [PATCH] ECS: Fix container memory validation (#7280)

---
 moto/ecs/models.py               | 13 ++++++++++---
 tests/test_ecs/test_ecs_boto3.py | 26 ++++++++++++++++++++++----
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/moto/ecs/models.py b/moto/ecs/models.py
index d983f8551..1d98a4c2b 100644
--- a/moto/ecs/models.py
+++ b/moto/ecs/models.py
@@ -1178,8 +1178,9 @@ class EC2ContainerServiceBackend(BaseBackend):
                 raise EcsClientException(
                     f"Tasks using the Fargate launch type do not support pidMode '{pid_mode}'. The supported value for pidMode is 'task'."
                 )
-
-        self._validate_container_defs(memory, container_definitions)
+        self._validate_container_defs(
+            memory, container_definitions, requires_compatibilities
+        )
 
         if family in self.task_definitions:
             last_id = self._get_last_task_definition_revision_id(family)
@@ -1214,13 +1215,19 @@ class EC2ContainerServiceBackend(BaseBackend):
         return task_definition
 
     @staticmethod
-    def _validate_container_defs(memory: Optional[str], container_definitions: List[Dict[str, Any]]) -> None:  # type: ignore[misc]
+    def _validate_container_defs(memory: Optional[str], container_definitions: List[Dict[str, Any]], requires_compatibilities: Optional[List[str]]) -> None:  # type: ignore[misc]
         # The capitalised keys are passed by Cloudformation
         for cd in container_definitions:
             if "name" not in cd and "Name" not in cd:
                 raise TaskDefinitionMissingPropertyError("name")
             if "image" not in cd and "Image" not in cd:
                 raise TaskDefinitionMissingPropertyError("image")
+            if (
+                requires_compatibilities
+                and "EC2" in requires_compatibilities
+                and ("memory" not in cd and "Memory" not in cd and not memory)
+            ):
+                raise TaskDefinitionMemoryError(cd["name"])
             if (
                 "memory" not in cd
                 and "Memory" not in cd
diff --git a/tests/test_ecs/test_ecs_boto3.py b/tests/test_ecs/test_ecs_boto3.py
index 75a5cebfa..29f3848d1 100644
--- a/tests/test_ecs/test_ecs_boto3.py
+++ b/tests/test_ecs/test_ecs_boto3.py
@@ -324,16 +324,16 @@ def test_register_task_definition_fargate_with_pid_mode():
 
 
 @mock_aws
-def test_register_task_definition_memory_validation():
+def test_register_task_definition_memory_validation_ec2():
     client = boto3.client("ecs", region_name=ECS_REGION)
     container_name = "hello_world"
     bad_definition1 = dict(
         family="test_ecs_task",
         containerDefinitions=[
-            {"name": container_name, "image": "hello-world:latest", "memory": 400},
+            {"name": container_name, "image": "hello-world:latest"},
             {"name": f"{container_name}2", "image": "hello-world:latest"},
         ],
-        requiresCompatibilities=["FARGATE"],
+        requiresCompatibilities=["EC2"],
     )
 
     with pytest.raises(ClientError) as exc:
@@ -344,10 +344,28 @@ def test_register_task_definition_memory_validation():
     assert ex.response["Error"]["Code"] == "ClientException"
     assert (
         ex.response["Error"]["Message"]
-        == f"Invalid setting for container '{container_name}2'. At least one of 'memory' or 'memoryReservation' must be specified."
+        == f"Invalid setting for container '{container_name}'. At least one of 'memory' or 'memoryReservation' must be specified."
     )
 
 
+@mock_aws
+def test_register_task_definition_memory_validation_fargate():
+    client = boto3.client("ecs", region_name=ECS_REGION)
+    container_name = "hello_world"
+    good_definition1 = dict(
+        family="test_ecs_task",
+        memory="1024",
+        containerDefinitions=[
+            {"name": container_name, "image": "hello-world:latest"},
+            {"name": f"{container_name}2", "image": "hello-world:latest"},
+        ],
+        requiresCompatibilities=["FARGATE"],
+    )
+
+    response = client.register_task_definition(**good_definition1)
+    assert response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+
 @mock_aws
 @pytest.mark.parametrize(
     "ecs_def,missing_prop",