chore(toolkit): rollback implementation and tests (#33074)

introduces rollback to the programatic toolkit, plus tests

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: kaizencc

packages/@aws-cdk/toolkit/lib/toolkit/toolkit.ts

@@ -4,6 +4,7 @@ import * as chalk from 'chalk';
 import * as chokidar from 'chokidar';
 import * as fs from 'fs-extra';
 import { ToolkitServices } from './private';
+import { formatErrorMessage } from '../../../../aws-cdk/lib/util/error';
 import { AssetBuildTime, DeployOptions, RequireApproval } from '../actions/deploy';
 import { buildParameterMap, removePublishedAssets } from '../actions/deploy/private';
 import { DestroyOptions } from '../actions/destroy';
@@ -349,7 +350,7 @@ export class Toolkit extends CloudAssemblySourceBuilder implements AsyncDisposab
             stack,
             deployName: stack.stackName,
             roleArn: options.roleArn,
-            toolkitStackName: options.toolkitStackName,
+            toolkitStackName: this.toolkitStackName,
             reuseAssets: options.reuseAssets,
             notificationArns,
             tags,
@@ -616,16 +617,46 @@ export class Toolkit extends CloudAssemblySourceBuilder implements AsyncDisposab
     const ioHost = withAction(this.ioHost, 'rollback');
     const timer = Timer.start();
     const assembly = await this.assemblyFromSource(cx);
-    const stacks = await assembly.selectStacksV2(options.stacks);
+    const stacks = assembly.selectStacksV2(options.stacks);
     await this.validateStacksMetadata(stacks, ioHost);
     const synthTime = timer.end();
     await ioHost.notify(info(`\n✨  Synthesis time: ${synthTime.asSec}s\n`, 'CDK_TOOLKIT_I5001', {
       time: synthTime.asMs,
     }));
 
-    // temporary
-    // eslint-disable-next-line @cdklabs/no-throw-default-error
-    throw new Error('Not implemented yet');
+    if (stacks.stackCount === 0) {
+      await ioHost.notify(error('No stacks selected'));
+      return;
+    }
+
+    let anyRollbackable = false;
+
+    for (const stack of stacks.stackArtifacts) {
+      await ioHost.notify(info(`Rolling back ${chalk.bold(stack.displayName)}`));
+      const startRollbackTime = Timer.start();
+      const deployments = await this.deploymentsForAction('rollback');
+      try {
+        const result = await deployments.rollbackStack({
+          stack,
+          roleArn: options.roleArn,
+          toolkitStackName: this.toolkitStackName,
+          force: options.orphanFailedResources,
+          validateBootstrapStackVersion: options.validateBootstrapStackVersion,
+          orphanLogicalIds: options.orphanLogicalIds,
+        });
+        if (!result.notInRollbackableState) {
+          anyRollbackable = true;
+        }
+        const elapsedRollbackTime = startRollbackTime.end();
+        await ioHost.notify(info(`\n✨  Rollback time: ${elapsedRollbackTime.asSec}s\n`));
+      } catch (e: any) {
+        await ioHost.notify(error(`\n ❌  ${chalk.bold(stack.displayName)} failed: ${formatErrorMessage(e)}`));
+        throw new ToolkitError('Rollback failed (use --force to orphan failing resources)');
+      }
+    }
+    if (!anyRollbackable) {
+      throw new ToolkitError('No stacks were in a state that could be rolled back');
+    }
   }
 
   /**

packages/@aws-cdk/toolkit/test/actions/deploy.test.ts

@@ -172,6 +172,8 @@ describe('deploy', () => {
       await toolkit.deploy(cx);
 
       // THEN
+      // We called rollback
+      expect(toolkit.rollback).toHaveBeenCalledTimes(1);
       successfulDeployment();
     });
 

packages/@aws-cdk/toolkit/test/actions/rollback.test.ts

@@ -0,0 +1,69 @@
+import { StackSelectionStrategy } from '../../lib';
+import { Toolkit } from '../../lib/toolkit';
+import { builderFixture, TestIoHost } from '../_helpers';
+
+const ioHost = new TestIoHost();
+const toolkit = new Toolkit({ ioHost });
+
+let mockRollbackStack = jest.fn().mockResolvedValue({
+  notInRollbackableState: false,
+  success: true,
+});
+
+jest.mock('../../lib/api/aws-cdk', () => {
+  return {
+    ...jest.requireActual('../../lib/api/aws-cdk'),
+    Deployments: jest.fn().mockImplementation(() => ({
+      rollbackStack: mockRollbackStack,
+    })),
+  };
+});
+
+beforeEach(() => {
+  ioHost.notifySpy.mockClear();
+  ioHost.requestSpy.mockClear();
+  jest.clearAllMocks();
+});
+
+describe('rollback', () => {
+  test('successful rollback', async () => {
+    // WHEN
+    const cx = await builderFixture(toolkit, 'two-empty-stacks');
+    await toolkit.rollback(cx, { stacks: { strategy: StackSelectionStrategy.ALL_STACKS } });
+
+    // THEN
+    successfulRollback();
+  });
+
+  test('rollback not in rollbackable state', async () => {
+    // GIVEN
+    mockRollbackStack.mockImplementation(() => ({
+      notInRollbackableState: true,
+      success: false,
+    }));
+    // WHEN
+    const cx = await builderFixture(toolkit, 'two-empty-stacks');
+    await expect(async () => toolkit.rollback(cx, {
+      stacks: { strategy: StackSelectionStrategy.ALL_STACKS },
+    })).rejects.toThrow(/No stacks were in a state that could be rolled back/);
+  });
+
+  test('rollback not in rollbackable state', async () => {
+    // GIVEN
+    mockRollbackStack.mockRejectedValue({});
+
+    // WHEN
+    const cx = await builderFixture(toolkit, 'two-empty-stacks');
+    await expect(async () => toolkit.rollback(cx, {
+      stacks: { strategy: StackSelectionStrategy.ALL_STACKS },
+    })).rejects.toThrow(/Rollback failed/);
+  });
+});
+
+function successfulRollback() {
+  expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
+    action: 'rollback',
+    level: 'info',
+    message: expect.stringContaining('Rollback time:'),
+  }));
+}