feat(@aws-cdk/aws-ec2): make GenericMachineImage use CfnMapping in an…

… agnostic stack (aws#8759)
wchaws · Jan 16, 2021 · 464c25e · 464c25e
1 parent ea67c54
commit 464c25e
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 27 deletions.
diff --git a/packages/@aws-cdk/aws-ec2/lib/machine-image.ts b/packages/@aws-cdk/aws-ec2/lib/machine-image.ts
@@ -1,6 +1,6 @@
 import * as ssm from '@aws-cdk/aws-ssm';
 import * as cxschema from '@aws-cdk/cloud-assembly-schema';
-import { Construct, ContextProvider, Stack, Token } from '@aws-cdk/core';
+import { Construct, ContextProvider, CfnMapping, Aws } from '@aws-cdk/core';
 import * as cxapi from '@aws-cdk/cx-api';
 import { UserData } from './user-data';
 import { WindowsVersion } from './windows-versions';
@@ -363,22 +363,17 @@ export interface GenericWindowsImageProps {
  * manually specify an AMI map.
  */
 export class GenericLinuxImage implements IMachineImage {
-  constructor(private readonly amiMap: {[region: string]: string}, private readonly props: GenericLinuxImageProps = {}) {
+  private mapping: { [k1: string]: { [k2: string]: any } } = {};
+  constructor(readonly amiMap: { [region: string]: string }, private readonly props: GenericLinuxImageProps = {}) {
+    for (const [region, ami] of Object.entries(amiMap)) {
+      this.mapping[region] = { ami };
+    }
   }
 
   public getImage(scope: Construct): MachineImageConfig {
-    const region = Stack.of(scope).region;
-    if (Token.isUnresolved(region)) {
-      throw new Error('Unable to determine AMI from AMI map since stack is region-agnostic');
-    }
-
-    const ami = region !== 'test-region' ? this.amiMap[region] : 'ami-12345';
-    if (!ami) {
-      throw new Error(`Unable to find AMI in AMI map: no AMI specified for region '${region}'`);
-    }
-
+    const amiMap = new CfnMapping(scope, 'AmiMap', { mapping: this.mapping });
     return {
-      imageId: ami,
+      imageId: amiMap.findInMap(Aws.REGION, 'ami'),
       userData: this.props.userData ?? UserData.forLinux(),
       osType: OperatingSystemType.LINUX,
     };
@@ -391,22 +386,17 @@ export class GenericLinuxImage implements IMachineImage {
  * Allows you to create a generic Windows EC2 , manually specify an AMI map.
  */
 export class GenericWindowsImage implements IMachineImage {
-  constructor(private readonly amiMap: {[region: string]: string}, private readonly props: GenericWindowsImageProps = {}) {
+  private mapping: { [k1: string]: { [k2: string]: any } } = {};
+  constructor(readonly amiMap: {[region: string]: string}, private readonly props: GenericWindowsImageProps = {}) {
+    for (const [region, ami] of Object.entries(amiMap)) {
+      this.mapping[region] = { ami };
+    }
   }
 
   public getImage(scope: Construct): MachineImageConfig {
-    const region = Stack.of(scope).region;
-    if (Token.isUnresolved(region)) {
-      throw new Error('Unable to determine AMI from AMI map since stack is region-agnostic');
-    }
-
-    const ami = region !== 'test-region' ? this.amiMap[region] : 'ami-12345';
-    if (!ami) {
-      throw new Error(`Unable to find AMI in AMI map: no AMI specified for region '${region}'`);
-    }
-
+    const amiMap = new CfnMapping(scope, 'AmiMap', { mapping: this.mapping });
     return {
-      imageId: ami,
+      imageId: amiMap.findInMap(Aws.REGION, 'ami'),
       userData: this.props.userData ?? UserData.forWindows(),
       osType: OperatingSystemType.WINDOWS,
     };

diff --git a/packages/@aws-cdk/aws-ec2/test/machine-image.test.ts b/packages/@aws-cdk/aws-ec2/test/machine-image.test.ts
@@ -1,3 +1,4 @@
+import { expect as cdkExpect, matchTemplate, MatchStyle } from '@aws-cdk/assert';
 import { App, Stack } from '@aws-cdk/core';
 import * as ec2 from '../lib';
 
@@ -19,7 +20,18 @@ test('can make and use a Windows image', () => {
 
   // THEN
   const details = image.getImage(stack);
-  expect(details.imageId).toEqual('ami-1234');
+  const expected = {
+    Mappings: {
+      AmiMap: {
+        testregion: {
+          ami: 'ami-1234',
+        },
+      },
+    },
+  };
+
+  cdkExpect(stack).to(matchTemplate(expected, MatchStyle.EXACT));
+  expect(stack.resolve(details.imageId)).toEqual({ 'Fn::FindInMap': ['AmiMap', { Ref: 'AWS::Region' }, 'ami'] });
   expect(details.osType).toEqual(ec2.OperatingSystemType.WINDOWS);
 });
 

diff --git a/packages/@aws-cdk/aws-ec2/test/vpc.test.ts b/packages/@aws-cdk/aws-ec2/test/vpc.test.ts
@@ -895,7 +895,7 @@ nodeunitShim({
       // THEN
       expect(stack).to(countResources('AWS::EC2::Instance', 3));
       expect(stack).to(haveResource('AWS::EC2::Instance', {
-        ImageId: 'ami-1',
+        ImageId: { 'Fn::FindInMap': ['TheVPCPublicSubnet1NatInstanceAmiMap7BC63154', { Ref: 'AWS::Region' }, 'ami'] },
         InstanceType: 'q86.mega',
         SourceDestCheck: false,
       }));