From 6f614b769615f7caadc4c3b917ddf3203280399e Mon Sep 17 00:00:00 2001
From: Calvin Combs <comcalvi@amazon.com>
Date: Tue, 1 Feb 2022 12:12:43 -0800
Subject: [PATCH] review comments

---
 .../lib/api/cloudformation-deployments.ts     | 87 +++++++-----------
 packages/aws-cdk/lib/cdk-toolkit.ts           |  2 +-
 .../api/cloudformation-deployments.test.ts    | 89 +++++++++----------
 packages/aws-cdk/test/util.ts                 | 38 ++++----
 4 files changed, 97 insertions(+), 119 deletions(-)

diff --git a/packages/aws-cdk/lib/api/cloudformation-deployments.ts b/packages/aws-cdk/lib/api/cloudformation-deployments.ts
index 4a22cb30836be..3ea84c78c31ff 100644
--- a/packages/aws-cdk/lib/api/cloudformation-deployments.ts
+++ b/packages/aws-cdk/lib/api/cloudformation-deployments.ts
@@ -2,12 +2,12 @@ import * as path from 'path';
 import * as cxapi from '@aws-cdk/cx-api';
 import { AssetManifest } from 'cdk-assets';
 import * as fs from 'fs-extra';
-import { LazyListStackResources, ListStackResources } from '../api/evaluate-cloudformation-template';
 import { Tag } from '../cdk-toolkit';
 import { debug, warning } from '../logging';
 import { publishAssets } from '../util/asset-publishing';
 import { Mode, SdkProvider, ISDK } from './aws-auth';
 import { deployStack, DeployStackResult, destroyStack } from './deploy-stack';
+import { LazyListStackResources, ListStackResources } from './evaluate-cloudformation-template';
 import { ToolkitInfo } from './toolkit-info';
 import { CloudFormationStack, Template } from './util/cloudformation';
 import { StackActivityProgress } from './util/cloudformation/stack-activity-monitor';
@@ -299,7 +299,11 @@ export class CloudFormationDeployments {
   public async readCurrentTemplateWithNestedStacks(rootStackArtifact: cxapi.CloudFormationStackArtifact): Promise<Template> {
     const sdk = await this.prepareSdkWithLookupOrDeployRole(rootStackArtifact);
     const deployedTemplate = await this.readCurrentTemplate(rootStackArtifact, sdk);
-    await this.replaceNestedStacksInRootStack(rootStackArtifact, deployedTemplate, sdk);
+    await this.addNestedTemplatesToGeneratedAndDeployedStacks(rootStackArtifact, sdk, {
+      generatedNestedTemplate: rootStackArtifact.template,
+      deployedNestedTemplate: deployedTemplate,
+      nestedStackName: rootStackArtifact.stackName,
+    });
     return deployedTemplate;
   }
 
@@ -382,75 +386,53 @@ export class CloudFormationDeployments {
   }
 
   private async readCurrentStackTemplate(stackName: string, stackSdk: ISDK) : Promise<Template> {
-    // if `stackName !== stackArtifact.stackName`, then `stackArtifact` is an ancestor to a nested stack with name `stackName`.
     const cfn = stackSdk.cloudFormation();
     const stack = await CloudFormationStack.lookup(cfn, stackName);
     return stack.template();
   }
 
-  private async replaceNestedStacksInRootStack(
-    rootStackArtifact: cxapi.CloudFormationStackArtifact, deployedTemplate: any, sdk: ISDK,
-  ): Promise<void> {
-    await this.replaceNestedStacksInParentTemplate(
-      rootStackArtifact, rootStackArtifact.template, deployedTemplate, rootStackArtifact.stackName, sdk,
-    );
-  }
-
-  private async replaceNestedStacksInParentTemplate(
+  private async addNestedTemplatesToGeneratedAndDeployedStacks(
     rootStackArtifact: cxapi.CloudFormationStackArtifact,
-    generatedParentTemplate: any,
-    deployedParentTemplate: any,
-    parentStackName: string | undefined,
     sdk: ISDK,
+    parentTemplates: NestedStackTemplates,
   ): Promise<void> {
-    const listStackResources = parentStackName ? new LazyListStackResources(sdk, parentStackName) : undefined;
-    for (const [nestedStackLogicalId, nestedStackResource] of Object.entries(generatedParentTemplate.Resources ?? {})) {
-      if (!this.isCdkManagedNestedStack(nestedStackResource)) {
+    const listStackResources = parentTemplates.nestedStackName ? new LazyListStackResources(sdk, parentTemplates.nestedStackName) : undefined;
+    for (const [nestedStackLogicalId, generatedNestedStackResource] of Object.entries(parentTemplates.generatedNestedTemplate.Resources ?? {})) {
+      if (!this.isCdkManagedNestedStack(generatedNestedStackResource)) {
         continue;
       }
 
-      const assetPath = nestedStackResource.Metadata['aws:asset:path'];
-      const nestedStackTemplates = await this.getNestedStackTemplates(
-        rootStackArtifact, assetPath, nestedStackLogicalId, parentStackName, listStackResources, sdk,
-      );
+      const assetPath = generatedNestedStackResource.Metadata['aws:asset:path'];
+      const nestedStackTemplates = await this.getNestedStackTemplates(rootStackArtifact, assetPath, nestedStackLogicalId, listStackResources, sdk);
 
-      //if (!generatedParentTemplate.Resources[nestedStackLogicalId].Properties) {
-      //  generatedParentTemplate.Resources[nestedStackLogicalId].Properties = {};
-      //}
-      generatedParentTemplate.Resources[nestedStackLogicalId].Properties.NestedTemplate = nestedStackTemplates.generatedNestedTemplate;
+      generatedNestedStackResource.Properties.NestedTemplate = nestedStackTemplates.generatedNestedTemplate;
 
-      if (!deployedParentTemplate.Resources) {
-        deployedParentTemplate.Resources = {};
-      }
-      if (!deployedParentTemplate.Resources[nestedStackLogicalId]) {
-        deployedParentTemplate.Resources[nestedStackLogicalId] = {};
-        deployedParentTemplate.Resources[nestedStackLogicalId].Type = 'AWS::CloudFormation::Stack';
-      }
-      if (!deployedParentTemplate.Resources[nestedStackLogicalId].Properties) {
-        deployedParentTemplate.Resources[nestedStackLogicalId].Properties = {};
-      }
-      deployedParentTemplate.Resources[nestedStackLogicalId].Properties.NestedTemplate = nestedStackTemplates.deployedNestedTemplate;
+      const deployedParentTemplate = parentTemplates.deployedNestedTemplate;
+      deployedParentTemplate.Resources = deployedParentTemplate.Resources ?? {};
+      const deployedNestedStackResource = deployedParentTemplate.Resources[nestedStackLogicalId] ?? {};
+      deployedParentTemplate.Resources[nestedStackLogicalId] = deployedNestedStackResource;
+      deployedNestedStackResource.Type = deployedNestedStackResource.Type ?? 'AWS::CloudFormation::Stack';
+      deployedNestedStackResource.Properties = deployedNestedStackResource.Properties ?? {};
+      deployedNestedStackResource.Properties.NestedTemplate = nestedStackTemplates.deployedNestedTemplate;
 
-      await this.replaceNestedStacksInParentTemplate(
+      await this.addNestedTemplatesToGeneratedAndDeployedStacks(
         rootStackArtifact,
-        generatedParentTemplate.Resources[nestedStackLogicalId].Properties.NestedTemplate,
-        deployedParentTemplate.Resources[nestedStackLogicalId].Properties.NestedTemplate ?? {},
-        nestedStackTemplates.nestedStackName,
         sdk,
+        nestedStackTemplates,
       );
     }
   }
 
   private async getNestedStackTemplates(
     rootStackArtifact: cxapi.CloudFormationStackArtifact, nestedTemplateAssetPath: string, nestedStackLogicalId: string,
-    parentStackName: string | undefined, listStackResources: ListStackResources | undefined, sdk: ISDK,
+    listStackResources: ListStackResources | undefined, sdk: ISDK,
   ): Promise<NestedStackTemplates> {
     const nestedTemplatePath = path.join(rootStackArtifact.assembly.directory, nestedTemplateAssetPath);
-    const nestedStackArn = await this.getNestedStackArn(nestedStackLogicalId, parentStackName, listStackResources);
 
     // CFN generates the nested stack name in the form `ParentStackName-NestedStackLogicalID-SomeHashWeCan'tCompute,
     // the arn is of the form: arn:aws:cloudformation:region:123456789012:stack/NestedStackName/AnotherHashWeDon'tNeed
     // so we get the ARN and manually extract the name.
+    const nestedStackArn = await this.getNestedStackArn(nestedStackLogicalId, listStackResources);
     const nestedStackName = nestedStackArn?.slice(nestedStackArn.indexOf('/') + 1, nestedStackArn.lastIndexOf('/'));
 
     return {
@@ -463,21 +445,17 @@ export class CloudFormationDeployments {
   }
 
   private async getNestedStackArn(
-    nestedStackLogicalId: string, parentStackName?: string, listStackResources?: ListStackResources,
+    nestedStackLogicalId: string, listStackResources?: ListStackResources,
   ): Promise<string | undefined> {
     try {
       const stackResources = await listStackResources?.listStackResources();
-
-      if (stackResources) {
-        return stackResources.find(sr => sr.LogicalResourceId === nestedStackLogicalId)?.PhysicalResourceId;
-      }
+      return stackResources?.find(sr => sr.LogicalResourceId === nestedStackLogicalId)?.PhysicalResourceId;
     } catch (e) {
-      if (e.message !== `Stack with id ${parentStackName} does not exist`) {
+      if (!(e.message.startsWith('Stack with id ') && e.message.endsWith(' does not exist'))) {
         throw e;
       }
+      return;
     }
-
-    return undefined;
   }
 
   private isCdkManagedNestedStack(stackResource: any): stackResource is NestedStackResource {
@@ -567,11 +545,12 @@ function isAssetManifestArtifact(art: cxapi.CloudArtifact): art is cxapi.AssetMa
 }
 
 interface NestedStackTemplates {
-  readonly generatedNestedTemplate: any,
-  readonly deployedNestedTemplate: any,
-  readonly nestedStackName: string | undefined,
+  readonly generatedNestedTemplate: any;
+  readonly deployedNestedTemplate: any;
+  readonly nestedStackName: string | undefined;
 }
 
 interface NestedStackResource {
   readonly Metadata: { 'aws:asset:path': string };
+  readonly Properties: any;
 }
diff --git a/packages/aws-cdk/lib/cdk-toolkit.ts b/packages/aws-cdk/lib/cdk-toolkit.ts
index 56c5302544a84..4639689361919 100644
--- a/packages/aws-cdk/lib/cdk-toolkit.ts
+++ b/packages/aws-cdk/lib/cdk-toolkit.ts
@@ -104,7 +104,7 @@ export class CdkToolkit {
       // Compare N stacks against deployed templates
       for (const stack of stacks.stackArtifacts) {
         stream.write(format('Stack %s\n', chalk.bold(stack.displayName)));
-        const currentTemplate = await this.props.cloudFormation.readCurrentTemplateWithNestedStacks(stack);
+        const currentTemplate = await this.props.cloudFormation.readCurrentTemplateWithNestedStacks(stack); // TODO; renmae
         diffs += options.securityOnly
           ? numberFromBool(printSecurityDiff(currentTemplate, stack, RequireApproval.Broadening))
           : printStackDiff(currentTemplate, stack, strict, contextLines, stream);
diff --git a/packages/aws-cdk/test/api/cloudformation-deployments.test.ts b/packages/aws-cdk/test/api/cloudformation-deployments.test.ts
index 1d23e120a5ece..fc3ceee0ac4c3 100644
--- a/packages/aws-cdk/test/api/cloudformation-deployments.test.ts
+++ b/packages/aws-cdk/test/api/cloudformation-deployments.test.ts
@@ -5,8 +5,9 @@ import { CloudFormationDeployments } from '../../lib/api/cloudformation-deployme
 import { deployStack } from '../../lib/api/deploy-stack';
 import { ToolkitInfo } from '../../lib/api/toolkit-info';
 import { CloudFormationStack } from '../../lib/api/util/cloudformation';
-import { instanceMockFrom, testStack } from '../util';
+import { testStack } from '../util';
 import { mockBootstrapStack, MockSdkProvider } from '../util/mock-sdk';
+import { FakeCloudformationStack } from './fake-cloudformation-stack';
 
 let sdkProvider: MockSdkProvider;
 let deployments: CloudFormationDeployments;
@@ -23,11 +24,12 @@ beforeEach(() => {
   sdkProvider.stubCloudFormation({
     listStackResources: ({ StackName: stackName }) => {
       numberOfTimesListStackResourcesWasCalled++;
-      if (!currentCfnStackResources[stackName]) {
+      const stackResources = currentCfnStackResources[stackName];
+      if (!stackResources) {
         throw new Error(`Stack with id ${stackName} does not exist`);
       }
       return {
-        StackResourceSummaries: currentCfnStackResources[stackName],
+        StackResourceSummaries: stackResources,
       };
     },
   });
@@ -163,11 +165,14 @@ test('if toolkit stack cannot be found but SSM parameter name is present deploym
 });
 
 test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the template', async () => {
-  const cfnStack = instanceMockFrom((CloudFormationStack as any));
-  CloudFormationStack.lookup = jest.fn().mockImplementation((_, stackName: string) => {
+  const cfnStack = new FakeCloudformationStack({
+    stackName: 'ParentStackWithOutputAndParameter',
+    stackId: 'StackId',
+  });
+  CloudFormationStack.lookup = (async (_, stackName: string) => {
     switch (stackName) {
-      case 'OutputParamRoot':
-        (cfnStack as any).template = jest.fn().mockReturnValue({
+      case 'ParentStackWithOutputAndParameter':
+        cfnStack.template = async () => ({
           Resources: {
             NestedStack: {
               Type: 'AWS::CloudFormation::Stack',
@@ -183,7 +188,7 @@ test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the temp
         break;
 
       case 'NestedStack':
-        (cfnStack as any).template = jest.fn().mockReturnValue({
+        cfnStack.template = async () => ({
           Resources: {
             NestedResource: {
               Type: 'AWS::Something',
@@ -208,14 +213,14 @@ test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the temp
         break;
 
       default:
-        throw new Error('unknown stack name ' + stackName + ' found in cloudformation-deployments.test.ts');
+        throw new Error('unknown stack name ' + stackName + ' found');
     }
 
     return cfnStack;
   });
 
   const rootStack = testStack({
-    stackName: 'OutputParamRoot',
+    stackName: 'ParentStackWithOutputAndParameter',
     template: {
       Resources: {
         NestedStack: {
@@ -231,16 +236,11 @@ test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the temp
     },
   });
 
-  pushStackResourceSummaries('OutputParamRoot',
+  pushStackResourceSummaries('ParentStackWithOutputAndParameter',
     stackSummaryOf('NestedStack', 'AWS::CloudFormation::Stack',
       'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/NestedStack/abcd',
     ),
   );
-  pushStackResourceSummaries('NestedStack',
-    stackSummaryOf('NestedResource', 'AWS::Something',
-      'arn:aws:something:bermuda-triangle-1337:123456789012:property',
-    ),
-  );
 
   // WHEN
   const deployedTemplate = await deployments.readCurrentTemplateWithNestedStacks(rootStack);
@@ -320,11 +320,14 @@ test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the temp
 });
 
 test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling structure works', async () => {
-  const cfnStack = instanceMockFrom((CloudFormationStack as any));
-  CloudFormationStack.lookup = jest.fn().mockImplementation((_, stackName: string) => {
+  const cfnStack = new FakeCloudformationStack({
+    stackName: 'MultiLevelRoot',
+    stackId: 'StackId',
+  });
+  CloudFormationStack.lookup = (async (_, stackName: string) => {
     switch (stackName) {
       case 'MultiLevelRoot':
-        (cfnStack as any).template = jest.fn().mockReturnValue({
+        cfnStack.template = async () => ({
           Resources: {
             NestedStack: {
               Type: 'AWS::CloudFormation::Stack',
@@ -340,7 +343,7 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
         break;
 
       case 'NestedStack':
-        (cfnStack as any).template = jest.fn().mockReturnValue({
+        cfnStack.template = async () => ({
           Resources: {
             SomeResource: {
               Type: 'AWS::Something',
@@ -371,7 +374,7 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
         break;
 
       case 'NestedStackA':
-        (cfnStack as any).template = jest.fn().mockReturnValue({
+        cfnStack.template = async () => ({
           Resources: {
             SomeResource: {
               Type: 'AWS::Something',
@@ -384,7 +387,7 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
         break;
 
       case 'NestedStackB':
-        (cfnStack as any).template = jest.fn().mockReturnValue({
+        cfnStack.template = async () => ({
           Resources: {
             SomeResource: {
               Type: 'AWS::Something',
@@ -432,9 +435,6 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
     stackSummaryOf('NestedStackB', 'AWS::CloudFormation::Stack',
       'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/NestedStackB/abcd',
     ),
-    stackSummaryOf('SomeResource', 'AWS::Something',
-      'arn:aws:something:bermuda-triangle-1337:123456789012:property',
-    ),
   );
   pushStackResourceSummaries('NestedStackA',
     stackSummaryOf('NestedStack', 'AWS::CloudFormation::Stack',
@@ -446,16 +446,6 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
       'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/GrandChildB/abcd',
     ),
   );
-  pushStackResourceSummaries('GrandChildA',
-    stackSummaryOf('SomeResource', 'AWS::Something',
-      'arn:aws:something:bermuda-triangle-1337:123456789012:property',
-    ),
-  );
-  pushStackResourceSummaries('GrandChildB',
-    stackSummaryOf('SomeResource', 'AWS::Something',
-      'arn:aws:something:bermuda-triangle-1337:123456789012:property',
-    ),
-  );
 
   // WHEN
   const deployedTemplate = await deployments.readCurrentTemplateWithNestedStacks(rootStack);
@@ -588,9 +578,12 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
 
 test('readCurrentTemplateWithNestedStacks() on an undeployed parent stack with an (also undeployed) nested stack works', async () => {
   // GIVEN
-  const cfnStack = instanceMockFrom((CloudFormationStack as any));
-  CloudFormationStack.lookup = jest.fn().mockImplementation((_cfn: CloudFormation, _stackName: string) => {
-    (cfnStack as any).template = jest.fn().mockReturnValue({ });
+  const cfnStack = new FakeCloudformationStack({
+    stackName: 'UndeployedParent',
+    stackId: 'StackId',
+  });
+  CloudFormationStack.lookup = (async (_cfn, _stackName: string) => {
+    cfnStack.template = async () => ({});
 
     return cfnStack;
   });
@@ -638,9 +631,12 @@ test('readCurrentTemplateWithNestedStacks() on an undeployed parent stack with a
 
 test('readCurrentTemplateWithNestedStacks() caches calls to listStackResources()', async () => {
   // GIVEN
-  const cfnStack = instanceMockFrom((CloudFormationStack as any));
-  CloudFormationStack.lookup = jest.fn().mockImplementation((_cfn: CloudFormation, _stackName: string) => {
-    (cfnStack as any).template = jest.fn().mockReturnValue({
+  const cfnStack = new FakeCloudformationStack({
+    stackName: 'CachingRoot',
+    stackId: 'StackId',
+  });
+  CloudFormationStack.lookup = (async (_cfn, _stackName: string) => {
+    cfnStack.template = async () => ({
       Resources:
       {
         NestedStackA: {
@@ -711,10 +707,13 @@ test('readCurrentTemplateWithNestedStacks() caches calls to listStackResources()
 
 test('readCurrentTemplateWithNestedStacks() succesfully ignores stacks without metadata', async () => {
   // GIVEN
-  const cfnStack = instanceMockFrom((CloudFormationStack as any));
-  CloudFormationStack.lookup = jest.fn().mockImplementation((_cfn: CloudFormation, stackName: string) => {
+  const cfnStack = new FakeCloudformationStack({
+    stackName: 'MetadataRoot',
+    stackId: 'StackId',
+  });
+  CloudFormationStack.lookup = (async (_, stackName: string) => {
     if (stackName === 'MetadataRoot') {
-      (cfnStack as any).template = jest.fn().mockReturnValue({
+      cfnStack.template = async () => ({
         Resources: {
           WithMetadata: {
             Type: 'AWS::CloudFormation::Stack',
@@ -729,7 +728,7 @@ test('readCurrentTemplateWithNestedStacks() succesfully ignores stacks without m
       });
 
     } else {
-      (cfnStack as any).template = jest.fn().mockReturnValue({
+      cfnStack.template = async () => ({
         Resources: {
           SomeResource: {
             Type: 'AWS::Something',
diff --git a/packages/aws-cdk/test/util.ts b/packages/aws-cdk/test/util.ts
index 63781e9b882d9..fb714bc8427d7 100644
--- a/packages/aws-cdk/test/util.ts
+++ b/packages/aws-cdk/test/util.ts
@@ -49,25 +49,6 @@ function clone(obj: any) {
   return JSON.parse(JSON.stringify(obj));
 }
 
-function addNestedStacks(templatePath: string, outdir: string, rootStackTemplate?: any) {
-  let template = rootStackTemplate;
-
-  if (!template) {
-    const templatePathWithDir = path.join('diff-nested-stacks-templates', templatePath);
-    template = JSON.parse(fs.readFileSync(path.join(__dirname, templatePathWithDir)).toString());
-    fs.writeFileSync(path.join(outdir, templatePath), JSON.stringify(template, undefined, 2));
-  }
-
-  for (const logicalId in template.Resources) {
-    if (template.Resources[logicalId].Type === 'AWS::CloudFormation::Stack') {
-      if (template.Resources[logicalId].Metadata && template.Resources[logicalId].Metadata['aws:asset:path']) {
-        const nestedTemplatePath = template.Resources[logicalId].Metadata['aws:asset:path'];
-        addNestedStacks(nestedTemplatePath, outdir);
-      }
-    }
-  }
-}
-
 function addAttributes(assembly: TestAssembly, builder: cxapi.CloudAssemblyBuilder) {
   for (const stack of assembly.stacks) {
     const templateFile = `${stack.stackName}.template.json`;
@@ -104,6 +85,25 @@ function addAttributes(assembly: TestAssembly, builder: cxapi.CloudAssemblyBuild
   }
 }
 
+function addNestedStacks(templatePath: string, outdir: string, rootStackTemplate?: any) {
+  let template = rootStackTemplate;
+
+  if (!template) {
+    const templatePathWithDir = path.join('diff-nested-stacks-templates', templatePath);
+    template = JSON.parse(fs.readFileSync(path.join(__dirname, templatePathWithDir)).toString());
+    fs.writeFileSync(path.join(outdir, templatePath), JSON.stringify(template, undefined, 2));
+  }
+
+  for (const logicalId in template.Resources) {
+    if (template.Resources[logicalId].Type === 'AWS::CloudFormation::Stack') {
+      if (template.Resources[logicalId].Metadata && template.Resources[logicalId].Metadata['aws:asset:path']) {
+        const nestedTemplatePath = template.Resources[logicalId].Metadata['aws:asset:path'];
+        addNestedStacks(nestedTemplatePath, outdir);
+      }
+    }
+  }
+}
+
 export function testAssembly(assembly: TestAssembly): cxapi.CloudAssembly {
   const builder = new cxapi.CloudAssemblyBuilder();
   addAttributes(assembly, builder);