feat: expose takeSnapshot() as public API

packages/vitest/src/browser/public/takeSnapshot.browser.test.ts

@@ -0,0 +1,98 @@
+import { beforeEach, expect, test } from 'vitest';
+import { commands } from 'vitest/browser';
+import { takeSnapshot } from './takeSnapshot';
+
+beforeEach(() => {
+  document.body.innerHTML = '';
+});
+
+test('saves snapshot on server', async ({ task }) => {
+  const h1 = document.createElement('h1');
+  h1.textContent = 'Example heading';
+  document.body.appendChild(h1);
+
+  await takeSnapshot('example');
+
+  const snapshots = await commands.__chromatic_getSnapshots(task.id);
+  expect(snapshots).toHaveProperty('example');
+
+  expect(snapshots.example).toMatchInlineSnapshot(`
+    {
+      "attributes": {},
+      "childNodes": [
+        {
+          "id": "number",
+          "textContent": "Example heading",
+          "type": 3,
+        },
+      ],
+      "id": "number",
+      "tagName": "h1",
+      "type": 2,
+    }
+  `);
+});
+
+test('saves multiple snapshots', async ({ task }) => {
+  await takeSnapshot('example-1');
+  await takeSnapshot('example-2');
+
+  const snapshots = await commands.__chromatic_getSnapshots(task.id);
+
+  expect(snapshots).toHaveProperty('example-1');
+  expect(snapshots).toHaveProperty('example-2');
+});
+
+test('implicit snapshot names increment', async ({ task }) => {
+  await takeSnapshot();
+
+  {
+    const snapshots = await commands.__chromatic_getSnapshots(task.id);
+
+    expect(snapshots).toHaveProperty('Snapshot #1');
+    expect(Object.keys(snapshots)).toHaveLength(1);
+  }
+
+  await takeSnapshot();
+
+  {
+    const snapshots = await commands.__chromatic_getSnapshots(task.id);
+
+    expect.soft(snapshots).toHaveProperty('Snapshot #1');
+    expect.soft(snapshots).toHaveProperty('Snapshot #2');
+    expect.soft(Object.keys(snapshots)).toHaveLength(2);
+  }
+});
+
+test('blob URLs are replaced with data URLs', async ({ task }) => {
+  const blob = await fetch(
+    'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAASCAQAAADIvofAAAAABGdBTUEAALGPC/xhBQAAACBjSFJNAAB6JgAAgIQAAPoAAACA6AAAdTAAAOpgAAA6mAAAF3CculE8AAAAAmJLR0QA/4ePzL8AAAAHdElNRQfhBhAPKSstM+EuAAAAvUlEQVQY05WQIW4CYRgF599gEZgeoAKBWIfCNSmVvQMe3wv0ChhIViKwtTQEAYJwhgpISBA0JSxNIdlB7LIGTJ/8kpeZ7wW5TcT9o/QNBtvOrrWMrtg0sSGOFeELbHlCDsQ+ukeYiHNFJPHBDRKlQKVEbFkLUT3AiAxI6VGCXsWXAoQLBUl5E7HjUFwiyI4zf/wWoB3CFnxX5IeGdY8IGU/iwE9jcZrLy4pnEat+FL4hf/cbqREKo/Cf6W5zASVMeh234UtGAAAAJXRFWHRkYXRlOmNyZWF0ZQAyMDE3LTA2LTE2VDE1OjQxOjQzLTA3OjAwd1xNIQAAACV0RVh0ZGF0ZTptb2RpZnkAMjAxNy0wNi0xNlQxNTo0MTo0My0wNzowMAYB9Z0AAAAASUVORK5CYII='
+  ).then((res) => res.blob());
+
+  const img = document.createElement('img');
+  const loaded = new Promise((resolve) => (img.onload = resolve));
+  img.src = window.URL.createObjectURL(blob);
+  document.body.appendChild(img);
+
+  await loaded;
+  await takeSnapshot('example');
+
+  const snapshots = await commands.__chromatic_getSnapshots(task.id);
+  expect(snapshots).toHaveProperty('example');
+
+  const snapshot = snapshots.example;
+  expect(snapshot).toMatchInlineSnapshot(`
+    {
+      "attributes": {
+        "src": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAASCAQAAADIvofAAAAABGdBTUEAALGPC/xhBQAAACBjSFJNAAB6JgAAgIQAAPoAAACA6AAAdTAAAOpgAAA6mAAAF3CculE8AAAAAmJLR0QA/4ePzL8AAAAHdElNRQfhBhAPKSstM+EuAAAAvUlEQVQY05WQIW4CYRgF599gEZgeoAKBWIfCNSmVvQMe3wv0ChhIViKwtTQEAYJwhgpISBA0JSxNIdlB7LIGTJ/8kpeZ7wW5TcT9o/QNBtvOrrWMrtg0sSGOFeELbHlCDsQ+ukeYiHNFJPHBDRKlQKVEbFkLUT3AiAxI6VGCXsWXAoQLBUl5E7HjUFwiyI4zf/wWoB3CFnxX5IeGdY8IGU/iwE9jcZrLy4pnEat+FL4hf/cbqREKo/Cf6W5zASVMeh234UtGAAAAJXRFWHRkYXRlOmNyZWF0ZQAyMDE3LTA2LTE2VDE1OjQxOjQzLTA3OjAwd1xNIQAAACV0RVh0ZGF0ZTptb2RpZnkAMjAxNy0wNi0xNlQxNTo0MTo0My0wNzowMAYB9Z0AAAAASUVORK5CYII=",
+      },
+      "childNodes": [],
+      "id": "number",
+      "rootId": 301,
+      "tagName": "img",
+      "type": 2,
+    }
+  `);
+
+  expect(img.src).toMatch(/^blob:/);
+});

packages/vitest/src/browser/public/takeSnapshot.test.ts

@@ -0,0 +1,87 @@
+import { expect, test } from 'vitest';
+import { runFixture } from '../../../test/utils/node';
+
+/** See {@link file://./../../../test/fixtures/take-snapshot.test.ts} */
+const include = ['take-snapshot.test.ts'];
+
+test('provides descriptive error when called in non-registered test', async () => {
+  const { stderr } = await runFixture(
+    {
+      include,
+      provide: { testName: 'one' },
+    },
+    { disabled: true }
+  );
+
+  expect(stderr).toMatchInlineSnapshot(`
+    "
+    ⎯⎯⎯⎯⎯⎯⎯ Failed Tests 1 ⎯⎯⎯⎯⎯⎯⎯
+
+     FAIL   chromium  take-snapshot.test.ts > test #1
+    TypeError: takeSnapshot() cannot be called in a test that is not registered for Chromatic plugin.
+    Make sure chromium project has chromaticPlugin() enabled.
+     ❯ take-snapshot.test.ts:7:8
+          5|   document.body.innerHTML = '<h1>Example heading</h1>';
+          6|
+          7|   await takeSnapshot();
+           |        ^
+          8|
+          9|   expect.fail('Should not reach this point');
+
+    ⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯[1/1]⎯"
+  `);
+});
+
+test('provides descriptive error when called outside of a test()', async () => {
+  const { stderr } = await runFixture({ include, provide: { testName: 'two' } });
+
+  expect(stderr).toMatchInlineSnapshot(`
+    "
+    ⎯⎯⎯⎯⎯⎯ Failed Suites 1 ⎯⎯⎯⎯⎯⎯⎯
+
+     FAIL   chromium  take-snapshot.test.ts > suite
+    TypeError: takeSnapshot() must be called within a test()
+     ❯ take-snapshot.test.ts:14:10
+         12| describe.runIf(inject('testName') === 'two')('suite', async () => {
+         13|   beforeAll(async () => {
+         14|     await takeSnapshot();
+           |          ^
+         15|   });
+         16|
+
+    ⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯[1/1]⎯"
+  `);
+});
+
+test('provides descriptive error when not awaited', async () => {
+  const { stderr } = await runFixture({ include, provide: { testName: 'three' } });
+
+  expect(stderr).toMatchInlineSnapshot(`
+    "
+    ⎯⎯⎯⎯⎯⎯⎯ Failed Tests 1 ⎯⎯⎯⎯⎯⎯⎯
+
+     FAIL   chromium  take-snapshot.test.ts > test #3
+    Error: takeSnapshot() call was not awaited!
+     ❯ take-snapshot.test.ts:23:2
+         21|   document.body.innerHTML = '<h1>Example heading</h1>';
+         22|
+         23|   takeSnapshot(); // Leave the promise floating, no await
+           |  ^
+         24|
+         25|   // another
+
+    ⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯[1/2]⎯
+
+     FAIL   chromium  take-snapshot.test.ts > test #3
+    Error: takeSnapshot() call was not awaited!
+     ❯ take-snapshot.test.ts:26:2
+         24|
+         25|   // another
+         26|   takeSnapshot();
+           |  ^
+         27| });
+         28|
+
+    ⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯⎯[2/2]⎯"
+  `);
+});

packages/vitest/src/browser/public/takeSnapshot.ts

@@ -5,20 +5,65 @@ import { serializedNodeWithId } from '@rrweb/types';
 import { getCurrentTest } from '../getCurrentTest';
 import type {} from '../../node/commands';
 
+interface Options {
+  ignoreUnawaited?: boolean;
+}
+
 /**
  * Take visual regression snapshot of the current state of the DOM.
  */
-export async function takeSnapshot(name?: string) {
+async function takeSnapshot(name?: string): Promise<void>;
+
+/** @internal Pass options when used by automatic snapshots */
+async function takeSnapshot(name: string | undefined, options: Options): Promise<void>;
+
+async function takeSnapshot(name?: string, options?: Options): Promise<void> {
   const test = getCurrentTest();
-  assert(test, 'takeSnapshot() must be called within a test()');
 
-  const domSnapshot = snapshot(document, { recordCanvas: true });
+  if (!test) {
+    throw new TypeError('takeSnapshot() must be called within a test()');
+  }
+
+  if (!test.meta.__chromatic_isRegistered) {
+    throw new TypeError(
+      'takeSnapshot() cannot be called in a test that is not registered for Chromatic plugin.' +
+        `\nMake sure ${test.file.projectName || 'root'} project has chromaticPlugin() enabled.`
+    );
+  }
+
+  test.meta.__chromatic_isTakeSnapshotCalled = true;
 
+  const domSnapshot = snapshot(document, { recordCanvas: true });
   assert(domSnapshot, 'Failed to capture DOM snapshot');
 
-  await replaceBlobUrls(domSnapshot);
+  const save = async () => {
+    await replaceBlobUrls(domSnapshot);
+    await commands.__chromatic_uploadDOMSnapshot(test.id, domSnapshot, name);
+  };
+
+  // Ignore is set when called by automatic snapshots
+  if (options?.ignoreUnawaited) {
+    return await save();
+  }
+
+  /**
+   * Provide descriptive error if the user forgets to await the takeSnapshot() call.
+   * See {@link file://./takeSnapshot.test.ts} for examples.
+   */
+  const error = new Error('takeSnapshot() call was not awaited!');
+  Error.captureStackTrace?.(error, takeSnapshot);
 
-  await commands.__chromatic_uploadDOMSnapshot(test.id, domSnapshot, name);
+  const pendingCall = { promise: save(), error };
+  test.meta.__chromatic_pendingTakeSnapshots ||= [];
+  test.meta.__chromatic_pendingTakeSnapshots.push(pendingCall);
+
+  await pendingCall.promise.finally(() => {
+    const index = test.meta.__chromatic_pendingTakeSnapshots?.indexOf(pendingCall);
+
+    if (index !== undefined && index !== -1) {
+      test.meta.__chromatic_pendingTakeSnapshots?.splice(index, 1);
+    }
+  });
 }
 
 async function replaceBlobUrls(node: serializedNodeWithId) {
@@ -57,3 +102,5 @@ async function toDataURL(url: string): Promise<string> {
     reader.readAsDataURL(blob);
   });
 }
+
+export { takeSnapshot };

packages/vitest/src/browser/setupFile.ts

@@ -7,14 +7,24 @@ import type {} from '../node/commands';
 // Destructuring the context is important, so that possible user-provided fixtures work too
 beforeEach<InternalTestContext>(async ({ task }) => {
   task.meta.__chromatic_isRegistered = true;
+  task.meta.__chromatic_isTakeSnapshotCalled = false;
 
   await commands.__chromatic_interceptFetch(task.id);
 
   return async function afterEach() {
+    const { __chromatic_pendingTakeSnapshots: pendingTakeSnapshots } = task.meta;
+
+    if (pendingTakeSnapshots?.length) {
+      throw new AggregateError(
+        pendingTakeSnapshots.map((call) => call.error),
+        `${pendingTakeSnapshots.length} unawaited takeSnapshot() call(s)`
+      );
+    }
+
     // TODO: Replace with waitForIdleNetwork API in Milestone #4
     await new Promise((resolve) => setTimeout(resolve, 500));
 
-    await takeSnapshot();
+    await takeSnapshot(undefined, { ignoreUnawaited: true });
 
     await commands.__chromatic_writeTestResult(task.id);
 
@@ -25,6 +35,8 @@ beforeEach<InternalTestContext>(async ({ task }) => {
    * Clean internal task meta so that it doesn't show up on Vitest's reporters
    */
   function cleanup() {
+    task.meta.__chromatic_isTakeSnapshotCalled = undefined;
     task.meta.__chromatic_isRegistered = undefined;
+    task.meta.__chromatic_pendingTakeSnapshots = undefined;
   }
 });

packages/vitest/src/index.ts

@@ -1,3 +1,3 @@
 /* Entrypoint for browser context */
 
-export {};
+export { takeSnapshot } from './browser/public/takeSnapshot';

packages/vitest/src/node/commands.browser.test.ts

@@ -4,6 +4,7 @@ import { InternalTestContext } from '../types';
 
 test<InternalTestContext>('setupFile is registered', async (context) => {
   expect(context.task.meta.__chromatic_isRegistered).toBe(true);
+  expect(context.task.meta.__chromatic_isTakeSnapshotCalled).toBe(false);
 });
 
 test('browser commands are available', () => {
@@ -13,6 +14,7 @@ test('browser commands are available', () => {
 
   expect(Object.fromEntries(chromaticCommands)).toMatchInlineSnapshot(`
     {
+      "__chromatic_getSnapshots": [Function],
       "__chromatic_interceptFetch": [Function],
       "__chromatic_reset": [Function],
       "__chromatic_uploadDOMSnapshot": [Function],

packages/vitest/src/node/commands.ts

@@ -105,6 +105,14 @@ export function createCommands(options: ResolvedOptions) {
       resourceArchivers.clear();
       snapshots.clear();
     },
+
+    /**
+     * Get currently save snapshots. Used only during testing.
+     * @internal
+     */
+    async __chromatic_getSnapshots(_, id: string) {
+      return Object.fromEntries(snapshots.get(id) || []);
+    },
   } satisfies Record<ChromaticNamespace, BrowserCommand<any>>;
 
   async function onTestCleanup(id: TestID) {

packages/vitest/src/node/plugin.test.ts

@@ -22,6 +22,7 @@ test('adds browser commands', async () => {
 
   expect(config.browser.commands).toMatchInlineSnapshot(`
     {
+      "__chromatic_getSnapshots": [Function],
       "__chromatic_interceptFetch": [Function],
       "__chromatic_reset": [Function],
       "__chromatic_uploadDOMSnapshot": [Function],

packages/vitest/src/types.ts

@@ -36,6 +36,12 @@ export interface ResolvedOptions
 type InternalMeta = Record<ChromaticNamespace, unknown> & {
   /** Indicates whether Visual Regression tracking is registered */
   __chromatic_isRegistered?: boolean;
+
+  /** Indicates whether `takeSnapshot()` has been called */
+  __chromatic_isTakeSnapshotCalled?: boolean;
+
+  /** Pending `takeSnapshot()` promises */
+  __chromatic_pendingTakeSnapshots?: { promise: Promise<void>; error: Error }[];
 };
 
 /** @internal */

packages/vitest/test/fixtures/take-snapshot.test.ts

@@ -0,0 +1,27 @@
+import { beforeAll, describe, expect, test, inject } from 'vitest';
+import { takeSnapshot } from '../../src';
+
+test.runIf(inject('testName') === 'one')('test #1', async () => {
+  document.body.innerHTML = '<h1>Example heading</h1>';
+
+  await takeSnapshot();
+
+  expect.fail('Should not reach this point');
+});
+
+describe.runIf(inject('testName') === 'two')('suite', async () => {
+  beforeAll(async () => {
+    await takeSnapshot();
+  });
+
+  test('test #2', async () => {});
+});
+
+test.runIf(inject('testName') === 'three')('test #3', async () => {
+  document.body.innerHTML = '<h1>Example heading</h1>';
+
+  takeSnapshot(); // Leave the promise floating, no await
+
+  // another
+  takeSnapshot();
+});