From 135b5208dce849f49b6f540f7199ece057a7ef22 Mon Sep 17 00:00:00 2001
From: Calvin Combs <66279577+comcalvi@users.noreply.github.com>
Date: Wed, 13 Mar 2024 13:20:12 -0700
Subject: [PATCH] feat(CLI): improved nested stack diff (#29172)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

### Issue # (if applicable)


### Reason for this change

The existing nested stack diff places a fake property, `NestedTemplate`, in the templates to be diffed. This prevents displaying resource replacement information in the diff, like we do for top level stacks. This PR does *not* add changeset replacement information from changesets, but it does add replacement information from the spec.

### Description of changes

Reworked nested stack diff to treat nested stacks as top level stacks. This improves the visual UX and sets us up for using changesets with nested stacks.

#### Before

<img width="957" alt="Screenshot 2024-02-19 at 1 47 59 PM" src="https://github.com/aws/aws-cdk/assets/66279577/a94275c4-e7c3-4d2c-a924-ee61c36bea4d">


#### After
<img width="957" alt="Screenshot 2024-02-19 at 1 48 48 PM" src="https://github.com/aws/aws-cdk/assets/66279577/5263aaf9-ef2f-4228-b413-81e780c4b8f8">



### Description of how you validated changes

Unit tests + manual tests.

### Checklist
- [x] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md)

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 .../cloudformation-diff/lib/diff-template.ts  |   6 +-
 .../cloudformation-diff/lib/format.ts         |   8 +-
 packages/aws-cdk/lib/api/deployments.ts       |  10 +-
 .../api/evaluate-cloudformation-template.ts   |  34 +-
 .../aws-cdk/lib/api/hotswap-deployments.ts    |  22 +-
 .../aws-cdk/lib/api/nested-stack-helpers.ts   |  85 ++--
 packages/aws-cdk/lib/cdk-toolkit.ts           |  16 +-
 packages/aws-cdk/lib/diff.ts                  |  35 +-
 .../api/cloudformation-deployments.test.ts    | 381 ++++++++++--------
 .../api/hotswap/nested-stacks-hotswap.test.ts | 251 ++++++------
 packages/aws-cdk/test/diff.test.ts            | 357 +++++++++++++---
 ...with-two-nested-stacks-stack.template.json |   7 +-
 ...mbda-two-stacks-stack.nested.template.json |   4 +-
 packages/aws-cdk/test/util.ts                 |   4 +-
 14 files changed, 757 insertions(+), 463 deletions(-)

diff --git a/packages/@aws-cdk/cloudformation-diff/lib/diff-template.ts b/packages/@aws-cdk/cloudformation-diff/lib/diff-template.ts
index 1902d757f486d..284bb3a8d5f46 100644
--- a/packages/@aws-cdk/cloudformation-diff/lib/diff-template.ts
+++ b/packages/@aws-cdk/cloudformation-diff/lib/diff-template.ts
@@ -54,7 +54,7 @@ export function fullDiff(
   normalize(newTemplate);
   const theDiff = diffTemplate(currentTemplate, newTemplate);
   if (changeSet) {
-    filterFalsePositivies(theDiff, changeSet);
+    filterFalsePositives(theDiff, changeSet);
     addImportInformation(theDiff, changeSet);
   }
   if (isImport) {
@@ -64,7 +64,7 @@ export function fullDiff(
   return theDiff;
 }
 
-function diffTemplate(
+export function diffTemplate(
   currentTemplate: { [key: string]: any },
   newTemplate: { [key: string]: any },
 ): types.TemplateDiff {
@@ -235,7 +235,7 @@ function addImportInformation(diff: types.TemplateDiff, changeSet?: CloudFormati
   }
 }
 
-function filterFalsePositivies(diff: types.TemplateDiff, changeSet: CloudFormation.DescribeChangeSetOutput) {
+function filterFalsePositives(diff: types.TemplateDiff, changeSet: CloudFormation.DescribeChangeSetOutput) {
   const replacements = findResourceReplacements(changeSet);
   diff.resources.forEachDifference((logicalId: string, change: types.ResourceDifference) => {
     if (change.resourceType.includes('AWS::Serverless')) {
diff --git a/packages/@aws-cdk/cloudformation-diff/lib/format.ts b/packages/@aws-cdk/cloudformation-diff/lib/format.ts
index 7935f774fd468..724af468c2f45 100644
--- a/packages/@aws-cdk/cloudformation-diff/lib/format.ts
+++ b/packages/@aws-cdk/cloudformation-diff/lib/format.ts
@@ -30,7 +30,7 @@ export interface FormatStream extends NodeJS.WritableStream {
 export function formatDifferences(
   stream: FormatStream,
   templateDiff: TemplateDiff,
-  logicalToPathMap: { [logicalId: string]: string } = { },
+  logicalToPathMap: { [logicalId: string]: string } = {},
   context: number = 3) {
   const formatter = new Formatter(stream, logicalToPathMap, templateDiff, context);
 
@@ -59,7 +59,7 @@ export function formatDifferences(
 export function formatSecurityChanges(
   stream: NodeJS.WritableStream,
   templateDiff: TemplateDiff,
-  logicalToPathMap: {[logicalId: string]: string} = {},
+  logicalToPathMap: { [logicalId: string]: string } = {},
   context?: number) {
   const formatter = new Formatter(stream, logicalToPathMap, templateDiff, context);
 
@@ -254,7 +254,7 @@ class Formatter {
           const oldStr = JSON.stringify(oldObject, null, 2);
           const newStr = JSON.stringify(newObject, null, 2);
           const diff = _diffStrings(oldStr, newStr, this.context);
-          for (let i = 0 ; i < diff.length ; i++) {
+          for (let i = 0; i < diff.length; i++) {
             this.print('%s   %s %s', linePrefix, i === 0 ? '└─' : '  ', diff[i]);
           }
         } else {
@@ -466,7 +466,7 @@ function _diffStrings(oldStr: string, newStr: string, context: number): string[]
   function _findIndent(lines: string[]): number {
     let indent = Number.MAX_SAFE_INTEGER;
     for (const line of lines) {
-      for (let i = 1 ; i < line.length ; i++) {
+      for (let i = 1; i < line.length; i++) {
         if (line.charAt(i) !== ' ') {
           indent = indent > i - 1 ? i - 1 : indent;
           break;
diff --git a/packages/aws-cdk/lib/api/deployments.ts b/packages/aws-cdk/lib/api/deployments.ts
index 925cdebd45e15..b7bea90c56c28 100644
--- a/packages/aws-cdk/lib/api/deployments.ts
+++ b/packages/aws-cdk/lib/api/deployments.ts
@@ -7,7 +7,7 @@ import { CredentialsOptions, SdkForEnvironment, SdkProvider } from './aws-auth/s
 import { deployStack, DeployStackResult, destroyStack, DeploymentMethod } from './deploy-stack';
 import { EnvironmentResources, EnvironmentResourcesRegistry } from './environment-resources';
 import { HotswapMode } from './hotswap/common';
-import { loadCurrentTemplateWithNestedStacks, loadCurrentTemplate, flattenNestedStackNames, TemplateWithNestedStackCount } from './nested-stack-helpers';
+import { loadCurrentTemplateWithNestedStacks, loadCurrentTemplate, RootTemplateWithNestedStacks } from './nested-stack-helpers';
 import { CloudFormationStack, Template, ResourcesToImport, ResourceIdentifierSummaries } from './util/cloudformation';
 import { StackActivityProgress } from './util/cloudformation/stack-activity-monitor';
 import { replaceEnvPlaceholders } from './util/placeholders';
@@ -327,13 +327,9 @@ export class Deployments {
   public async readCurrentTemplateWithNestedStacks(
     rootStackArtifact: cxapi.CloudFormationStackArtifact,
     retrieveProcessedTemplate: boolean = false,
-  ): Promise<TemplateWithNestedStackCount> {
+  ): Promise<RootTemplateWithNestedStacks> {
     const sdk = (await this.prepareSdkWithLookupOrDeployRole(rootStackArtifact)).stackSdk;
-    const templateWithNestedStacks = await loadCurrentTemplateWithNestedStacks(rootStackArtifact, sdk, retrieveProcessedTemplate);
-    return {
-      deployedTemplate: templateWithNestedStacks.deployedTemplate,
-      nestedStackCount: flattenNestedStackNames(templateWithNestedStacks.nestedStackNames).length,
-    };
+    return loadCurrentTemplateWithNestedStacks(rootStackArtifact, sdk, retrieveProcessedTemplate);
   }
 
   public async readCurrentTemplate(stackArtifact: cxapi.CloudFormationStackArtifact): Promise<Template> {
diff --git a/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts b/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
index 329bf2ec9ba59..5a1cc675bff85 100644
--- a/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
+++ b/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
@@ -1,7 +1,7 @@
 import * as AWS from 'aws-sdk';
 import { PromiseResult } from 'aws-sdk/lib/request';
 import { ISDK } from './aws-auth';
-import { NestedStackNames } from './nested-stack-helpers';
+import { NestedStackTemplates } from './nested-stack-helpers';
 
 export interface ListStackResources {
   listStackResources(): Promise<AWS.CloudFormation.StackResourceSummary[]>;
@@ -102,7 +102,7 @@ export interface EvaluateCloudFormationTemplateProps {
   readonly partition: string;
   readonly urlSuffix: (region: string) => string;
   readonly sdk: ISDK;
-  readonly nestedStackNames?: { [nestedStackLogicalId: string]: NestedStackNames };
+  readonly nestedStacks?: { [nestedStackLogicalId: string]: NestedStackTemplates };
 }
 
 export class EvaluateCloudFormationTemplate {
@@ -114,7 +114,7 @@ export class EvaluateCloudFormationTemplate {
   private readonly partition: string;
   private readonly urlSuffix: (region: string) => string;
   private readonly sdk: ISDK;
-  private readonly nestedStackNames: { [nestedStackLogicalId: string]: NestedStackNames };
+  private readonly nestedStacks: { [nestedStackLogicalId: string]: NestedStackTemplates };
   private readonly stackResources: ListStackResources;
   private readonly lookupExport: LookupExport;
 
@@ -136,7 +136,7 @@ export class EvaluateCloudFormationTemplate {
     this.sdk = props.sdk;
 
     // We need names of nested stack so we can evaluate cross stack references
-    this.nestedStackNames = props.nestedStackNames ?? {};
+    this.nestedStacks = props.nestedStacks ?? {};
 
     // The current resources of the Stack.
     // We need them to figure out the physical name of a resource in case it wasn't specified by the user.
@@ -163,7 +163,7 @@ export class EvaluateCloudFormationTemplate {
       partition: this.partition,
       urlSuffix: this.urlSuffix,
       sdk: this.sdk,
-      nestedStackNames: this.nestedStackNames,
+      nestedStacks: this.nestedStacks,
     });
   }
 
@@ -386,17 +386,15 @@ export class EvaluateCloudFormationTemplate {
     }
 
     if (foundResource.ResourceType == 'AWS::CloudFormation::Stack' && attribute?.startsWith('Outputs.')) {
-      // need to resolve attributes from another stack's Output section
-      const dependantStackName = this.findNestedStack(logicalId, this.nestedStackNames);
-      if (!dependantStackName) {
+      const dependantStack = this.findNestedStack(logicalId, this.nestedStacks);
+      if (!dependantStack || !dependantStack.physicalName) {
         //this is a newly created nested stack and cannot be hotswapped
         return undefined;
       }
-      const dependantStackTemplate = this.template.Resources[logicalId];
       const evaluateCfnTemplate = await this.createNestedEvaluateCloudFormationTemplate(
-        dependantStackName,
-        dependantStackTemplate?.Properties?.NestedTemplate,
-        dependantStackTemplate.newValue?.Properties?.Parameters);
+        dependantStack.physicalName,
+        dependantStack.generatedTemplate,
+        dependantStack.generatedTemplate.Parameters!);
 
       // Split Outputs.<refName> into 'Outputs' and '<refName>' and recursively call evaluate
       return evaluateCfnTemplate.evaluateCfnExpression({ 'Fn::GetAtt': attribute.split(/\.(.*)/s) });
@@ -406,14 +404,14 @@ export class EvaluateCloudFormationTemplate {
     return this.formatResourceAttribute(foundResource, attribute);
   }
 
-  private findNestedStack(logicalId: string, nestedStackNames: {
-    [nestedStackLogicalId: string]: NestedStackNames;
-  }): string | undefined {
-    for (const [nestedStackLogicalId, { nestedChildStackNames, nestedStackPhysicalName }] of Object.entries(nestedStackNames)) {
+  private findNestedStack(logicalId: string, nestedStacks: {
+    [nestedStackLogicalId: string]: NestedStackTemplates;
+  }): NestedStackTemplates | undefined {
+    for (const nestedStackLogicalId of Object.keys(nestedStacks)) {
       if (nestedStackLogicalId === logicalId) {
-        return nestedStackPhysicalName;
+        return nestedStacks[nestedStackLogicalId];
       }
-      const checkInNestedChildStacks = this.findNestedStack(logicalId, nestedChildStackNames);
+      const checkInNestedChildStacks = this.findNestedStack(logicalId, nestedStacks[nestedStackLogicalId].nestedStackTemplates);
       if (checkInNestedChildStacks) return checkInNestedChildStacks;
     }
     return undefined;
diff --git a/packages/aws-cdk/lib/api/hotswap-deployments.ts b/packages/aws-cdk/lib/api/hotswap-deployments.ts
index 51404d624b92e..efebb414d6f92 100644
--- a/packages/aws-cdk/lib/api/hotswap-deployments.ts
+++ b/packages/aws-cdk/lib/api/hotswap-deployments.ts
@@ -11,7 +11,7 @@ import { isHotswappableEcsServiceChange } from './hotswap/ecs-services';
 import { isHotswappableLambdaFunctionChange } from './hotswap/lambda-functions';
 import { skipChangeForS3DeployCustomResourcePolicy, isHotswappableS3BucketDeploymentChange } from './hotswap/s3-bucket-deployments';
 import { isHotswappableStateMachineChange } from './hotswap/stepfunctions-state-machines';
-import { loadCurrentTemplateWithNestedStacks, NestedStackNames } from './nested-stack-helpers';
+import { NestedStackTemplates, loadCurrentTemplateWithNestedStacks } from './nested-stack-helpers';
 import { CloudFormationStack } from './util/cloudformation';
 import { print } from '../logging';
 
@@ -78,12 +78,12 @@ export async function tryHotswapDeployment(
     partition: (await sdk.currentAccount()).partition,
     urlSuffix: (region) => sdk.getEndpointSuffix(region),
     sdk,
-    nestedStackNames: currentTemplate.nestedStackNames,
+    nestedStacks: currentTemplate.nestedStacks,
   });
 
-  const stackChanges = cfn_diff.fullDiff(currentTemplate.deployedTemplate, stackArtifact.template);
+  const stackChanges = cfn_diff.fullDiff(currentTemplate.deployedRootTemplate, stackArtifact.template);
   const { hotswappableChanges, nonHotswappableChanges } = await classifyResourceChanges(
-    stackChanges, evaluateCfnTemplate, sdk, currentTemplate.nestedStackNames,
+    stackChanges, evaluateCfnTemplate, sdk, currentTemplate.nestedStacks,
   );
 
   logNonHotswappableChanges(nonHotswappableChanges, hotswapMode);
@@ -109,7 +109,7 @@ async function classifyResourceChanges(
   stackChanges: cfn_diff.TemplateDiff,
   evaluateCfnTemplate: EvaluateCloudFormationTemplate,
   sdk: ISDK,
-  nestedStackNames: { [nestedStackName: string]: NestedStackNames },
+  nestedStackNames: { [nestedStackName: string]: NestedStackTemplates },
 ): Promise<ClassifiedResourceChanges> {
   const resourceDifferences = getStackResourceDifferences(stackChanges);
 
@@ -225,12 +225,12 @@ function filterDict<T>(dict: { [key: string]: T }, func: (t: T) => boolean): { [
 async function findNestedHotswappableChanges(
   logicalId: string,
   change: cfn_diff.ResourceDifference,
-  nestedStackNames: { [nestedStackName: string]: NestedStackNames },
+  nestedStackTemplates: { [nestedStackName: string]: NestedStackTemplates },
   evaluateCfnTemplate: EvaluateCloudFormationTemplate,
   sdk: ISDK,
 ): Promise<ClassifiedResourceChanges> {
-  const nestedStackName = nestedStackNames[logicalId].nestedStackPhysicalName;
-  if (!nestedStackName) {
+  const nestedStack = nestedStackTemplates[logicalId];
+  if (!nestedStack.physicalName) {
     return {
       hotswappableChanges: [],
       nonHotswappableChanges: [{
@@ -244,14 +244,14 @@ async function findNestedHotswappableChanges(
   }
 
   const evaluateNestedCfnTemplate = await evaluateCfnTemplate.createNestedEvaluateCloudFormationTemplate(
-    nestedStackName, change.newValue?.Properties?.NestedTemplate, change.newValue?.Properties?.Parameters,
+    nestedStack.physicalName, nestedStack.generatedTemplate, change.newValue?.Properties?.Parameters,
   );
 
   const nestedDiff = cfn_diff.fullDiff(
-    change.oldValue?.Properties?.NestedTemplate, change.newValue?.Properties?.NestedTemplate,
+    nestedStackTemplates[logicalId].deployedTemplate, nestedStackTemplates[logicalId].generatedTemplate,
   );
 
-  return classifyResourceChanges(nestedDiff, evaluateNestedCfnTemplate, sdk, nestedStackNames[logicalId].nestedChildStackNames);
+  return classifyResourceChanges(nestedDiff, evaluateNestedCfnTemplate, sdk, nestedStackTemplates[logicalId].nestedStackTemplates);
 }
 
 /** Returns 'true' if a pair of changes is for the same resource. */
diff --git a/packages/aws-cdk/lib/api/nested-stack-helpers.ts b/packages/aws-cdk/lib/api/nested-stack-helpers.ts
index acbbb1e2a9245..3242794aa1039 100644
--- a/packages/aws-cdk/lib/api/nested-stack-helpers.ts
+++ b/packages/aws-cdk/lib/api/nested-stack-helpers.ts
@@ -5,59 +5,38 @@ import { ISDK } from './aws-auth';
 import { LazyListStackResources, ListStackResources } from './evaluate-cloudformation-template';
 import { CloudFormationStack, Template } from './util/cloudformation';
 
-export interface TemplateWithNestedStackNames {
+export interface NestedStackTemplates {
+  readonly physicalName: string | undefined;
   readonly deployedTemplate: Template;
-  readonly nestedStackNames: { [nestedStackLogicalId: string]: NestedStackNames };
+  readonly generatedTemplate: Template;
+  readonly nestedStackTemplates: { [nestedStackLogicalId: string]: NestedStackTemplates};
 }
 
-export interface NestedStackNames {
-  readonly nestedStackPhysicalName: string | undefined;
-  readonly nestedChildStackNames: { [logicalId: string]: NestedStackNames };
-}
-
-export interface TemplateWithNestedStackCount {
-  readonly deployedTemplate: Template;
-  readonly nestedStackCount: number;
+export interface RootTemplateWithNestedStacks {
+  readonly deployedRootTemplate: Template;
+  readonly nestedStacks: { [nestedStackLogicalId: string]: NestedStackTemplates };
 }
 
 /**
- * Reads the currently deployed template from CloudFormation and adds a
- * property, `NestedTemplate`, to any nested stacks that appear in either
- * the deployed template or the newly synthesized template. `NestedTemplate`
- * is populated with contents of the nested template by mutating the
- * `template` property of `rootStackArtifact`. This is done for all
- * nested stack resources to arbitrary depths.
+ * Reads the currently deployed template and all of its nested stack templates from CloudFormation.
  */
 export async function loadCurrentTemplateWithNestedStacks(
   rootStackArtifact: cxapi.CloudFormationStackArtifact, sdk: ISDK,
   retrieveProcessedTemplate: boolean = false,
-): Promise<TemplateWithNestedStackNames> {
-  const deployedTemplate = await loadCurrentTemplate(rootStackArtifact, sdk, retrieveProcessedTemplate);
-  const nestedStackNames = await addNestedTemplatesToGeneratedAndDeployedStacks(rootStackArtifact, sdk, {
+): Promise<RootTemplateWithNestedStacks> {
+  const deployedRootTemplate = await loadCurrentTemplate(rootStackArtifact, sdk, retrieveProcessedTemplate);
+  const nestedStacks = await loadNestedStacks(rootStackArtifact, sdk, {
     generatedTemplate: rootStackArtifact.template,
-    deployedTemplate: deployedTemplate,
+    deployedTemplate: deployedRootTemplate,
     deployedStackName: rootStackArtifact.stackName,
   });
 
   return {
-    deployedTemplate,
-    nestedStackNames,
+    deployedRootTemplate,
+    nestedStacks,
   };
 }
 
-export function flattenNestedStackNames(nestedStackNames: { [nestedStackLogicalId: string]: NestedStackNames }): string[] {
-  const nameList = [];
-  for (const key of Object.keys(nestedStackNames)) {
-    nameList.push(key);
-
-    if (Object.keys(nestedStackNames[key].nestedChildStackNames).length !== 0) {
-      flattenNestedStacksHelper(nestedStackNames[key].nestedChildStackNames, nameList);
-    }
-  }
-
-  return nameList;
-}
-
 /**
  * Returns the currently deployed template from CloudFormation that corresponds to `stackArtifact`.
  */
@@ -76,13 +55,13 @@ async function loadCurrentStackTemplate(
   return stack.template();
 }
 
-async function addNestedTemplatesToGeneratedAndDeployedStacks(
+async function loadNestedStacks(
   rootStackArtifact: cxapi.CloudFormationStackArtifact,
   sdk: ISDK,
   parentTemplates: StackTemplates,
-): Promise<{ [nestedStackLogicalId: string]: NestedStackNames }> {
+): Promise<{ [nestedStackLogicalId: string]: NestedStackTemplates }> {
   const listStackResources = parentTemplates.deployedStackName ? new LazyListStackResources(sdk, parentTemplates.deployedStackName) : undefined;
-  const nestedStackNames: { [nestedStackLogicalId: string]: NestedStackNames } = {};
+  const nestedStacks: { [nestedStackLogicalId: string]: NestedStackTemplates } = {};
   for (const [nestedStackLogicalId, generatedNestedStackResource] of Object.entries(parentTemplates.generatedTemplate.Resources ?? {})) {
     if (!isCdkManagedNestedStack(generatedNestedStackResource)) {
       continue;
@@ -91,19 +70,11 @@ async function addNestedTemplatesToGeneratedAndDeployedStacks(
     const assetPath = generatedNestedStackResource.Metadata['aws:asset:path'];
     const nestedStackTemplates = await getNestedStackTemplates(rootStackArtifact, assetPath, nestedStackLogicalId, listStackResources, sdk);
 
-    generatedNestedStackResource.Properties.NestedTemplate = nestedStackTemplates.generatedTemplate;
-
-    const deployedParentTemplate = parentTemplates.deployedTemplate;
-    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.deployedTemplate;
-
-    nestedStackNames[nestedStackLogicalId] = {
-      nestedStackPhysicalName: nestedStackTemplates.deployedStackName,
-      nestedChildStackNames: await addNestedTemplatesToGeneratedAndDeployedStacks(
+    nestedStacks[nestedStackLogicalId] = {
+      deployedTemplate: nestedStackTemplates.deployedTemplate,
+      generatedTemplate: nestedStackTemplates.generatedTemplate,
+      physicalName: nestedStackTemplates.deployedStackName,
+      nestedStackTemplates: await loadNestedStacks(
         rootStackArtifact,
         sdk,
         nestedStackTemplates,
@@ -111,7 +82,7 @@ async function addNestedTemplatesToGeneratedAndDeployedStacks(
     };
   }
 
-  return nestedStackNames;
+  return nestedStacks;
 }
 
 async function getNestedStackTemplates(
@@ -153,16 +124,6 @@ function isCdkManagedNestedStack(stackResource: any): stackResource is NestedSta
   return stackResource.Type === 'AWS::CloudFormation::Stack' && stackResource.Metadata && stackResource.Metadata['aws:asset:path'];
 }
 
-function flattenNestedStacksHelper(nestedStackNames: { [logicalId: string]: NestedStackNames }, nameList: string[]) {
-  for (const key of Object.keys(nestedStackNames)) {
-    nameList.push(key);
-
-    if (Object.keys(nestedStackNames[key].nestedChildStackNames).length !== 0) {
-      flattenNestedStacksHelper(nestedStackNames[key].nestedChildStackNames, nameList);
-    }
-  }
-}
-
 interface StackTemplates {
   readonly generatedTemplate: any;
   readonly deployedTemplate: any;
diff --git a/packages/aws-cdk/lib/cdk-toolkit.ts b/packages/aws-cdk/lib/cdk-toolkit.ts
index c67b02dbc9cf2..05aeee002c068 100644
--- a/packages/aws-cdk/lib/cdk-toolkit.ts
+++ b/packages/aws-cdk/lib/cdk-toolkit.ts
@@ -156,7 +156,7 @@ export class CdkToolkit {
       const template = deserializeStructure(await fs.readFile(options.templatePath, { encoding: 'UTF-8' }));
       diffs = options.securityOnly
         ? numberFromBool(printSecurityDiff(template, stacks.firstStack, RequireApproval.Broadening, changeSet))
-        : printStackDiff(template, stacks.firstStack, strict, contextLines, quiet, changeSet, stream);
+        : printStackDiff(template, stacks.firstStack.template, strict, contextLines, quiet, changeSet, stream);
     } else {
       // Compare N stacks against deployed templates
       for (const stack of stacks.stackArtifacts) {
@@ -164,11 +164,11 @@ export class CdkToolkit {
           stream.write(format('Stack %s\n', chalk.bold(stack.displayName)));
         }
 
-        const templateWithNames = await this.props.deployments.readCurrentTemplateWithNestedStacks(
+        const templateWithNestedStacks = await this.props.deployments.readCurrentTemplateWithNestedStacks(
           stack, options.compareAgainstProcessedTemplate,
         );
-        const currentTemplate = templateWithNames.deployedTemplate;
-        const nestedStackCount = templateWithNames.nestedStackCount;
+        const currentTemplate = templateWithNestedStacks.deployedRootTemplate;
+        const nestedStacks = templateWithNestedStacks.nestedStacks;
 
         const resourcesToImport = await this.tryGetResources(await this.props.deployments.resolveEnvironment(stack));
         if (resourcesToImport) {
@@ -203,10 +203,10 @@ export class CdkToolkit {
         // pass a boolean to print if the stack is a migrate stack in order to set all resource diffs to import
         const stackCount =
         options.securityOnly
-          ? (numberFromBool(printSecurityDiff(currentTemplate, stack, RequireApproval.Broadening, changeSet)) > 0 ? 1 : 0)
-          : (printStackDiff(currentTemplate, stack, strict, contextLines, quiet, changeSet, stream, !!resourcesToImport) > 0 ? 1 : 0);
+          ? (numberFromBool(printSecurityDiff(currentTemplate, stack, RequireApproval.Broadening, changeSet)))
+          : (printStackDiff(currentTemplate, stack, strict, contextLines, quiet, changeSet, stream, nestedStacks, !!resourcesToImport));
 
-        diffs += stackCount + nestedStackCount;
+        diffs += stackCount;
       }
     }
 
@@ -1531,4 +1531,4 @@ function buildParameterMap(parameters: {
   }
 
   return parameterMap;
-}
\ No newline at end of file
+}
diff --git a/packages/aws-cdk/lib/diff.ts b/packages/aws-cdk/lib/diff.ts
index fa897430ad37b..b49f288fc2534 100644
--- a/packages/aws-cdk/lib/diff.ts
+++ b/packages/aws-cdk/lib/diff.ts
@@ -1,8 +1,10 @@
+import { format } from 'util';
 import * as cxschema from '@aws-cdk/cloud-assembly-schema';
 import * as cfnDiff from '@aws-cdk/cloudformation-diff';
 import * as cxapi from '@aws-cdk/cx-api';
 import { CloudFormation } from 'aws-sdk';
 import * as chalk from 'chalk';
+import { NestedStackTemplates } from './api/nested-stack-helpers';
 import { print, warning } from './logging';
 
 /**
@@ -14,7 +16,7 @@ import { print, warning } from './logging';
  * @param context     lines of context to use in arbitrary JSON diff
  * @param quiet       silences \'There were no differences\' messages
  *
- * @returns the count of differences that were rendered.
+ * @returns the number of stacks in this stack tree that have differences, including the top-level root stack
  */
 export function printStackDiff(
   oldTemplate: any,
@@ -23,7 +25,8 @@ export function printStackDiff(
   context: number,
   quiet: boolean,
   changeSet?: CloudFormation.DescribeChangeSetOutput,
-  stream?: cfnDiff.FormatStream,
+  stream: cfnDiff.FormatStream = process.stderr,
+  nestedStackTemplates?: { [nestedStackLogicalId: string]: NestedStackTemplates },
   isImport?: boolean): number {
 
   let diff = cfnDiff.fullDiff(oldTemplate, newTemplate.template, changeSet, isImport);
@@ -49,8 +52,10 @@ export function printStackDiff(
     });
   }
 
+  let stackDiffCount = 0;
   if (!diff.isEmpty) {
-    cfnDiff.formatDifferences(stream || process.stderr, diff, {
+    stackDiffCount++;
+    cfnDiff.formatDifferences(stream, diff, {
       ...logicalIdMapFromTemplate(oldTemplate),
       ...buildLogicalToPathMap(newTemplate),
     }, context);
@@ -61,7 +66,29 @@ export function printStackDiff(
     print(chalk.yellow(`Omitted ${filteredChangesCount} changes because they are likely mangled non-ASCII characters. Use --strict to print them.`));
   }
 
-  return diff.differenceCount;
+  for (const nestedStackLogicalId of Object.keys(nestedStackTemplates ?? {})) {
+    if (!nestedStackTemplates) {
+      break;
+    }
+    const nestedStack = nestedStackTemplates[nestedStackLogicalId];
+    if (!quiet) {
+      stream.write(format('Stack %s\n', chalk.bold(nestedStack.physicalName ?? nestedStackLogicalId)));
+    }
+
+    (newTemplate as any)._template = nestedStack.generatedTemplate;
+    stackDiffCount += printStackDiff(
+      nestedStack.deployedTemplate,
+      newTemplate,
+      strict,
+      context,
+      quiet,
+      undefined,
+      stream,
+      nestedStack.nestedStackTemplates,
+    );
+  }
+
+  return stackDiffCount;
 }
 
 export enum RequireApproval {
diff --git a/packages/aws-cdk/test/api/cloudformation-deployments.test.ts b/packages/aws-cdk/test/api/cloudformation-deployments.test.ts
index ac1fed6176e1b..cbaf7c3d8746c 100644
--- a/packages/aws-cdk/test/api/cloudformation-deployments.test.ts
+++ b/packages/aws-cdk/test/api/cloudformation-deployments.test.ts
@@ -246,39 +246,16 @@ test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the temp
   );
 
   // WHEN
-  const nestedStackCount = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStackCount;
-  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedTemplate;
+  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedRootTemplate;
+  const nestedStacks = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStacks;
 
   // THEN
-  expect(nestedStackCount).toEqual(1);
   expect(deployedTemplate).toEqual({
     Resources: {
       NestedStack: {
         Type: 'AWS::CloudFormation::Stack',
         Properties: {
           TemplateURL: 'https://www.magic-url.com',
-          NestedTemplate: {
-            Resources: {
-              NestedResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Property: 'old-value',
-                },
-              },
-            },
-            Outputs: {
-              NestedOutput: {
-                Value: {
-                  Ref: 'NestedResource',
-                },
-              },
-            },
-            Parameters: {
-              NestedParam: {
-                Type: 'String',
-              },
-            },
-          },
         },
         Metadata: {
           'aws:asset:path': 'one-output-one-param-stack.nested.template.json',
@@ -293,33 +270,62 @@ test('readCurrentTemplateWithNestedStacks() can handle non-Resources in the temp
         Type: 'AWS::CloudFormation::Stack',
         Properties: {
           TemplateURL: 'https://www.magic-url.com',
-          NestedTemplate: {
-            Resources: {
-              NestedResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Property: 'new-value',
-                },
-              },
+        },
+        Metadata: {
+          'aws:asset:path': 'one-output-one-param-stack.nested.template.json',
+        },
+      },
+    },
+  });
+
+  expect(nestedStacks).toEqual({
+    NestedStack: {
+      deployedTemplate: {
+        Outputs: {
+          NestedOutput: {
+            Value: {
+              Ref: 'NestedResource',
             },
-            Outputs: {
-              NestedOutput: {
-                Value: {
-                  Ref: 'NestedResource',
-                },
-              },
+          },
+        },
+        Parameters: {
+          NestedParam: {
+            Type: 'String',
+          },
+        },
+        Resources: {
+          NestedResource: {
+            Properties: {
+              Property: 'old-value',
             },
-            Parameters: {
-              NestedParam: {
-                Type: 'Number',
-              },
+            Type: 'AWS::Something',
+          },
+        },
+      },
+      generatedTemplate: {
+        Outputs: {
+          NestedOutput: {
+            Value: {
+              Ref: 'NestedResource',
             },
           },
         },
-        Metadata: {
-          'aws:asset:path': 'one-output-one-param-stack.nested.template.json',
+        Parameters: {
+          NestedParam: {
+            Type: 'Number',
+          },
+        },
+        Resources: {
+          NestedResource: {
+            Properties: {
+              Property: 'new-value',
+            },
+            Type: 'AWS::Something',
+          },
         },
       },
+      nestedStackTemplates: {},
+      physicalName: 'NestedStack',
     },
   });
 });
@@ -453,65 +459,16 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
   );
 
   // WHEN
-  const nestedStackCount = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStackCount;
-  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedTemplate;
+  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedRootTemplate;
+  const nestedStacks = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStacks;
 
   // THEN
-  expect(nestedStackCount).toEqual(3);
   expect(deployedTemplate).toEqual({
     Resources: {
       NestedStack: {
         Type: 'AWS::CloudFormation::Stack',
         Properties: {
           TemplateURL: 'https://www.magic-url.com',
-          NestedTemplate: {
-            Resources: {
-              GrandChildStackA: {
-                Type: 'AWS::CloudFormation::Stack',
-                Properties: {
-                  TemplateURL: 'https://www.magic-url.com',
-                  NestedTemplate: {
-                    Resources: {
-                      SomeResource: {
-                        Type: 'AWS::Something',
-                        Properties: {
-                          Property: 'old-value',
-                        },
-                      },
-                    },
-                  },
-                },
-                Metadata: {
-                  'aws:asset:path': 'one-resource-stack.nested.template.json',
-                },
-              },
-              GrandChildStackB: {
-                Type: 'AWS::CloudFormation::Stack',
-                Properties: {
-                  TemplateURL: 'https://www.magic-url.com',
-                  NestedTemplate: {
-                    Resources: {
-                      SomeResource: {
-                        Type: 'AWS::Something',
-                        Properties: {
-                          Property: 'old-value',
-                        },
-                      },
-                    },
-                  },
-                },
-                Metadata: {
-                  'aws:asset:path': 'one-resource-stack.nested.template.json',
-                },
-              },
-              SomeResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Property: 'old-value',
-                },
-              },
-            },
-          },
         },
         Metadata: {
           'aws:asset:path': 'one-resource-two-stacks-stack.nested.template.json',
@@ -526,59 +483,123 @@ test('readCurrentTemplateWithNestedStacks() with a 3-level nested + sibling stru
         Type: 'AWS::CloudFormation::Stack',
         Properties: {
           TemplateURL: 'https://www.magic-url.com',
-          NestedTemplate: {
+        },
+        Metadata: {
+          'aws:asset:path': 'one-resource-two-stacks-stack.nested.template.json',
+        },
+      },
+    },
+  });
+
+  expect(nestedStacks).toEqual({
+    NestedStack: {
+      deployedTemplate: {
+        Resources: {
+          GrandChildStackA: {
+            Metadata: {
+              'aws:asset:path': 'one-resource-stack.nested.template.json',
+            },
+            Properties: {
+              TemplateURL: 'https://www.magic-url.com',
+            },
+            Type: 'AWS::CloudFormation::Stack',
+          },
+          GrandChildStackB: {
+            Metadata: {
+              'aws:asset:path': 'one-resource-stack.nested.template.json',
+            },
+            Properties: {
+              TemplateURL: 'https://www.magic-url.com',
+            },
+            Type: 'AWS::CloudFormation::Stack',
+          },
+          SomeResource: {
+            Properties: {
+              Property: 'old-value',
+            },
+            Type: 'AWS::Something',
+          },
+        },
+      },
+      generatedTemplate: {
+        Resources: {
+          GrandChildStackA: {
+            Metadata: {
+              'aws:asset:path': 'one-resource-stack.nested.template.json',
+            },
+            Properties: {
+              TemplateURL: 'https://www.magic-url.com',
+            },
+            Type: 'AWS::CloudFormation::Stack',
+          },
+          GrandChildStackB: {
+            Metadata: {
+              'aws:asset:path': 'one-resource-stack.nested.template.json',
+            },
+            Properties: {
+              TemplateURL: 'https://www.magic-url.com',
+            },
+            Type: 'AWS::CloudFormation::Stack',
+          },
+          SomeResource: {
+            Properties: {
+              Property: 'new-value',
+            },
+            Type: 'AWS::Something',
+          },
+        },
+      },
+      nestedStackTemplates: {
+        GrandChildStackA: {
+          deployedTemplate: {
             Resources: {
-              GrandChildStackA: {
-                Type: 'AWS::CloudFormation::Stack',
+              SomeResource: {
                 Properties: {
-                  TemplateURL: 'https://www.magic-url.com',
-                  NestedTemplate: {
-                    Resources: {
-                      SomeResource: {
-                        Type: 'AWS::Something',
-                        Properties: {
-                          Property: 'new-value',
-                        },
-                      },
-                    },
-                  },
-                },
-                Metadata: {
-                  'aws:asset:path': 'one-resource-stack.nested.template.json',
+                  Property: 'old-value',
                 },
+                Type: 'AWS::Something',
               },
-              GrandChildStackB: {
-                Type: 'AWS::CloudFormation::Stack',
+            },
+          },
+          generatedTemplate: {
+            Resources: {
+              SomeResource: {
                 Properties: {
-                  TemplateURL: 'https://www.magic-url.com',
-                  NestedTemplate: {
-                    Resources: {
-                      SomeResource: {
-                        Type: 'AWS::Something',
-                        Properties: {
-                          Property: 'new-value',
-                        },
-                      },
-                    },
-                  },
-                },
-                Metadata: {
-                  'aws:asset:path': 'one-resource-stack.nested.template.json',
+                  Property: 'new-value',
                 },
+                Type: 'AWS::Something',
               },
+            },
+          },
+          nestedStackTemplates: {},
+          physicalName: 'GrandChildStackA',
+        },
+        GrandChildStackB: {
+          deployedTemplate: {
+            Resources: {
               SomeResource: {
+                Properties: {
+                  Property: 'old-value',
+                },
                 Type: 'AWS::Something',
+              },
+            },
+          },
+          generatedTemplate: {
+            Resources: {
+              SomeResource: {
                 Properties: {
                   Property: 'new-value',
                 },
+                Type: 'AWS::Something',
               },
             },
           },
-        },
-        Metadata: {
-          'aws:asset:path': 'one-resource-two-stacks-stack.nested.template.json',
+          nestedStackTemplates: {},
+          physicalName: 'GrandChildStackB',
         },
       },
+      physicalName: 'NestedStack',
     },
   });
 });
@@ -612,26 +633,47 @@ test('readCurrentTemplateWithNestedStacks() on an undeployed parent stack with a
   });
 
   // WHEN
-  const nestedStackCount = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStackCount;
-  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedTemplate;
+  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedRootTemplate;
+  const nestedStacks = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStacks;
 
   // THEN
-  expect(nestedStackCount).toEqual(2);
-  expect(deployedTemplate).toEqual({
-    Resources: {
-      NestedStack: {
-        Type: 'AWS::CloudFormation::Stack',
-        Properties: {
-          NestedTemplate: {
+  expect(deployedTemplate).toEqual({});
+  expect(nestedStacks).toEqual({
+    NestedStack: {
+      deployedTemplate: {},
+      generatedTemplate: {
+        Resources: {
+          SomeResource: {
+            Type: 'AWS::Something',
+            Properties: {
+              Property: 'new-value',
+            },
+          },
+          NestedStack: {
+            Type: 'AWS::CloudFormation::Stack',
+            Properties: {
+              TemplateURL: 'https://www.magic-url.com',
+            },
+            Metadata: {
+              'aws:asset:path': 'one-resource-stack.nested.template.json',
+            },
+          },
+        },
+      },
+      nestedStackTemplates: {
+        NestedStack: {
+          deployedTemplate: {},
+          generatedTemplate: {
             Resources: {
-              NestedStack: {
-                Type: 'AWS::CloudFormation::Stack',
+              SomeResource: {
+                Type: 'AWS::Something',
                 Properties: {
-                  NestedTemplate: {},
+                  Property: 'new-value',
                 },
               },
             },
           },
+          nestedStackTemplates: {},
         },
       },
     },
@@ -787,27 +829,16 @@ test('readCurrentTemplateWithNestedStacks() succesfully ignores stacks without m
   ));
 
   // WHEN
-  const nestedStackCount = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStackCount;
-  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedTemplate;
+  const deployedTemplate = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).deployedRootTemplate;
+  const nestedStacks = (await deployments.readCurrentTemplateWithNestedStacks(rootStack)).nestedStacks;
 
   // THEN
-  expect(nestedStackCount).toEqual(1);
   expect(deployedTemplate).toEqual({
     Resources: {
       WithMetadata: {
         Type: 'AWS::CloudFormation::Stack',
         Properties: {
           TemplateURL: 'https://www.magic-url.com',
-          NestedTemplate: {
-            Resources: {
-              SomeResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Property: 'old-value',
-                },
-              },
-            },
-          },
         },
         Metadata: {
           'aws:asset:path': 'one-resource-stack.nested.template.json',
@@ -835,16 +866,6 @@ test('readCurrentTemplateWithNestedStacks() succesfully ignores stacks without m
         Type: 'AWS::CloudFormation::Stack',
         Properties: {
           TemplateURL: 'https://www.magic-url.com',
-          NestedTemplate: {
-            Resources: {
-              SomeResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Property: 'new-value',
-                },
-              },
-            },
-          },
         },
         Metadata: {
           'aws:asset:path': 'one-resource-stack.nested.template.json',
@@ -852,6 +873,32 @@ test('readCurrentTemplateWithNestedStacks() succesfully ignores stacks without m
       },
     },
   });
+  expect(nestedStacks).toEqual({
+    WithMetadata: {
+      deployedTemplate: {
+        Resources: {
+          SomeResource: {
+            Properties: {
+              Property: 'old-value',
+            },
+            Type: 'AWS::Something',
+          },
+        },
+      },
+      generatedTemplate: {
+        Resources: {
+          SomeResource: {
+            Properties: {
+              Property: 'new-value',
+            },
+            Type: 'AWS::Something',
+          },
+        },
+      },
+      physicalName: 'one-resource-stack',
+      nestedStackTemplates: {},
+    },
+  });
 });
 
 function pushStackResourceSummaries(stackName: string, ...items: CloudFormation.StackResourceSummary[]) {
diff --git a/packages/aws-cdk/test/api/hotswap/nested-stacks-hotswap.test.ts b/packages/aws-cdk/test/api/hotswap/nested-stacks-hotswap.test.ts
index cb499f24bd3b7..318638b72000c 100644
--- a/packages/aws-cdk/test/api/hotswap/nested-stacks-hotswap.test.ts
+++ b/packages/aws-cdk/test/api/hotswap/nested-stacks-hotswap.test.ts
@@ -8,7 +8,6 @@ let mockUpdateLambdaCode: (params: Lambda.Types.UpdateFunctionCodeRequest) => La
 let mockPublishVersion: jest.Mock<Lambda.FunctionConfiguration, Lambda.PublishVersionRequest[]>;
 let hotswapMockSdkProvider: setup.HotswapMockSdkProvider;
 
-// TODO: more tests for parent vs child containing hotswappable changes
 describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hotswapMode) => {
   test('can hotswap a lambda function in a 1-level nested stack', async () => {
     // GIVEN
@@ -18,14 +17,14 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'LambdaRoot',
       template: {
         Resources: {
           NestedStack: {
             Type: 'AWS::CloudFormation::Stack',
             Properties: {
-              TemplateURL: 'https://www.magic-url.com',
+              TemplateURL: 'https://www.amazoff.com',
             },
             Metadata: {
               'aws:asset:path': 'one-lambda-stack.nested.template.json',
@@ -35,7 +34,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'NestedStack',
       template: {
@@ -50,7 +49,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -63,10 +62,10 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    const cdkStackArtifact = testStack({ stackName: 'LambdaRoot', template: rootStack.template });
-
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+    oldRootStack.template.Resources.NestedStack.Properties.TemplateURL = 'https://www.amazon.com';
+    const newRootStack = testStack({ stackName: 'LambdaRoot', template: oldRootStack.template });
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
     // THEN
     expect(deployStackResult).not.toBeUndefined();
@@ -85,14 +84,14 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'TwoLevelLambdaRoot',
       template: {
         Resources: {
           ChildStack: {
             Type: 'AWS::CloudFormation::Stack',
             Properties: {
-              TemplateURL: 'https://www.magic-url.com',
+              TemplateURL: 'https://www.amazoff.com',
             },
             Metadata: {
               'aws:asset:path': 'one-lambda-one-stack-stack.nested.template.json',
@@ -102,8 +101,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
-    setup.addTemplateToCloudFormationLookupMock(testStack({
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
+
+    const oldChildStack = testStack({
       stackName: 'ChildStack',
       template: {
         Resources: {
@@ -117,13 +117,13 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'child-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
           GrandChildStack: {
             Type: 'AWS::CloudFormation::Stack',
             Properties: {
-              TemplateURL: 'https://www.magic-url.com',
+              TemplateURL: 'https://www.amazoff.com',
             },
             Metadata: {
               'aws:asset:path': 'one-lambda-stack.nested.template.json',
@@ -131,7 +131,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
           },
         },
       },
-    }));
+    });
+    setup.addTemplateToCloudFormationLookupMock(oldChildStack);
+
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'GrandChildStack',
       template: {
@@ -146,7 +148,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -164,10 +166,15 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    const cdkStackArtifact = testStack({ stackName: 'TwoLevelLambdaRoot', template: rootStack.template });
-
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+    oldRootStack.template.Resources.ChildStack.Properties.TemplateURL = 'https://www.amazon.com';
+    oldChildStack.template.Resources.GrandChildStack.Properties.TemplateURL = 'https://www.amazon.com';
+
+    // write the new templates to disk
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
+    testStack({ stackName: oldChildStack.stackName, template: oldChildStack.template });
+
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
     // THEN
     expect(deployStackResult).not.toBeUndefined();
@@ -191,7 +198,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'SiblingLambdaRoot',
       template: {
         Resources: {
@@ -214,14 +221,14 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'root-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'NestedStack',
       template: {
@@ -236,7 +243,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -249,11 +256,12 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    rootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
-    const cdkStackArtifact = testStack({ stackName: 'SiblingLambdaRoot', template: rootStack.template });
-
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+    oldRootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
+    oldRootStack.template.Resources.NestedStack.Properties.TemplateURL = 'https://www.amazon.com';
+    // write the updated templates to disk
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
     // THEN
     expect(deployStackResult).not.toBeUndefined();
@@ -279,7 +287,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'NonHotswappableRoot',
       template: {
         Resources: {
@@ -302,14 +310,14 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'root-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'NestedStack',
       template: {
@@ -325,7 +333,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -338,12 +346,13 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    rootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
-    const cdkStackArtifact = testStack({ stackName: 'NonHotswappableRoot', template: rootStack.template });
+    oldRootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
+    oldRootStack.template.Resources.NestedStack.Properties.TemplateURL = 'https://www.amazon.com';
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
 
     if (hotswapMode === HotswapMode.FALL_BACK) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).toBeUndefined();
@@ -351,7 +360,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
 
     } else if (hotswapMode === HotswapMode.HOTSWAP_ONLY) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).not.toBeUndefined();
@@ -373,7 +382,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'NestedStackDeletionRoot',
       template: {
         Resources: {
@@ -396,14 +405,14 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'root-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'NestedStack',
       template: {
@@ -418,7 +427,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -431,20 +440,20 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    rootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
-    delete rootStack.template.Resources.NestedStack;
-    const cdkStackArtifact = testStack({ stackName: 'NestedStackDeletionRoot', template: rootStack.template });
+    oldRootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
+    delete oldRootStack.template.Resources.NestedStack;
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
 
     if (hotswapMode === HotswapMode.FALL_BACK) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).toBeUndefined();
       expect(mockUpdateLambdaCode).not.toHaveBeenCalled();
     } else if (hotswapMode === HotswapMode.HOTSWAP_ONLY) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).not.toBeUndefined();
@@ -466,7 +475,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'NestedStackCreationRoot',
       template: {
         Resources: {
@@ -480,30 +489,31 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'root-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
 
-    rootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
-    rootStack.template.Resources.NestedStack = {
+    oldRootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
+    oldRootStack.template.Resources.NestedStack = {
       Type: 'AWS::CloudFormation::Stack',
       Properties: {
-        TemplateURL: 'https://www.magic-url.com',
+        TemplateURL: 'https://www.amazon.com',
       },
       Metadata: {
         'aws:asset:path': 'one-lambda-stack.nested.template.json',
       },
     };
-    const cdkStackArtifact = testStack({ stackName: 'NestedStackCreationRoot', template: rootStack.template });
+    // we need this because testStack() immediately writes the template to disk, so changing the template afterwards is not going to update the file.
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
 
     if (hotswapMode === HotswapMode.FALL_BACK) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).toBeUndefined();
@@ -511,7 +521,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
 
     } else if (hotswapMode === HotswapMode.HOTSWAP_ONLY) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).not.toBeUndefined();
@@ -533,7 +543,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'NestedStackTypeChangeRoot',
       template: {
         Resources: {
@@ -547,7 +557,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'root-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
           FutureNestedStack: {
@@ -560,30 +570,31 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'spooky-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
 
-    rootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
-    rootStack.template.Resources.FutureNestedStack = {
+    oldRootStack.template.Resources.Func.Properties.Code.S3Bucket = 'new-bucket';
+    oldRootStack.template.Resources.FutureNestedStack = {
       Type: 'AWS::CloudFormation::Stack',
       Properties: {
-        TemplateURL: 'https://www.magic-url.com',
+        TemplateURL: 'https://www.amazon.com',
       },
       Metadata: {
         'aws:asset:path': 'one-lambda-stack.nested.template.json',
       },
     };
-    const cdkStackArtifact = testStack({ stackName: 'NestedStackTypeChangeRoot', template: rootStack.template });
+    // write the updated template to disk
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
 
     if (hotswapMode === HotswapMode.FALL_BACK) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).toBeUndefined();
@@ -591,7 +602,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
 
     } else if (hotswapMode === HotswapMode.HOTSWAP_ONLY) {
       // WHEN
-      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+      const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
       // THEN
       expect(deployStackResult).not.toBeUndefined();
@@ -620,11 +631,11 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
         },
       },
       Metadata: {
-        'aws:asset:path': 'old-path',
+        'aws:asset:path': 'old-lambda-path',
       },
     };
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'MultiLayerRoot',
       template: {
         Resources: {
@@ -642,8 +653,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
-    setup.addTemplateToCloudFormationLookupMock(testStack({
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
+
+    const oldChildStack = testStack({
       stackName: 'ChildStack',
       template: {
         Resources: {
@@ -668,7 +680,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
           Func: lambdaFunctionResource,
         },
       },
-    }));
+    });
+    setup.addTemplateToCloudFormationLookupMock(oldChildStack);
+
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'GrandChildStackA',
       template: {
@@ -708,11 +722,16 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       setup.stackSummaryOf('Func', 'AWS::Lambda::Function', 'grandchild-B-function'),
     );
 
-    rootStack.template.Resources.Func.Properties.Code.S3Key = 'new-key';
-    const cdkStackArtifact = testStack({ stackName: 'MultiLayerRoot', template: rootStack.template });
-
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+    oldRootStack.template.Resources.Func.Properties.Code.S3Key = 'new-key';
+    oldRootStack.template.Resources.ChildStack.Properties.TemplateURL = 'https://www.amazon.com';
+    oldChildStack.template.Resources.GrandChildStackA.Properties.TemplateURL = 'https://www.amazon.com';
+    oldChildStack.template.Resources.GrandChildStackB.Properties.TemplateURL = 'https://www.amazon.com';
+
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
+    //testStack({ stackName: oldChildStack.stackName, template: oldChildStack.template });
+
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack);
 
     // THEN
     expect(deployStackResult).not.toBeUndefined();
@@ -796,7 +815,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -809,10 +828,10 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    const cdkStackArtifact = testStack({ stackName: 'LambdaRoot', template: rootStack.template });
+    rootStack.template.Resources.NestedStack.Properties.TemplateURL = 'https://www.amazon.com';
 
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact, {
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, rootStack, {
       S3BucketParam: 'bucket-param-value',
       S3KeyParam: 'key-param-value',
     });
@@ -826,8 +845,8 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
     });
   });
 
-  test('can hotswap a lambda function in a 2-level nested stack with dependency on a output of 2nd level sibling stack', async () => {
-    // GIVEN: RootStack has one child stack `FirstLevelRootStack` which further has two child stacks
+  test('can hotswap a lambda function in a 2-level nested stack with dependency on an output of 2nd level sibling stack', async () => {
+    // GIVEN: RootStack has one child stack `FirstLevelNestedStack` which further has two child stacks
     // `NestedLambdaStack` and `NestedSiblingStack`. `NestedLambdaStack` takes two parameters s3Key
     // and s3Bucket and use them for a Lambda function.
     // RootStack resolves s3Bucket from a root template parameter and passed to FirstLevelRootStack which
@@ -838,11 +857,11 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       updateFunctionCode: mockUpdateLambdaCode,
     });
 
-    const rootStack = testStack({
+    const oldRootStack = testStack({
       stackName: 'RootStack',
       template: {
         Resources: {
-          FirstLevelRootStack: {
+          FirstLevelNestedStack: {
             Type: 'AWS::CloudFormation::Stack',
             Properties: {
               TemplateURL: 'https://www.magic-url.com',
@@ -866,8 +885,8 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       },
     });
 
-    const firstLevelRootStack = testStack({
-      stackName: 'FirstLevelRootStack',
+    const oldFirstLevelNestedStack = testStack({
+      stackName: 'FirstLevelNestedStack',
       template: {
         Resources: {
           NestedLambdaStack: {
@@ -925,7 +944,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
           },
         },
         Metadata: {
-          'aws:asset:path': 'old-path',
+          'aws:asset:path': 'old-lambda-path',
         },
       },
     });
@@ -936,40 +955,39 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
         Outputs: {
           NestedOutput: { Value: 's3-key-value-from-output' },
         },
-        Metadata: {
-          'aws:asset:path': 'old-path',
-        },
       },
     });
 
-    setup.addTemplateToCloudFormationLookupMock(rootStack);
-    setup.addTemplateToCloudFormationLookupMock(firstLevelRootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldRootStack);
+    setup.addTemplateToCloudFormationLookupMock(oldFirstLevelNestedStack);
     setup.addTemplateToCloudFormationLookupMock(nestedLambdaStack);
     setup.addTemplateToCloudFormationLookupMock(nestedSiblingStack);
 
-    setup.pushNestedStackResourceSummaries('RootStack',
-      setup.stackSummaryOf('FirstLevelRootStack', 'AWS::CloudFormation::Stack',
-        'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/FirstLevelRootStack/abcd',
+    setup.pushNestedStackResourceSummaries(oldRootStack.stackName,
+      setup.stackSummaryOf(oldFirstLevelNestedStack.stackName, 'AWS::CloudFormation::Stack',
+        `arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/${oldFirstLevelNestedStack.stackName}/abcd`,
       ),
     );
-
-    setup.pushNestedStackResourceSummaries('FirstLevelRootStack',
-      setup.stackSummaryOf('NestedLambdaStack', 'AWS::CloudFormation::Stack',
-        'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/NestedLambdaStack/abcd',
+    setup.pushNestedStackResourceSummaries(oldFirstLevelNestedStack.stackName,
+      setup.stackSummaryOf(nestedLambdaStack.stackName, 'AWS::CloudFormation::Stack',
+        `arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/${nestedLambdaStack.stackName}/abcd`,
       ),
-      setup.stackSummaryOf('NestedSiblingStack', 'AWS::CloudFormation::Stack',
-        'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/NestedSiblingStack/abcd',
+      setup.stackSummaryOf(nestedSiblingStack.stackName, 'AWS::CloudFormation::Stack',
+        `arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/${nestedSiblingStack.stackName}/abcd`,
       ),
     );
-    setup.pushNestedStackResourceSummaries('NestedLambdaStack',
+    setup.pushNestedStackResourceSummaries(nestedLambdaStack.stackName,
       setup.stackSummaryOf('Func', 'AWS::Lambda::Function', 'nested-lambda-function'),
     );
-    setup.pushNestedStackResourceSummaries('NestedSiblingStack');
-
-    const cdkStackArtifact = testStack({ stackName: 'RootStack', template: rootStack.template });
+    setup.pushNestedStackResourceSummaries(nestedSiblingStack.stackName);
+    oldRootStack.template.Resources.FirstLevelNestedStack.Properties.TemplateURL = 'https://www.amazon.com';
+    oldFirstLevelNestedStack.template.Resources.NestedLambdaStack.Properties.TemplateURL = 'https://www.amazon.com';
+    oldFirstLevelNestedStack.template.Resources.NestedSiblingStack.Properties.TemplateURL = 'https://www.amazon.com';
+    const newRootStack = testStack({ stackName: oldRootStack.stackName, template: oldRootStack.template });
+    testStack({ stackName: oldFirstLevelNestedStack.stackName, template: oldFirstLevelNestedStack.template });
 
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact, {
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, newRootStack, {
       S3BucketParam: 'new-bucket',
     });
 
@@ -1046,7 +1064,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -1059,10 +1077,10 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    const cdkStackArtifact = testStack({ stackName: 'LambdaRoot', template: rootStack.template });
+    rootStack.template.Resources.NestedStack.Properties.TemplateURL = 'https://www.amazon.com';
 
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, rootStack);
 
     // THEN
     expect(deployStackResult).not.toBeUndefined();
@@ -1129,9 +1147,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
         },
       },
     });
-
     setup.addTemplateToCloudFormationLookupMock(rootStack);
-    setup.addTemplateToCloudFormationLookupMock(testStack({
+
+    const childStack = testStack({
       stackName: 'ChildStack',
       template: {
         Resources: {
@@ -1145,7 +1163,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
           GrandChildStack: {
@@ -1159,7 +1177,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
           },
         },
       },
-    }));
+    });
+    setup.addTemplateToCloudFormationLookupMock(childStack);
+
     setup.addTemplateToCloudFormationLookupMock(testStack({
       stackName: 'GrandChildStack',
       template: {
@@ -1174,7 +1194,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
               FunctionName: 'my-function',
             },
             Metadata: {
-              'aws:asset:path': 'old-path',
+              'aws:asset:path': 'old-lambda-path',
             },
           },
         },
@@ -1192,10 +1212,12 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
         'arn:aws:cloudformation:bermuda-triangle-1337:123456789012:stack/GrandChildStack/abcd',
       ),
     );
-    const cdkStackArtifact = testStack({ stackName: 'LambdaRoot', template: rootStack.template });
+
+    rootStack.template.Resources.ChildStack.Properties.TemplateURL = 'https://www.amazon.com';
+    childStack.template.Resources.GrandChildStack.Properties.TemplateURL = 'https://www.amazon.com';
 
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact, {
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, rootStack, {
       GrandChildS3BucketParam: 'child-bucket-param-value',
       GrandChildS3KeyParam: 'child-key-param-value',
       ChildS3BucketParam: 'bucket-param-value',
@@ -1232,7 +1254,7 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
           NestedStack: {
             Type: 'AWS::CloudFormation::Stack',
             Properties: {
-              TemplateURL: 'https://www.magic-url.com',
+              TemplateURL: 'https://www.amazoff.com',
             },
             Metadata: {
               'aws:asset:path': 'one-lambda-version-stack.nested.template.json',
@@ -1273,10 +1295,9 @@ describe.each([HotswapMode.FALL_BACK, HotswapMode.HOTSWAP_ONLY])('%p mode', (hot
       ),
     );
 
-    const cdkStackArtifact = testStack({ stackName: 'LambdaRoot', template: rootStack.template });
-
     // WHEN
-    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, cdkStackArtifact);
+    rootStack.template.Resources.NestedStack.Properties.TemplateURL = 'https://www.amazon.com';
+    const deployStackResult = await hotswapMockSdkProvider.tryHotswapDeployment(hotswapMode, rootStack);
 
     // THEN
     expect(deployStackResult).not.toBeUndefined();
diff --git a/packages/aws-cdk/test/diff.test.ts b/packages/aws-cdk/test/diff.test.ts
index a7b5905e12f87..8dd11bb0b5fce 100644
--- a/packages/aws-cdk/test/diff.test.ts
+++ b/packages/aws-cdk/test/diff.test.ts
@@ -7,11 +7,119 @@ import { instanceMockFrom, MockCloudExecutable } from './util';
 import { Deployments } from '../lib/api/deployments';
 import { CdkToolkit } from '../lib/cdk-toolkit';
 import * as cfn from '../lib/api/util/cloudformation';
+import { NestedStackTemplates } from '../lib/api/nested-stack-helpers';
+import * as fs from 'fs';
 
 let cloudExecutable: MockCloudExecutable;
 let cloudFormation: jest.Mocked<Deployments>;
 let toolkit: CdkToolkit;
 
+describe('imports', () => {
+  beforeEach(() => {
+    const outputToJson = {
+      '//': 'This file is generated by cdk migrate. It will be automatically deleted after the first successful deployment of this app to the environment of the original resources.',
+      'Source': 'localfile',
+      'Resources': [],
+    };
+    fs.writeFileSync('migrate.json', JSON.stringify(outputToJson, null, 2));
+
+    jest.spyOn(cfn, 'createDiffChangeSet').mockImplementationOnce(async () => {
+      return {
+        Changes: [
+          {
+            ResourceChange: {
+              Action: 'Import',
+              LogicalResourceId: 'Queue',
+            },
+          },
+          {
+            ResourceChange: {
+              Action: 'Import',
+              LogicalResourceId: 'Bucket',
+            },
+          },
+          {
+            ResourceChange: {
+              Action: 'Import',
+              LogicalResourceId: 'Queue2',
+            },
+          },
+        ],
+      };
+    });
+    cloudExecutable = new MockCloudExecutable({
+      stacks: [{
+        stackName: 'A',
+        template: {
+          Resources: {
+            Queue: {
+              Type: 'AWS::SQS::Queue',
+            },
+            Queue2: {
+              Type: 'AWS::SQS::Queue',
+            },
+            Bucket: {
+              Type: 'AWS::S3::Bucket',
+            },
+          },
+        },
+      }],
+    });
+
+    cloudFormation = instanceMockFrom(Deployments);
+
+    toolkit = new CdkToolkit({
+      cloudExecutable,
+      deployments: cloudFormation,
+      configuration: cloudExecutable.configuration,
+      sdkProvider: cloudExecutable.sdkProvider,
+    });
+
+    // Default implementations
+    cloudFormation.readCurrentTemplateWithNestedStacks.mockImplementation((_stackArtifact: CloudFormationStackArtifact) => {
+      return Promise.resolve({
+        deployedRootTemplate: {},
+        nestedStacks: {},
+      });
+    });
+    cloudFormation.deployStack.mockImplementation((options) => Promise.resolve({
+      noOp: true,
+      outputs: {},
+      stackArn: '',
+      stackArtifact: options.stack,
+    }));
+  });
+
+  afterEach(() => {
+    fs.rmSync('migrate.json');
+  });
+
+  test('imports', async () => {
+    // GIVEN
+    const buffer = new StringWritable();
+
+    // WHEN
+    const exitCode = await toolkit.diff({
+      stackNames: ['A'],
+      stream: buffer,
+      changeSet: true,
+    });
+
+    // THEN
+    const plainTextOutput = buffer.data.replace(/\x1B\[[0-?]*[ -/]*[@-~]/g, '');
+    expect(plainTextOutput).toContain(`Stack A
+Parameters and rules created during migration do not affect resource configuration.
+Resources
+[←] AWS::SQS::Queue Queue import
+[←] AWS::SQS::Queue Queue2 import
+[←] AWS::S3::Bucket Bucket import
+`);
+
+    expect(buffer.data.trim()).toContain('✨  Number of stacks with differences: 1');
+    expect(exitCode).toBe(0);
+  });
+});
+
 describe('non-nested stacks', () => {
   beforeEach(() => {
     cloudExecutable = new MockCloudExecutable({
@@ -56,13 +164,13 @@ describe('non-nested stacks', () => {
     cloudFormation.readCurrentTemplateWithNestedStacks.mockImplementation((stackArtifact: CloudFormationStackArtifact) => {
       if (stackArtifact.stackName === 'D') {
         return Promise.resolve({
-          deployedTemplate: { resource: 'D' },
-          nestedStackCount: 0,
+          deployedRootTemplate: { resource: 'D' },
+          nestedStacks: {},
         });
       }
       return Promise.resolve({
-        deployedTemplate: {},
-        nestedStackCount: 0,
+        deployedRootTemplate: {},
+        nestedStacks: {},
       });
     });
     cloudFormation.deployStack.mockImplementation((options) => Promise.resolve({
@@ -223,48 +331,82 @@ describe('nested stacks', () => {
         stackArtifact.template.Resources = {
           AdditionChild: {
             Type: 'AWS::CloudFormation::Stack',
-            Resources: {
-              SomeResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Prop: 'added-value',
-                },
-              },
+            Properties: {
+              TemplateURL: 'addition-child-url-old',
             },
           },
           DeletionChild: {
             Type: 'AWS::CloudFormation::Stack',
-            Resources: {
-              SomeResource: {
-                Type: 'AWS::Something',
-              },
+            Properties: {
+              TemplateURL: 'deletion-child-url-old',
             },
           },
           ChangedChild: {
             Type: 'AWS::CloudFormation::Stack',
-            Resources: {
-              SomeResource: {
-                Type: 'AWS::Something',
-                Properties: {
-                  Prop: 'new-value',
-                },
-              },
+            Properties: {
+              TemplateURL: 'changed-child-url-old',
+            },
+          },
+          UnchangedChild: {
+            Type: 'AWS::CloudFormation::Stack',
+            Properties: {
+              TemplateURL: 'changed-child-url-constant',
             },
           },
         };
         return Promise.resolve({
-          deployedTemplate: {
+          deployedRootTemplate: {
             Resources: {
               AdditionChild: {
                 Type: 'AWS::CloudFormation::Stack',
+                Properties: {
+                  TemplateURL: 'addition-child-url-new',
+                },
+              },
+              DeletionChild: {
+                Type: 'AWS::CloudFormation::Stack',
+                Properties: {
+                  TemplateURL: 'deletion-child-url-new',
+                },
+              },
+              ChangedChild: {
+                Type: 'AWS::CloudFormation::Stack',
+                Properties: {
+                  TemplateURL: 'changed-child-url-new',
+                },
+              },
+              UnchangedChild: {
+                Type: 'AWS::CloudFormation::Stack',
+                Properties: {
+                  TemplateURL: 'changed-child-url-constant',
+                },
+              },
+            },
+          },
+          nestedStacks: {
+            AdditionChild: {
+              deployedTemplate: {
                 Resources: {
                   SomeResource: {
                     Type: 'AWS::Something',
                   },
                 },
               },
-              DeletionChild: {
-                Type: 'AWS::CloudFormation::Stack',
+              generatedTemplate: {
+                Resources: {
+                  SomeResource: {
+                    Type: 'AWS::Something',
+                    Properties: {
+                      Prop: 'added-value',
+                    },
+                  },
+                },
+              },
+              nestedStackTemplates: {},
+              physicalName: 'AdditionChild',
+            },
+            DeletionChild: {
+              deployedTemplate: {
                 Resources: {
                   SomeResource: {
                     Type: 'AWS::Something',
@@ -274,8 +416,18 @@ describe('nested stacks', () => {
                   },
                 },
               },
-              ChangedChild: {
-                Type: 'AWS::CloudFormation::Stack',
+              generatedTemplate: {
+                Resources: {
+                  SomeResource: {
+                    Type: 'AWS::Something',
+                  },
+                },
+              },
+              nestedStackTemplates: {},
+              physicalName: 'DeletionChild',
+            },
+            ChangedChild: {
+              deployedTemplate: {
                 Resources: {
                   SomeResource: {
                     Type: 'AWS::Something',
@@ -285,19 +437,61 @@ describe('nested stacks', () => {
                   },
                 },
               },
+              generatedTemplate: {
+                Resources: {
+                  SomeResource: {
+                    Type: 'AWS::Something',
+                    Properties: {
+                      Prop: 'new-value',
+                    },
+                  },
+                },
+              },
+              nestedStackTemplates: {},
+              physicalName: 'ChangedChild',
+            },
+            newChild: {
+              deployedTemplate: {},
+              generatedTemplate: {
+                Resources: {
+                  SomeResource: {
+                    Type: 'AWS::Something',
+                    Properties: {
+                      Prop: 'new-value',
+                    },
+                  },
+                },
+              },
+              nestedStackTemplates: {
+                newGrandChild: {
+                  deployedTemplate: {},
+                  generatedTemplate: {
+                    Resources: {
+                      SomeResource: {
+                        Type: 'AWS::Something',
+                        Properties: {
+                          Prop: 'new-value',
+                        },
+                      },
+                    },
+                  },
+                  physicalName: undefined,
+                  nestedStackTemplates: {},
+                } as NestedStackTemplates,
+              },
+              physicalName: undefined,
             },
           },
-          nestedStackCount: 3,
         });
       }
       return Promise.resolve({
-        deployedTemplate: {},
-        nestedStackCount: 0,
+        deployedRootTemplate: {},
+        nestedStacks: {},
       });
     });
   });
 
-  test('diff can diff nested stacks', async () => {
+  test('diff can diff nested stacks and display the nested stack logical ID if has not been deployed or otherwise has no physical name', async () => {
     // GIVEN
     const buffer = new StringWritable();
 
@@ -305,6 +499,7 @@ describe('nested stacks', () => {
     const exitCode = await toolkit.diff({
       stackNames: ['Parent'],
       stream: buffer,
+      changeSet: false,
     });
 
     // THEN
@@ -313,23 +508,47 @@ describe('nested stacks', () => {
     expect(plainTextOutput.trim()).toEqual(`Stack Parent
 Resources
 [~] AWS::CloudFormation::Stack AdditionChild
- └─ [~] Resources
-     └─ [~] .SomeResource:
-         └─ [+] Added: .Properties
+ └─ [~] TemplateURL
+     ├─ [-] addition-child-url-new
+     └─ [+] addition-child-url-old
 [~] AWS::CloudFormation::Stack DeletionChild
- └─ [~] Resources
-     └─ [~] .SomeResource:
-         └─ [-] Removed: .Properties
+ └─ [~] TemplateURL
+     ├─ [-] deletion-child-url-new
+     └─ [+] deletion-child-url-old
 [~] AWS::CloudFormation::Stack ChangedChild
- └─ [~] Resources
-     └─ [~] .SomeResource:
-         └─ [~] .Properties:
-             └─ [~] .Prop:
-                 ├─ [-] old-value
-                 └─ [+] new-value
+ └─ [~] TemplateURL
+     ├─ [-] changed-child-url-new
+     └─ [+] changed-child-url-old
+
+Stack AdditionChild
+Resources
+[~] AWS::Something SomeResource
+ └─ [+] Prop
+     └─ added-value
+
+Stack DeletionChild
+Resources
+[~] AWS::Something SomeResource
+ └─ [-] Prop
+     └─ value-to-be-removed
+
+Stack ChangedChild
+Resources
+[~] AWS::Something SomeResource
+ └─ [~] Prop
+     ├─ [-] old-value
+     └─ [+] new-value
 
+Stack newChild
+Resources
+[+] AWS::Something SomeResource
+
+Stack newGrandChild
+Resources
+[+] AWS::Something SomeResource
 
-✨  Number of stacks with differences: 4`);
+
+✨  Number of stacks with differences: 6`);
 
     expect(exitCode).toBe(0);
   });
@@ -352,23 +571,47 @@ Resources
     expect(plainTextOutput.trim()).toEqual(`Stack Parent
 Resources
 [~] AWS::CloudFormation::Stack AdditionChild
- └─ [~] Resources
-     └─ [~] .SomeResource:
-         └─ [+] Added: .Properties
+ └─ [~] TemplateURL
+     ├─ [-] addition-child-url-new
+     └─ [+] addition-child-url-old
 [~] AWS::CloudFormation::Stack DeletionChild
- └─ [~] Resources
-     └─ [~] .SomeResource:
-         └─ [-] Removed: .Properties
+ └─ [~] TemplateURL
+     ├─ [-] deletion-child-url-new
+     └─ [+] deletion-child-url-old
 [~] AWS::CloudFormation::Stack ChangedChild
- └─ [~] Resources
-     └─ [~] .SomeResource:
-         └─ [~] .Properties:
-             └─ [~] .Prop:
-                 ├─ [-] old-value
-                 └─ [+] new-value
+ └─ [~] TemplateURL
+     ├─ [-] changed-child-url-new
+     └─ [+] changed-child-url-old
+
+Stack AdditionChild
+Resources
+[~] AWS::Something SomeResource
+ └─ [+] Prop
+     └─ added-value
+
+Stack DeletionChild
+Resources
+[~] AWS::Something SomeResource
+ └─ [-] Prop
+     └─ value-to-be-removed
+
+Stack ChangedChild
+Resources
+[~] AWS::Something SomeResource
+ └─ [~] Prop
+     ├─ [-] old-value
+     └─ [+] new-value
+
+Stack newChild
+Resources
+[+] AWS::Something SomeResource
+
+Stack newGrandChild
+Resources
+[+] AWS::Something SomeResource
 
 
-✨  Number of stacks with differences: 4`);
+✨  Number of stacks with differences: 6`);
 
     expect(exitCode).toBe(0);
     expect(changeSetSpy).not.toHaveBeenCalled();
diff --git a/packages/aws-cdk/test/nested-stack-templates/one-stack-with-two-nested-stacks-stack.template.json b/packages/aws-cdk/test/nested-stack-templates/one-stack-with-two-nested-stacks-stack.template.json
index bdf9ed5cb2c51..e0bcd1cb3ab43 100644
--- a/packages/aws-cdk/test/nested-stack-templates/one-stack-with-two-nested-stacks-stack.template.json
+++ b/packages/aws-cdk/test/nested-stack-templates/one-stack-with-two-nested-stacks-stack.template.json
@@ -3,7 +3,7 @@
     "NestedLambdaStack": {
       "Type": "AWS::CloudFormation::Stack",
       "Properties": {
-        "TemplateURL": "https://www.magic-url.com",
+        "TemplateURL": "https://www.amazon.com",
         "Parameters": {
           "referenceToS3BucketParam": {
             "Ref": "S3BucketParam"
@@ -23,7 +23,7 @@
     "NestedSiblingStack": {
       "Type": "AWS::CloudFormation::Stack",
       "Properties": {
-        "TemplateURL": "https://www.magic-url.com"
+        "TemplateURL": "https://www.amazon.com"
       },
       "Metadata": {
         "aws:asset:path": "one-output-stack.nested.template.json"
@@ -32,7 +32,8 @@
   },
   "Parameters": {
     "S3BucketParam": {
-      "Type": "String"
+      "Type": "String",
+      "Description": "S3 bucket for asset"
     }
   }
 }
diff --git a/packages/aws-cdk/test/nested-stack-templates/one-unnamed-lambda-two-stacks-stack.nested.template.json b/packages/aws-cdk/test/nested-stack-templates/one-unnamed-lambda-two-stacks-stack.nested.template.json
index 2ac6271955cf4..a1adcd776d469 100644
--- a/packages/aws-cdk/test/nested-stack-templates/one-unnamed-lambda-two-stacks-stack.nested.template.json
+++ b/packages/aws-cdk/test/nested-stack-templates/one-unnamed-lambda-two-stacks-stack.nested.template.json
@@ -15,7 +15,7 @@
     "GrandChildStackA": {
       "Type": "AWS::CloudFormation::Stack",
       "Properties": {
-        "TemplateURL": "https://www.magic-url.com"
+        "TemplateURL": "https://www.amazon.com"
       },
       "Metadata": {
         "aws:asset:path": "one-unnamed-lambda-stack.nested.template.json"
@@ -24,7 +24,7 @@
     "GrandChildStackB": {
       "Type": "AWS::CloudFormation::Stack",
       "Properties": {
-        "TemplateURL": "https://www.magic-url.com"
+        "TemplateURL": "https://www.amazon.com"
       },
       "Metadata": {
         "aws:asset:path": "one-unnamed-lambda-stack.nested.template.json"
diff --git a/packages/aws-cdk/test/util.ts b/packages/aws-cdk/test/util.ts
index ecc262e948153..879d6572f369b 100644
--- a/packages/aws-cdk/test/util.ts
+++ b/packages/aws-cdk/test/util.ts
@@ -37,9 +37,9 @@ export class MockCloudExecutable extends CloudExecutable {
   public readonly configuration: Configuration;
   public readonly sdkProvider: MockSdkProvider;
 
-  constructor(assembly: TestAssembly) {
+  constructor(assembly: TestAssembly, sdkProviderArg?: MockSdkProvider) {
     const configuration = new Configuration();
-    const sdkProvider = new MockSdkProvider();
+    const sdkProvider = sdkProviderArg ?? new MockSdkProvider();
 
     super({
       configuration,