diff --git a/fbpcp/entity/container_instance.py b/fbpcp/entity/container_instance.py index 9e34a6d1..ea0373ed 100644 --- a/fbpcp/entity/container_instance.py +++ b/fbpcp/entity/container_instance.py @@ -26,3 +26,5 @@ class ContainerInstance: instance_id: str ip_address: Optional[str] = None status: ContainerInstanceStatus = ContainerInstanceStatus.UNKNOWN + cpu: Optional[int] = None # Number of vCPU + memory: Optional[int] = None # Memory in GB diff --git a/fbpcp/mapper/aws.py b/fbpcp/mapper/aws.py index 25246d88..18f10040 100644 --- a/fbpcp/mapper/aws.py +++ b/fbpcp/mapper/aws.py @@ -70,8 +70,20 @@ def map_ecstask_to_containerinstance(task: Dict[str, Any]) -> ContainerInstance: status = ContainerInstanceStatus.FAILED else: status = ContainerInstanceStatus.UNKNOWN - - return ContainerInstance(task["taskArn"], ip_v4, status) + task_overrides = None + if ( + "overrides" in task + and "cpu" in task["overrides"] + and "memory" in task["overrides"] + ): + task_overrides = task["overrides"] + elif "cpu" in task and "memory" in task: + task_overrides = task + vcpu = map_unit_to_vcpu(int(task_overrides["cpu"])) if task_overrides else None + memory_in_gb = ( + map_mb_to_gb(int(task_overrides["memory"])) if task_overrides else None + ) + return ContainerInstance(task["taskArn"], ip_v4, status, vcpu, memory_in_gb) def map_esccluster_to_clusterinstance(cluster: Dict[str, Any]) -> Cluster: diff --git a/tests/gateway/test_ecs.py b/tests/gateway/test_ecs.py index b4532143..0ef801d3 100644 --- a/tests/gateway/test_ecs.py +++ b/tests/gateway/test_ecs.py @@ -96,6 +96,8 @@ def test_run_task(self) -> None: self.TEST_TASK_ARN, self.TEST_IP_ADDRESS, ContainerInstanceStatus.STARTED, + self.TEST_CPU, + self.TEST_MEMORY, ) # Act task = self.gw.run_task( diff --git a/tests/mapper/test_aws.py b/tests/mapper/test_aws.py index af94d4e8..c9b75d9c 100644 --- a/tests/mapper/test_aws.py +++ b/tests/mapper/test_aws.py @@ -107,11 +107,15 @@ def test_map_ecstask_to_containerinstance(self): self.TEST_TASK_ARN, None, ContainerInstanceStatus.COMPLETED, + self.TEST_CPU, + self.TEST_MEMORY, ), ContainerInstance( self.TEST_TASK_ARN, None, ContainerInstanceStatus.FAILED, + self.TEST_CPU, + self.TEST_MEMORY, ), ContainerInstance( self.TEST_TASK_ARN,