From 82466bea5b1fa4851471c69e1214a4420b73acfa Mon Sep 17 00:00:00 2001
From: Kaizen Conroy <zen.conroy@gmail.com>
Date: Fri, 17 Jan 2025 14:35:56 -0500
Subject: [PATCH] use testiohost

---
 .../@aws-cdk/toolkit/test/_helpers/test-io-host.ts   | 12 ++++++++----
 .../@aws-cdk/toolkit/test/actions/deploy.test.ts     | 11 ++++-------
 packages/@aws-cdk/toolkit/test/actions/synth.test.ts | 11 ++++++-----
 .../test/api/cloud-assembly/source-builder.test.ts   |  3 ++-
 4 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/packages/@aws-cdk/toolkit/test/_helpers/test-io-host.ts b/packages/@aws-cdk/toolkit/test/_helpers/test-io-host.ts
index d8d0dd1641a3d..8b80aeed1602d 100644
--- a/packages/@aws-cdk/toolkit/test/_helpers/test-io-host.ts
+++ b/packages/@aws-cdk/toolkit/test/_helpers/test-io-host.ts
@@ -5,20 +5,24 @@ import { IIoHost, IoMessage, IoMessageLevel, IoRequest, isMessageRelevantForLeve
  * Optional set a level to filter out all irrelevant messages.
  */
 export class TestIoHost implements IIoHost {
-  public readonly spy: jest.Mock<any, any, any>;
+  public readonly notifySpy: jest.Mock<any, any, any>;
+  public readonly requestSpy: jest.Mock<any, any, any>;
 
   constructor(public level: IoMessageLevel = 'info') {
-    this.spy = jest.fn();
+    this.notifySpy = jest.fn();
+    this.requestSpy = jest.fn();
   }
 
   public async notify<T>(msg: IoMessage<T>): Promise<void> {
     if (isMessageRelevantForLevel(msg, this.level)) {
-      this.spy(msg);
+      this.notifySpy(msg);
     }
   }
 
   public async requestResponse<T, U>(msg: IoRequest<T, U>): Promise<U> {
-    await this.notify(msg);
+    if (isMessageRelevantForLevel(msg, this.level)) {
+      this.requestSpy(msg);
+    }
     return msg.defaultResponse;
   }
 }
diff --git a/packages/@aws-cdk/toolkit/test/actions/deploy.test.ts b/packages/@aws-cdk/toolkit/test/actions/deploy.test.ts
index 59d09b427484d..31b2e9e86a9e7 100644
--- a/packages/@aws-cdk/toolkit/test/actions/deploy.test.ts
+++ b/packages/@aws-cdk/toolkit/test/actions/deploy.test.ts
@@ -5,8 +5,6 @@ import { Toolkit } from '../../lib/toolkit';
 import { TestIoHost } from '../_helpers';
 
 const ioHost = new TestIoHost();
-const notifySpy = jest.spyOn(ioHost, 'notify');
-const requestResponseSpy = jest.spyOn(ioHost, 'requestResponse');
 const cdk = new Toolkit({ ioHost });
 
 jest.mock('../../lib/api/aws-cdk', () => {
@@ -38,8 +36,7 @@ const cxFromBuilder = async () => {
 };
 
 beforeEach(() => {
-  requestResponseSpy.mockClear();
-  notifySpy.mockClear();
+  ioHost.notifySpy.mockClear();
   jest.clearAllMocks();
 });
 
@@ -50,7 +47,7 @@ describe('deploy', () => {
     await cdk.deploy(cx);
 
     // THEN
-    expect(notifySpy).toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
       action: 'deploy',
       level: 'info',
       message: expect.stringContaining('Deployment time:'),
@@ -73,7 +70,7 @@ describe('deploy', () => {
     });
 
     // THEN
-    expect(requestResponseSpy).toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.requestSpy).toHaveBeenCalledWith(expect.objectContaining({
       action: 'deploy',
       level: 'info',
       code: 'CDK_TOOLKIT_I5060',
@@ -97,7 +94,7 @@ describe('deploy', () => {
     });
 
     // THEN
-    expect(requestResponseSpy).not.toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.requestSpy).not.toHaveBeenCalledWith(expect.objectContaining({
       action: 'deploy',
       level: 'info',
       code: 'CDK_TOOLKIT_I5060',
diff --git a/packages/@aws-cdk/toolkit/test/actions/synth.test.ts b/packages/@aws-cdk/toolkit/test/actions/synth.test.ts
index 3b64136b1d397..79b1af63c9835 100644
--- a/packages/@aws-cdk/toolkit/test/actions/synth.test.ts
+++ b/packages/@aws-cdk/toolkit/test/actions/synth.test.ts
@@ -5,7 +5,8 @@ const ioHost = new TestIoHost();
 const toolkit = new Toolkit({ ioHost });
 
 beforeEach(() => {
-  ioHost.spy.mockClear();
+  ioHost.notifySpy.mockClear();
+  ioHost.requestSpy.mockClear();
 });
 
 describe('synth', () => {
@@ -15,7 +16,7 @@ describe('synth', () => {
     await toolkit.synth(cx);
 
     // THEN
-    expect(ioHost.spy).toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
       action: 'synth',
       level: 'info',
       message: expect.stringContaining('Successfully synthesized'),
@@ -28,7 +29,7 @@ describe('synth', () => {
     await toolkit.synth(cx);
 
     // THEN
-    expect(ioHost.spy).toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
       action: 'synth',
       level: 'info',
       message: expect.stringContaining('Successfully synthesized'),
@@ -41,7 +42,7 @@ describe('synth', () => {
     await toolkit.synth(cx);
 
     // THEN
-    expect(ioHost.spy).toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
       action: 'synth',
       level: 'info',
       code: 'CDK_TOOLKIT_I0001',
@@ -62,7 +63,7 @@ describe('synth', () => {
     await toolkit.synth(await appFixture(toolkit, 'two-empty-stacks'));
 
     // THEN
-    expect(ioHost.spy).toHaveBeenCalledWith(expect.objectContaining({
+    expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
       action: 'synth',
       level: 'info',
       code: 'CDK_TOOLKIT_I0002',
diff --git a/packages/@aws-cdk/toolkit/test/api/cloud-assembly/source-builder.test.ts b/packages/@aws-cdk/toolkit/test/api/cloud-assembly/source-builder.test.ts
index 9e5ec588963c4..ac3862de03c37 100644
--- a/packages/@aws-cdk/toolkit/test/api/cloud-assembly/source-builder.test.ts
+++ b/packages/@aws-cdk/toolkit/test/api/cloud-assembly/source-builder.test.ts
@@ -5,7 +5,8 @@ const ioHost = new TestIoHost();
 const toolkit = new Toolkit({ ioHost });
 
 beforeEach(() => {
-  ioHost.spy.mockClear();
+  ioHost.notifySpy.mockClear();
+  ioHost.requestSpy.mockClear();
 });
 
 describe('fromAssemblyBuilder', () => {