From 6c74ef8f0a623fc7a5e666bb8ad0f61749b39f21 Mon Sep 17 00:00:00 2001
From: kackyt <t_kaki@nextappli.com>
Date: Thu, 9 May 2024 19:05:04 +0900
Subject: [PATCH] feat(stepfunctions-tasks): add cpu and memory parameters to
 EcsRunTask

---
 .../lib/ecs/run-task.ts                       | 16 ++++++
 .../test/ecs/run-tasks.test.ts                | 54 +++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/ecs/run-task.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/ecs/run-task.ts
index 72c0da31a447b..1bd026a772d62 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/ecs/run-task.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/ecs/run-task.ts
@@ -90,6 +90,20 @@ export interface EcsRunTaskProps extends sfn.TaskStateBaseProps {
    * @default false
    */
   readonly enableExecuteCommand?: boolean;
+
+  /**
+   * Cpu setting override
+   *
+   * @default - No override
+   */
+  readonly cpu?: string;
+
+  /**
+   * Memory setting override
+   *
+   * @default - No override
+   */
+  readonly memoryMiB?: string;
 }
 
 /**
@@ -312,6 +326,8 @@ export class EcsRunTask extends sfn.TaskStateBase implements ec2.IConnectable {
         NetworkConfiguration: this.networkConfiguration,
         Overrides: renderOverrides(this.props.containerOverrides),
         PropagateTags: this.props.propagatedTagSource,
+        Cpu: this.props.cpu,
+        Memory: this.props.memoryMiB,
         ...this.props.launchTarget.bind(this, { taskDefinition: this.props.taskDefinition, cluster: this.props.cluster }).parameters,
         EnableExecuteCommand: this.props.enableExecuteCommand,
       }),
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts
index d7a0ad00510b9..b51a6249db91e 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts
@@ -209,6 +209,60 @@ test('Running a task with NONE as propagated tag source', () => {
     }).toStateJson())).toHaveProperty('Parameters.PropagateTags', 'NONE');
 });
 
+test('Running a task with cpu parameter', () => {
+  const taskDefinition = new ecs.TaskDefinition(stack, 'TD', {
+    memoryMiB: '1024',
+    cpu: '512',
+    compatibility: ecs.Compatibility.FARGATE,
+  });
+  const containerDefinition = taskDefinition.addContainer('TheContainer', {
+    containerName: 'ExplicitContainerName',
+    image: ecs.ContainerImage.fromRegistry('foo/bar'),
+    memoryLimitMiB: 256,
+  });
+
+  expect(stack.resolve(
+    new tasks.EcsRunTask(stack, 'task', {
+      cluster,
+      taskDefinition,
+      cpu: '1024',
+      containerOverrides: [
+        {
+          containerDefinition,
+          environment: [{ name: 'SOME_KEY', value: sfn.JsonPath.stringAt('$.SomeKey') }],
+        },
+      ],
+      launchTarget: new tasks.EcsFargateLaunchTarget(),
+    }).toStateJson())).toHaveProperty('Parameters.Cpu', '1024');
+});
+
+test('Running a task with memory parameter', () => {
+  const taskDefinition = new ecs.TaskDefinition(stack, 'TD', {
+    memoryMiB: '1024',
+    cpu: '512',
+    compatibility: ecs.Compatibility.FARGATE,
+  });
+  const containerDefinition = taskDefinition.addContainer('TheContainer', {
+    containerName: 'ExplicitContainerName',
+    image: ecs.ContainerImage.fromRegistry('foo/bar'),
+    memoryLimitMiB: 256,
+  });
+
+  expect(stack.resolve(
+    new tasks.EcsRunTask(stack, 'task', {
+      cluster,
+      taskDefinition,
+      memoryMiB: '2048',
+      containerOverrides: [
+        {
+          containerDefinition,
+          environment: [{ name: 'SOME_KEY', value: sfn.JsonPath.stringAt('$.SomeKey') }],
+        },
+      ],
+      launchTarget: new tasks.EcsFargateLaunchTarget(),
+    }).toStateJson())).toHaveProperty('Parameters.Memory', '2048');
+});
+
 test('Running a Fargate Task', () => {
   const taskDefinition = new ecs.TaskDefinition(stack, 'TD', {
     memoryMiB: '512',