fix(core): allow CfnMapping.findInMap to use pseudo functions/params

packages/@aws-cdk/cdk/lib/cfn-mapping.ts

@@ -1,6 +1,7 @@
 import { CfnRefElement } from './cfn-element';
 import { Construct } from './construct';
 import { Fn } from './fn';
+import { Token } from './token';
 
 export interface CfnMappingProps {
   readonly mapping?: { [k1: string]: { [k2: string]: any } };
@@ -32,11 +33,13 @@ export class CfnMapping extends CfnRefElement {
    * @returns A reference to a value in the map based on the two keys.
    */
   public findInMap(key1: string, key2: string): string {
-    if (!(key1 in this.mapping)) {
+    // opportunistically check that the key exists (if the key does not contain tokens)
+    if (!Token.unresolved(key1) && !(key1 in this.mapping)) {
       throw new Error(`Mapping doesn't contain top-level key '${key1}'`);
     }
 
-    if (!(key2 in this.mapping[key1])) {
+    // opportunistically check that the key exists (if the key does not contain tokens)
+    if (!Token.unresolved(key2) && !(key2 in this.mapping[key1])) {
       throw new Error(`Mapping doesn't contain second-level key '${key2}'`);
     }
 

packages/@aws-cdk/cdk/test/test.mappings.ts

@@ -1,5 +1,5 @@
 import { Test } from 'nodeunit';
-import { CfnMapping, CfnResource, Stack } from '../lib';
+import { CfnMapping, CfnResource, Stack, Aws, Fn } from '../lib';
 
 export = {
   'mappings can be added as another type of entity, and mapping.findInMap can be used to get a token'(test: Test) {
@@ -43,4 +43,24 @@ export = {
 
     test.done();
   },
+
+  'allow using unresolved tokens in find-in-map'(test: Test) {
+    const stack = new Stack();
+
+    const mapping = new CfnMapping(stack, 'mapping', {
+      mapping: {
+        instanceCount: {
+          'us-east-1': 12
+        }
+      }
+    });
+
+    const v1 = mapping.findInMap('instanceCount', Aws.region);
+    const v2 = Fn.findInMap(mapping.logicalId, 'instanceCount', Aws.region);
+
+    const expected = { 'Fn::FindInMap': [ 'mapping', 'instanceCount', { Ref: 'AWS::Region' } ] };
+    test.deepEqual(stack.node.resolve(v1), expected);
+    test.deepEqual(stack.node.resolve(v2), expected);
+    test.done();
+  }
 };