From 3332d06982c6eb04e4f88c5b01ad745ed18d7e82 Mon Sep 17 00:00:00 2001
From: Rico Huijbers <rix0rrr@gmail.com>
Date: Thu, 13 Feb 2020 12:14:16 +0100
Subject: [PATCH] fix(ec2): `onePerAz` does not work for looked-up VPCs

Fix the use of `onePerAz` for VPCs obtained through `Vpc.fromLookup()`.

Fixes #3126.
---
 packages/@aws-cdk/aws-ec2/lib/vpc.ts          | 12 +++-
 .../aws-ec2/test/test.vpc.from-lookup.ts      | 72 ++++++++++++++++++-
 2 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/packages/@aws-cdk/aws-ec2/lib/vpc.ts b/packages/@aws-cdk/aws-ec2/lib/vpc.ts
index 55f2ec95495f3..78cb95d7cfe0f 100644
--- a/packages/@aws-cdk/aws-ec2/lib/vpc.ts
+++ b/packages/@aws-cdk/aws-ec2/lib/vpc.ts
@@ -416,8 +416,7 @@ abstract class VpcBase extends Resource implements IVpc {
     let subnets = allSubnets[subnetType];
 
     if (onePerAz && subnets.length > 0) {
-      // Restrict to at most one subnet group
-      subnets = subnets.filter(s => subnetGroupNameFromConstructId(s) === subnetGroupNameFromConstructId(subnets[0]));
+      subnets = retainOnePerAz(subnets);
     }
 
     // Force merge conflict here with https://github.com/aws/aws-cdk/pull/4089
@@ -462,6 +461,15 @@ abstract class VpcBase extends Resource implements IVpc {
   }
 }
 
+function retainOnePerAz(subnets: ISubnet[]): ISubnet[] {
+  const azsSeen = new Set<string>();
+  return subnets.filter(subnet => {
+    if (azsSeen.has(subnet.availabilityZone)) { return false; }
+    azsSeen.add(subnet.availabilityZone);
+    return true;
+  });
+}
+
 /**
  * Properties that reference an external Vpc
  */
diff --git a/packages/@aws-cdk/aws-ec2/test/test.vpc.from-lookup.ts b/packages/@aws-cdk/aws-ec2/test/test.vpc.from-lookup.ts
index 9c48031a32a5b..189f037370d25 100644
--- a/packages/@aws-cdk/aws-ec2/test/test.vpc.from-lookup.ts
+++ b/packages/@aws-cdk/aws-ec2/test/test.vpc.from-lookup.ts
@@ -1,7 +1,7 @@
 import { Construct, ContextProvider, GetContextValueOptions, GetContextValueResult, Lazy, Stack } from "@aws-cdk/core";
 import * as cxapi from '@aws-cdk/cx-api';
 import { Test } from 'nodeunit';
-import { Vpc } from "../lib";
+import { SubnetType, Vpc } from "../lib";
 
 export = {
   'Vpc.fromLookup()': {
@@ -102,6 +102,76 @@ export = {
       restoreContextProvider(previous);
       test.done();
     },
+
+    'selectSubnets onePerAz works on imported VPC'(test: Test) {
+      const previous = mockVpcContextProviderWith(test, {
+        vpcId: 'vpc-1234',
+        subnetGroups: [
+          {
+            name: 'Public',
+            type: cxapi.VpcSubnetGroupType.PUBLIC,
+            subnets: [
+              {
+                subnetId: 'pub-sub-in-us-east-1a',
+                availabilityZone: 'us-east-1a',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pub-sub-in-us-east-1b',
+                availabilityZone: 'us-east-1b',
+                routeTableId: 'rt-123',
+              },
+            ],
+          },
+          {
+            name: 'Private',
+            type: cxapi.VpcSubnetGroupType.PRIVATE,
+            subnets: [
+              {
+                subnetId: 'pri-sub-1-in-us-east-1c',
+                availabilityZone: 'us-east-1c',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-2-in-us-east-1c',
+                availabilityZone: 'us-east-1c',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-1-in-us-east-1d',
+                availabilityZone: 'us-east-1d',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-2-in-us-east-1d',
+                availabilityZone: 'us-east-1d',
+                routeTableId: 'rt-123',
+              },
+            ],
+          },
+        ],
+      }, options => {
+        test.deepEqual(options.filter, {
+          isDefault: 'true',
+        });
+
+        test.equal(options.subnetGroupNameTag, undefined);
+      });
+
+      const stack = new Stack();
+      const vpc = Vpc.fromLookup(stack, 'Vpc', {
+        isDefault: true,
+      });
+
+      // WHEN
+      const subnets = vpc.selectSubnets({ subnetType: SubnetType.PRIVATE, onePerAz: true });
+
+      // THEN: we got 2 subnets and not 4
+      test.deepEqual(subnets.subnets.map(s => s.availabilityZone), ['us-east-1c', 'us-east-1d']);
+
+      restoreContextProvider(previous);
+      test.done();
+    },
   },
 };