From b524b237746c07387c0c16ebbeecf51bb9f3e9b6 Mon Sep 17 00:00:00 2001
From: Paris Holley <mail@parisholley.com>
Date: Wed, 25 Sep 2019 10:40:40 -0400
Subject: [PATCH] support lazy values in fargate task definition

---
 .../lib/fargate/fargate-task-definition.ts    | 17 ++++++++++---
 .../fargate/test.fargate-task-definition.ts   | 24 +++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/packages/@aws-cdk/aws-ecs/lib/fargate/fargate-task-definition.ts b/packages/@aws-cdk/aws-ecs/lib/fargate/fargate-task-definition.ts
index 66d50f1a53128..3c0d4af0de13d 100644
--- a/packages/@aws-cdk/aws-ecs/lib/fargate/fargate-task-definition.ts
+++ b/packages/@aws-cdk/aws-ecs/lib/fargate/fargate-task-definition.ts
@@ -1,4 +1,4 @@
-import { Construct, Resource } from '@aws-cdk/core';
+import { Construct, Lazy, Resource, Token } from '@aws-cdk/core';
 import { CommonTaskDefinitionProps, Compatibility, ITaskDefinition, NetworkMode, TaskDefinition } from '../base/task-definition';
 
 /**
@@ -77,10 +77,21 @@ export class FargateTaskDefinition extends TaskDefinition implements IFargateTas
   constructor(scope: Construct, id: string, props: FargateTaskDefinitionProps = {}) {
     super(scope, id, {
       ...props,
-      cpu: props.cpu !== undefined ? String(props.cpu) : '256',
-      memoryMiB: props.memoryLimitMiB !== undefined ? String(props.memoryLimitMiB) : '512',
+      cpu: props.cpu !== undefined ? stringifyNumber(props.cpu) : '256',
+      memoryMiB: props.memoryLimitMiB !== undefined ? stringifyNumber(props.memoryLimitMiB) : '512',
       compatibility: Compatibility.FARGATE,
       networkMode: NetworkMode.AWS_VPC,
     });
   }
 }
+
+/**
+ * Stringify a number directly or lazily if it's a Token
+ */
+function stringifyNumber(x: number) {
+  if (Token.isUnresolved(x)) {
+    return Lazy.stringValue({ produce: context => `${context.resolve(x)}` });
+  } else {
+    return `${x}`;
+  }
+}
diff --git a/packages/@aws-cdk/aws-ecs/test/fargate/test.fargate-task-definition.ts b/packages/@aws-cdk/aws-ecs/test/fargate/test.fargate-task-definition.ts
index 3548c56864c37..205029dd7bd95 100644
--- a/packages/@aws-cdk/aws-ecs/test/fargate/test.fargate-task-definition.ts
+++ b/packages/@aws-cdk/aws-ecs/test/fargate/test.fargate-task-definition.ts
@@ -23,6 +23,30 @@ export = {
       test.done();
     },
 
+    "support lazy cpu and memory values"(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+      const taskDefinition = new ecs.FargateTaskDefinition(stack, 'FargateTaskDef', {
+        cpu: cdk.Lazy.numberValue({produce: () => 128}),
+        memoryLimitMiB: cdk.Lazy.numberValue({produce: () => 1024})
+      });
+
+      taskDefinition.addVolume({
+        host: {
+          sourcePath: "/tmp/cache",
+        },
+        name: "scratch"
+      });
+
+      // THEN
+      expect(stack).to(haveResourceLike("AWS::ECS::TaskDefinition", {
+        Cpu: "128",
+        Memory: "1024"
+      }));
+
+      test.done();
+    },
+
     "with all properties set"(test: Test) {
       // GIVEN
       const stack = new cdk.Stack();