feat: Add snap_startTrace and snap_endTrace support to snaps-jest

packages/examples/packages/preinstalled/src/index.test.tsx

@@ -124,6 +124,32 @@ describe('onRpcRequest', () => {
       });
     });
   });
+
+  describe('startTrace + endTrace', () => {
+    it('starts and ends a trace', async () => {
+      const { request } = await installSnap();
+
+      const response = await request({
+        method: 'startTrace',
+      });
+
+      expect(response).toRespondWith({
+        /* eslint-disable @typescript-eslint/naming-convention */
+        _traceId: expect.any(String),
+        _spanId: expect.any(String),
+        /* eslint-enable @typescript-eslint/naming-convention */
+      });
+
+      const endResponse = await request({
+        method: 'endTrace',
+      });
+
+      expect(endResponse).toRespondWith(null);
+      expect(endResponse).toTrace({
+        name: 'Test Snap Trace',
+      });
+    });
+  });
 });
 
 describe('onSettingsPage', () => {

packages/snaps-jest/src/global.ts

@@ -107,6 +107,21 @@ interface SnapsMatchers {
    * });
    */
   toTrackEvent(event?: unknown): void;
+
+  /**
+   * Assert that the Snap started and ended a trace with the expected
+   * parameters. This is equivalent to calling
+   * `expect(response.tracked.traces).toContainEqual(span)`.
+   *
+   * @param trace - The expected trace parameters.
+   * @throws If the snap did not end a trace with the expected parameters.
+   * @example
+   * const response = await request({ method: 'foo' });
+   * expect(response).toTrace({
+   *   name: 'My Trace',
+   * });
+   */
+  toTrace(trace?: unknown): void;
 }
 
 // Extend the `expect` interface with the new matchers. This is used when

packages/snaps-jest/src/matchers.test.tsx

@@ -13,6 +13,7 @@ import {
   toRespondWith,
   toRespondWithError,
   toSendNotification,
+  toTrace,
   toTrackError,
   toTrackEvent,
 } from './matchers';
@@ -29,6 +30,7 @@ expect.extend({
   toRender,
   toTrackError,
   toTrackEvent,
+  toTrace,
 });
 
 describe('toRespondWith', () => {
@@ -949,3 +951,166 @@ describe('toTrackEvent', () => {
     });
   });
 });
+
+describe('toTrace', () => {
+  it('passes when the trace is correct', () => {
+    expect(
+      getMockResponse({
+        tracked: {
+          traces: [
+            {
+              id: '1',
+              name: 'foo',
+              parentContext: { bar: 'baz' },
+              startTime: 1234567890,
+              data: { qux: 'quux' },
+              tags: { corge: 'grault' },
+            },
+          ],
+        },
+      }),
+    ).toTrace({
+      id: '1',
+      name: 'foo',
+      parentContext: { bar: 'baz' },
+      startTime: 1234567890,
+      data: { qux: 'quux' },
+      tags: { corge: 'grault' },
+    });
+  });
+
+  it('passes when the partial trace is correct', () => {
+    expect(
+      getMockResponse({
+        tracked: {
+          traces: [
+            {
+              id: '1',
+              name: 'foo',
+              parentContext: { bar: 'baz' },
+              startTime: 1234567890,
+            },
+          ],
+        },
+      }),
+    ).toTrace({
+      id: '1',
+      name: 'foo',
+      parentContext: { bar: 'baz' },
+      startTime: 1234567890,
+    });
+  });
+
+  it('passes when any trace is tracked', () => {
+    expect(
+      getMockResponse({
+        tracked: {
+          traces: [
+            {
+              id: '1',
+              name: 'foo',
+              parentContext: { bar: 'baz' },
+              startTime: 1234567890,
+              data: { qux: 'quux' },
+              tags: { corge: 'grault' },
+            },
+          ],
+        },
+      }),
+    ).toTrace();
+  });
+
+  it('fails when the trace is incorrect', () => {
+    expect(() =>
+      expect(
+        getMockResponse({
+          tracked: {
+            traces: [
+              {
+                id: '1',
+                name: 'foo',
+                parentContext: { bar: 'baz' },
+                startTime: 1234567890,
+                data: { qux: 'quux' },
+                tags: { corge: 'grault' },
+              },
+            ],
+          },
+        }),
+      ).toTrace({
+        id: '2',
+        name: 'foo',
+        parentContext: { bar: 'baz' },
+        startTime: 1234567890,
+        data: { qux: 'quux' },
+        tags: { corge: 'grault' },
+      }),
+    ).toThrow('Received');
+  });
+
+  describe('not', () => {
+    it('passes when the trace does not match', () => {
+      expect(
+        getMockResponse({
+          tracked: {
+            traces: [
+              {
+                id: '1',
+                name: 'foo',
+                parentContext: { bar: 'baz' },
+                startTime: 1234567890,
+                data: { qux: 'quux' },
+                tags: { corge: 'grault' },
+              },
+            ],
+          },
+        }),
+      ).not.toTrace({
+        id: '2',
+        name: 'foo',
+        parentContext: { bar: 'baz' },
+        startTime: 1234567890,
+        data: { qux: 'quux' },
+        tags: { corge: 'grault' },
+      });
+    });
+
+    it('passes when there are no traces', () => {
+      expect(
+        getMockResponse({
+          tracked: {
+            traces: [],
+          },
+        }),
+      ).not.toTrace();
+    });
+
+    it('fails when the trace matches', () => {
+      expect(() =>
+        expect(
+          getMockResponse({
+            tracked: {
+              traces: [
+                {
+                  id: '1',
+                  name: 'foo',
+                  parentContext: { bar: 'baz' },
+                  startTime: 1234567890,
+                  data: { qux: 'quux' },
+                  tags: { corge: 'grault' },
+                },
+              ],
+            },
+          }),
+        ).not.toTrace({
+          id: '1',
+          name: 'foo',
+          parentContext: { bar: 'baz' },
+          startTime: 1234567890,
+          data: { qux: 'quux' },
+          tags: { corge: 'grault' },
+        }),
+      ).toThrow('Expected not to trace with data');
+    });
+  });
+});

packages/snaps-jest/src/matchers.ts

@@ -9,6 +9,7 @@ import type {
   Component,
   NotificationType,
   TrackableError,
+  TraceRequest,
 } from '@metamask/snaps-sdk';
 import type { JSXElement, SnapNode } from '@metamask/snaps-sdk/jsx';
 import { isJSXElementUnsafe, JSXElementStruct } from '@metamask/snaps-sdk/jsx';
@@ -376,7 +377,8 @@ export const toTrackError: MatcherFunction<
 
   const errorValidator = (error: SnapResponse['tracked']['errors'][number]) => {
     if (!errorData) {
-      // If no error data is provided, we just check for the existence of an error.
+      // If no error data is provided, we just check for the existence of an
+      // error.
       return true;
     }
 
@@ -411,7 +413,8 @@ export const toTrackEvent: MatcherFunction<[eventData?: Json | undefined]> =
       event: SnapResponse['tracked']['events'][number],
     ) => {
       if (!eventData) {
-        // If no event data is provided, we just check for the existence of an event.
+        // If no event data is provided, we just check for the existence of an
+        // event.
         return true;
       }
 
@@ -438,11 +441,48 @@ export const toTrackEvent: MatcherFunction<[eventData?: Json | undefined]> =
     return { message, pass };
   };
 
+export const toTrace: MatcherFunction<[traceData?: TraceRequest]> = function (
+  actual,
+  traceData,
+) {
+  assertActualIsSnapResponse(actual, 'toTrace');
+
+  const traceValidator = (trace: SnapResponse['tracked']['traces'][number]) => {
+    if (!traceData) {
+      // If no trace data is provided, we just check for the existence of a
+      // trace.
+      return true;
+    }
+
+    return this.equals(trace, traceData);
+  };
+
+  const { traces } = actual.tracked;
+  const pass = traces.some(traceValidator);
+
+  const message = pass
+    ? () =>
+        `${this.utils.matcherHint('.not.toTrace')}\n\n` +
+        `Expected not to trace with data: ${this.utils.printExpected(
+          traceData,
+        )}\n` +
+        `Received traces: ${this.utils.printReceived(traces)}`
+    : () =>
+        `${this.utils.matcherHint('.toTrace')}\n\n` +
+        `Expected to trace with data: ${this.utils.printExpected(
+          traceData,
+        )}\n` +
+        `Received traces: ${this.utils.printReceived(traces)}`;
+
+  return { message, pass };
+};
+
 expect.extend({
   toRespondWith,
   toRespondWithError,
   toSendNotification,
   toRender,
   toTrackError,
   toTrackEvent,
+  toTrace,
 });

packages/snaps-jest/src/test-utils/response.ts

@@ -25,13 +25,11 @@ export function getMockResponse({
   tracked = {
     errors: [],
     events: [],
+    traces: [],
   },
   getInterface,
 }: Omit<Partial<SnapResponseWithInterface>, 'tracked'> & {
-  tracked?: {
-    errors?: SnapResponse['tracked']['errors'];
-    events?: SnapResponse['tracked']['events'];
-  };
+  tracked?: Partial<SnapResponse['tracked']>;
 }): SnapResponse {
   return {
     id,
@@ -40,6 +38,7 @@ export function getMockResponse({
     tracked: {
       errors: [],
       events: [],
+      traces: [],
       ...tracked,
     },
     ...(getInterface ? { getInterface } : {}),

packages/snaps-simulation/src/methods/hooks/end-trace.test.ts

@@ -0,0 +1,31 @@
+import { getEndTraceImplementation } from './end-trace';
+import { createStore, startTrace } from '../../store';
+import { getMockOptions } from '../../test-utils';
+
+describe('getEndTraceImplementation', () => {
+  it('returns the implementation of the `endTrace` hook', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    store.dispatch(
+      startTrace({
+        name: 'Test Trace',
+        id: 'test-trace-id',
+      }),
+    );
+
+    const fn = getEndTraceImplementation(runSaga);
+
+    expect(
+      fn({
+        name: 'Test Trace',
+        id: 'test-trace-id',
+      }),
+    ).toBeNull();
+
+    expect(store.getState().trackables.pendingTraces).toHaveLength(0);
+    expect(store.getState().trackables.traces).toHaveLength(1);
+    expect(store.getState().trackables.traces[0]).toStrictEqual({
+      name: 'Test Trace',
+      id: 'test-trace-id',
+    });
+  });
+});