fix(toolkit-lib): deployment total time includes stray output (#238)

mrgrain · web-flow · commit ec93ebd39180 · 2025-03-18T11:41:29.000Z
Fixes an extra `Deployment &lt;number&gt;` at the end of the "Total Time"
message reporting the total deployment time.

Adds missing test cases for the `MessageSpan` feature (GenAI supported).

---
By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache-2.0 license
diff --git a/packages/@aws-cdk/tmp-toolkit-helpers/src/api/io/private/span.ts b/packages/@aws-cdk/tmp-toolkit-helpers/src/api/io/private/span.ts
@@ -106,23 +106,24 @@ export class SpanMaker<S extends object, E extends SpanEnd> {
    */
   public async begin(payload: VoidWhenEmpty<S>): Promise<IMessageSpan<E>>;
   public async begin(message: string, payload: S): Promise<IMessageSpan<E>>;
-  public async begin(first: any, second?: S): Promise<IMessageSpan<E>> {
+  public async begin(a: any, b?: S): Promise<IMessageSpan<E>> {
     const spanId = uuid.v4();
     const startTime = new Date().getTime();
 
     const notify = (msg: ActionLessMessage<unknown>): Promise<void> => {
       return this.ioHelper.notify(withSpanId(spanId, msg));
     };
 
-    const startMsg = second ? first : 'Starting %s ...';
-    const startPayload = second ?? first;
+    const startInput = parseArgs<S>(a, b);
+    const startMsg = startInput.message ?? `Starting ${this.definition.name} ...`;
+    const startPayload = startInput.payload;
 
     await notify(this.definition.start.msg(
-      util.format(startMsg, this.definition.name),
+      startMsg,
       startPayload,
     ));
 
-    const timingMsg = '\n✨  %s time: %ds\n';
+    const timingMsgTemplate = '\n✨  %s time: %ds\n';
     const time = () => {
       const elapsedTime = new Date().getTime() - startTime;
       return {
@@ -132,13 +133,7 @@ export class SpanMaker<S extends object, E extends SpanEnd> {
     };
 
     return {
-      elapsedTime: async (msg?: ActionLessMessage<object>): Promise<ElapsedTime> => {
-        if (msg) {
-          await notify({
-            ...msg,
-            data: msg.data,
-          });
-        }
+      elapsedTime: async (): Promise<ElapsedTime> => {
         return time();
       },
 
@@ -148,20 +143,22 @@ export class SpanMaker<S extends object, E extends SpanEnd> {
 
       timing: async(maker: IoMessageMaker<Duration>, message?: string): Promise<ElapsedTime> => {
         const duration = time();
-        const endMsg = message ? message : timingMsg;
-        await notify(maker.msg(util.format(endMsg, this.definition.name, duration.asSec), {
+        const timingMsg = message ? message : util.format(timingMsgTemplate, this.definition.name, duration.asSec);
+        await notify(maker.msg(timingMsg, {
           duration: duration.asMs,
         }));
         return duration;
       },
 
-      end: async (a: any, b?: ForceEmpty<Optional<E, keyof SpanEnd>>): Promise<ElapsedTime> => {
+      end: async (x: any, y?: ForceEmpty<Optional<E, keyof SpanEnd>>): Promise<ElapsedTime> => {
         const duration = time();
-        const endMsg = b ? a : timingMsg;
-        const endPayload = b ?? a;
+
+        const endInput = parseArgs<ForceEmpty<Optional<E, keyof SpanEnd>>>(x, y);
+        const endMsg = endInput.message ?? util.format(timingMsgTemplate, this.definition.name, duration.asSec);
+        const endPayload = endInput.payload;
 
         await notify(this.definition.end.msg(
-          util.format(endMsg, this.definition.name, duration.asSec), {
+          endMsg, {
             duration: duration.asMs,
             ...endPayload,
           } as E));
@@ -172,6 +169,22 @@ export class SpanMaker<S extends object, E extends SpanEnd> {
   }
 }
 
+function parseArgs<S extends object>(first: any, second?: S): { message: string | undefined; payload: S } {
+  const firstIsMessage = typeof first === 'string';
+
+  // When the first argument is a string or we have a second argument, then the first arg is the message
+  const message = (firstIsMessage || second) ? first : undefined;
+
+  // When the first argument is a string or we have a second argument,
+  // then the second arg is the payload, otherwise the first arg is the payload
+  const payload = (firstIsMessage || second) ? second : first;
+
+  return {
+    message,
+    payload,
+  };
+}
+
 function withSpanId(span: string, message: ActionLessMessage<unknown>): ActionLessMessage<unknown> {
   return {
     ...message,
diff --git a/packages/@aws-cdk/tmp-toolkit-helpers/test/api/io/span.test.ts b/packages/@aws-cdk/tmp-toolkit-helpers/test/api/io/span.test.ts
@@ -0,0 +1,218 @@
+/* eslint-disable @typescript-eslint/unbound-method */
+import type { IoHelper, SpanDefinition } from '../../../src/api/io/private';
+import { SpanMaker } from '../../../src/api/io/private';
+import * as maker from '../../../src/api/io/private/message-maker';
+
+describe('SpanMaker', () => {
+  let ioHelper: jest.Mocked<IoHelper>;
+
+  beforeEach(() => {
+    ioHelper = {
+      notify: jest.fn(),
+    } as any;
+  });
+
+  test('begin creates span with unique id and sends start message', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{}, { duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{}>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'test',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'test',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    await spanMaker.begin('Test span', {});
+
+    // THEN
+    expect(ioHelper.notify).toHaveBeenCalledTimes(1);
+    const notifyCall = ioHelper.notify.mock.calls[0][0];
+    expect(notifyCall.message).toBe('Test span');
+    expect(notifyCall.span).toBeDefined(); // UUID should be set
+  });
+
+  test('end returns elapsed time and sends end message', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{}, { duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{}>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Starting span',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Ending span',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    const messageSpan = await spanMaker.begin('Test span', {});
+    const result = await messageSpan.end();
+
+    // THEN
+    expect(result.asMs).toBeGreaterThanOrEqual(0);
+    expect(ioHelper.notify).toHaveBeenCalledTimes(2);
+    const endCall = ioHelper.notify.mock.calls[1][0] as any;
+    expect(endCall.message).toContain('Test Span time');
+    expect(endCall.span).toBeDefined();
+    expect(endCall.data.duration).toBeGreaterThanOrEqual(0);
+  });
+
+  test('intermediate messages are sent with same span id', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{}, { duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{}>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Starting span',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Ending span',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    const messageSpan = await spanMaker.begin('Test span', {});
+    await messageSpan.notify({
+      message: 'Intermediate message',
+      code: 'CDK_TOOLKIT_I1234',
+      time: new Date(),
+      level: 'error',
+      data: undefined,
+    });
+    await messageSpan.end();
+
+    // THEN
+    expect(ioHelper.notify).toHaveBeenCalledTimes(3);
+    const spanId = ioHelper.notify.mock.calls[0][0].span;
+    expect(ioHelper.notify.mock.calls[1][0].span).toBe(spanId);
+    expect(ioHelper.notify.mock.calls[2][0].span).toBe(spanId);
+  });
+
+  test('end with payload overrides default elapsed time payload', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{}, { customField: string; duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{}>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Starting span',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ customField: string; duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Ending span',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    const messageSpan = await spanMaker.begin('Test span', {});
+    const result = await messageSpan.end({ customField: 'test value' });
+
+    // THEN
+    expect(result).toEqual({ asMs: expect.any(Number), asSec: expect.any(Number) });
+    const endCall = ioHelper.notify.mock.calls[1][0];
+    expect(endCall.data).toEqual(expect.objectContaining({ customField: 'test value' }));
+  });
+
+  test('begin with payload includes payload in start message', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{ startField: string }, { duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{ startField: string }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Starting span',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Ending span',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    await spanMaker.begin('Test span', { startField: 'test value' });
+
+    // THEN
+    const startCall = ioHelper.notify.mock.calls[0][0];
+    expect(startCall.data).toEqual({ startField: 'test value' });
+  });
+
+  test('endWithMessage allows custom end message', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{}, { duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{}>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Starting span',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Ending span',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    const messageSpan =await spanMaker.begin('Test span', {});
+    await messageSpan.end('Custom end message');
+
+    // THEN
+    const endCall = ioHelper.notify.mock.calls[1][0];
+    expect(endCall.message).toBe('Custom end message');
+  });
+
+  test('timing sends timing message with elapsed time', async () => {
+    // GIVEN
+    const definition: SpanDefinition<{}, { duration: number }> = {
+      name: 'Test Span',
+      start: maker.info<{}>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Starting span',
+        interface: 'stuff',
+      }),
+      end: maker.info<{ duration: number }>({
+        code: 'CDK_TOOLKIT_I1234',
+        description: 'Ending span',
+        interface: 'stuff',
+      }),
+    };
+    const spanMaker = new SpanMaker(ioHelper, definition);
+
+    // WHEN
+    const messageSpan = await spanMaker.begin('Test span', {});
+    const elapsedTime = await messageSpan.timing(maker.info({
+      code: 'CDK_TOOLKIT_I1234',
+      description: 'Timing message',
+      interface: 'stuff',
+    }), 'Custom timing message');
+
+    // THEN
+    expect(elapsedTime.asMs).toBeGreaterThanOrEqual(0);
+    expect(ioHelper.notify).toHaveBeenCalledTimes(2);
+    const timingCall = ioHelper.notify.mock.calls[1][0] as any;
+    expect(timingCall.message).toBe('Custom timing message');
+    expect(timingCall.span).toBeDefined();
+    expect(timingCall.data.duration).toBeGreaterThanOrEqual(0);
+  });
+});
