return one object only from renderTo functions, allow assertions on takeRender/takeSnapshot

Author: phryneas

src/assertable.ts

@@ -0,0 +1,26 @@
+import { RenderStream } from "./profile/profile.js";
+
+export const assertableSymbol = Symbol.for(
+  "@testing-library/react-render-stream:assertable"
+);
+
+/**
+ * A function or object that can be used in assertions, like e.g.
+ ```ts
+ expect(assertable).toRerender()
+ expect(assertable).not.toRerender()
+ expect(assertable).toRenderExactlyTimes(3)
+ ```
+ */
+export type Assertable = {
+  [assertableSymbol]: RenderStream<any>;
+};
+
+export function markAssertable<T extends {}>(
+  assertable: T,
+  stream: RenderStream<any>
+): T & Assertable {
+  return Object.assign(assertable, {
+    [assertableSymbol]: stream,
+  });
+}

src/index.ts

@@ -13,3 +13,5 @@ export type { SyncScreen } from "./profile/Render.js";
 
 export { renderToRenderStream } from "./renderToRenderStream.js";
 export { renderHookToSnapshotStream } from "./renderHookToSnapshotStream.js";
+
+export type { Assertable } from "./assertable.js";

src/jest/ProfiledComponent.ts

@@ -1,17 +1,19 @@
 import type { MatcherFunction } from "expect";
 import { WaitForRenderTimeoutError } from "@testing-library/react-render-stream";
 import type {
+  Assertable,
   NextRenderOptions,
   RenderStream,
 } from "@testing-library/react-render-stream";
+// explicitly imported the symbol from the internal file
+// this will bundle the `Symbol.for` call twice, but we keep it private
+import { assertableSymbol } from "../assertable.js";
 
 export const toRerender: MatcherFunction<[options?: NextRenderOptions]> =
   async function (actual, options) {
-    const _profiler = actual as RenderStream<any>;
+    const _profiler = actual as RenderStream<any> | Assertable;
     const profiler =
-      "Profiler" in _profiler
-        ? (_profiler.Profiler as RenderStream<any>)
-        : _profiler;
+      assertableSymbol in _profiler ? _profiler[assertableSymbol] : _profiler;
     const hint = this.utils.matcherHint("toRerender", "ProfiledComponent", "");
     let pass = true;
     try {
@@ -42,11 +44,9 @@ const failed = {};
 export const toRenderExactlyTimes: MatcherFunction<
   [times: number, options?: NextRenderOptions]
 > = async function (actual, times, optionsPerRender) {
-  const _profiler = actual as RenderStream<any>;
+  const _profiler = actual as RenderStream<any> | Assertable;
   const profiler =
-    "Profiler" in _profiler
-      ? (_profiler.Profiler as RenderStream<any>)
-      : _profiler;
+    assertableSymbol in _profiler ? _profiler[assertableSymbol] : _profiler;
   const options = { timeout: 100, ...optionsPerRender };
   const hint = this.utils.matcherHint("toRenderExactlyTimes");
   let pass = true;

src/jest/index.ts

@@ -1,23 +1,31 @@
 import { expect } from "@jest/globals";
 import { toRerender, toRenderExactlyTimes } from "./ProfiledComponent.js";
-import type { NextRenderOptions, RenderStream } from "../index.js";
+import type {
+  NextRenderOptions,
+  RenderStream,
+  Assertable,
+} from "@testing-library/react-render-stream";
 
 expect.extend({
   toRerender,
   toRenderExactlyTimes,
 });
-interface ApolloCustomMatchers<R = void, T = {}> {
-  toRerender: T extends RenderStream<any> | unknown // TODO
+interface CustomMatchers<R = void, T = {}> {
+  toRerender: T extends RenderStream<any> | Assertable
     ? (options?: NextRenderOptions) => Promise<R>
-    : { error: "matcher needs to be called on a ProfiledComponent instance" };
+    : {
+        error: "matcher needs to be called on a `takeRender` function, `takeSnapshot` function or `RenderStream` instance";
+      };
 
-  toRenderExactlyTimes: T extends RenderStream<any> | unknown // TODO
+  toRenderExactlyTimes: T extends RenderStream<any> | Assertable
     ? (count: number, options?: NextRenderOptions) => Promise<R>
-    : { error: "matcher needs to be called on a ProfiledComponent instance" };
+    : {
+        error: "matcher needs to be called on a `takeRender` function, `takeSnapshot` function or `RenderStream` instance";
+      };
 }
 
 declare global {
   namespace jest {
-    interface Matchers<R = void, T = {}> extends ApolloCustomMatchers<R, T> {}
+    interface Matchers<R = void, T = {}> extends CustomMatchers<R, T> {}
   }
 }

src/profile/profile.tsx

@@ -7,6 +7,7 @@ import type { ProfilerContextValue } from "./context.js";
 import { ProfilerContextProvider, useProfilerContext } from "./context.js";
 import { disableActWarnings } from "./disableActWarnings.js";
 import { render as baseRender, RenderOptions } from "@testing-library/react";
+import { Assertable, markAssertable } from "../assertable.js";
 
 export type ValidSnapshot =
   | void
@@ -59,7 +60,8 @@ export interface ProfiledComponentFields<Snapshot> {
    * If no render has happened yet, it will wait for the next render to happen.
    * @throws {WaitForRenderTimeoutError} if no render happens within the timeout
    */
-  takeRender(options?: NextRenderOptions): Promise<Render<Snapshot>>;
+  takeRender: Assertable &
+    ((options?: NextRenderOptions) => Promise<Render<Snapshot>>);
   /**
    * Returns the total number of renders.
    */
@@ -241,7 +243,9 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
     });
   }) as typeof baseRender;
 
-  const Profiler: RenderStreamWithRenderFn<Snapshot> = Object.assign(
+  let Profiler: RenderStreamWithRenderFn<Snapshot> = {} as any;
+  Profiler = Object.assign(
+    Profiler as {},
     {
       replaceSnapshot,
       mergeSnapshot,
@@ -269,7 +273,9 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
           ...options,
         });
       },
-      async takeRender(options: NextRenderOptions = {}) {
+      takeRender: markAssertable(async function takeRender(
+        options: NextRenderOptions = {}
+      ) {
         // In many cases we do not control the resolution of the suspended
         // promise which results in noisy tests when the profiler due to
         // repeated act warnings.
@@ -290,7 +296,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
             iteratorPosition++;
           }
         }
-      },
+      }, Profiler),
       getCurrentRender() {
         // The "current" render should point at the same render that the most
         // recent `takeRender` call returned, so we need to get the "previous"

src/renderHookToSnapshotStream.ts

@@ -1,36 +1,58 @@
 import { RenderHookOptions } from "@testing-library/react";
 import {
   createProfiler,
+  NextRenderOptions,
   ProfiledComponentFields,
   ValidSnapshot,
 } from "./profile/profile.js";
 import { Render } from "./profile/Render.js";
 import { createElement } from "react";
-
-type StringReplaceRenderWithSnapshot<T extends string> =
-  T extends `${infer Pre}Render${infer Post}` ? `${Pre}Snapshot${Post}` : T;
-
-type ResultReplaceRenderWithSnapshot<T> = T extends (
-  ...args: infer Args
-) => Render<infer Snapshot>
-  ? (...args: Args) => Snapshot
-  : T extends (...args: infer Args) => Promise<Render<infer Snapshot>>
-    ? (...args: Args) => Promise<Snapshot>
-    : T;
-
-type ProfiledHookFields<ReturnValue> =
-  ProfiledComponentFields<ReturnValue> extends infer PC
-    ? {
-        [K in keyof PC as StringReplaceRenderWithSnapshot<
-          K & string
-        >]: ResultReplaceRenderWithSnapshot<PC[K]>;
-      }
-    : never;
+import { Assertable, assertableSymbol, markAssertable } from "./assertable.js";
 
 /** @internal */
-export interface ProfiledHook<Props, ReturnValue extends ValidSnapshot>
-  extends ProfiledHookFields<ReturnValue> {
-  //Profiler: RenderStream<ReturnValue>;
+export interface ProfiledHook<Snapshot extends ValidSnapshot>
+  extends Assertable {
+  /**
+   * An array of all renders that have happened so far.
+   * Errors thrown during component render will be captured here, too.
+   */
+  renders: Array<
+    Render<Snapshot> | { phase: "snapshotError"; count: number; error: unknown }
+  >;
+  /**
+   * Peeks the next render from the current iterator position, without advancing the iterator.
+   * If no render has happened yet, it will wait for the next render to happen.
+   * @throws {WaitForRenderTimeoutError} if no render happens within the timeout
+   */
+  peekSnapshot(options?: NextRenderOptions): Promise<Snapshot>;
+  /**
+   * Iterates to the next render and returns it.
+   * If no render has happened yet, it will wait for the next render to happen.
+   * @throws {WaitForRenderTimeoutError} if no render happens within the timeout
+   */
+  takeSnapshot: Assertable &
+    ((options?: NextRenderOptions) => Promise<Snapshot>);
+  /**
+   * Returns the total number of renders.
+   */
+  totalSnapshotCount(): number;
+  /**
+   * Returns the current render.
+   * @throws {Error} if no render has happened yet
+   */
+  getCurrentSnapshot(): Snapshot;
+  /**
+   * Waits for the next render to happen.
+   * Does not advance the render iterator.
+   */
+  waitForNextSnapshot(options?: NextRenderOptions): Promise<Snapshot>;
+}
+
+interface HookSnapshotStream<Props, ReturnValue extends ValidSnapshot>
+  extends ProfiledHook<ReturnValue>,
+    Assertable {
+  rerender: (rerenderCallbackProps: Props) => void;
+  unmount: () => void;
 }
 
 export function renderHookToSnapshotStream<
@@ -39,13 +61,7 @@ export function renderHookToSnapshotStream<
 >(
   renderCallback: (props: Props) => ReturnValue,
   { initialProps, ...options }: RenderHookOptions<Props> = {}
-): [
-  stream: ProfiledHook<Props, ReturnValue>,
-  renderResult: {
-    rerender: (rerenderCallbackProps: Props) => void;
-    unmount: () => void;
-  },
-] {
+): HookSnapshotStream<Props, ReturnValue> {
   const { render, ...stream } = createProfiler<ReturnValue>();
 
   const ProfiledHook: React.FC<Props> = (props) => {
@@ -62,26 +78,23 @@ export function renderHookToSnapshotStream<
     return baseRerender(createElement(ProfiledHook, rerenderCallbackProps));
   }
 
-  return [
-    Object.assign({}, stream, {
-      renders: stream.renders,
-      totalSnapshotCount: stream.totalRenderCount,
-      async peekSnapshot(options) {
-        return (await stream.peekRender(options)).snapshot;
-      },
-      async takeSnapshot(options) {
-        return (await stream.takeRender(options)).snapshot;
-      },
-      getCurrentSnapshot() {
-        return stream.getCurrentRender().snapshot;
-      },
-      async waitForNextSnapshot(options) {
-        return (await stream.waitForNextRender(options)).snapshot;
-      },
-    } satisfies ProfiledHookFields<ReturnValue>),
-    {
-      rerender,
-      unmount,
+  return {
+    [assertableSymbol]: stream,
+    renders: stream.renders,
+    totalSnapshotCount: stream.totalRenderCount,
+    async peekSnapshot(options) {
+      return (await stream.peekRender(options)).snapshot;
+    },
+    takeSnapshot: markAssertable(async function takeSnapshot(options) {
+      return (await stream.takeRender(options)).snapshot;
+    }, stream),
+    getCurrentSnapshot() {
+      return stream.getCurrentRender().snapshot;
     },
-  ];
+    async waitForNextSnapshot(options) {
+      return (await stream.waitForNextRender(options)).snapshot;
+    },
+    rerender,
+    unmount,
+  };
 }

src/renderToRenderStream.ts

@@ -13,11 +13,11 @@ import {
 type RenderOptions<Snapshot extends ValidSnapshot = void> = BaseOptions &
   ProfilerOptions<Snapshot>;
 
-type RenderResult<Snapshot extends ValidSnapshot = void> = [
-  Stream: ProfiledComponentFields<Snapshot> &
-    ProfiledComponentOnlyFields<Snapshot>,
-  renderResultPromise: Promise<BaseResult>,
-];
+type RenderResult<Snapshot extends ValidSnapshot = void> =
+  ProfiledComponentFields<Snapshot> &
+    ProfiledComponentOnlyFields<Snapshot> & {
+      renderResultPromise: Promise<BaseResult>;
+    };
 
 /**
  * Render into a container which is appended to document.body. It should be used with cleanup.
@@ -40,5 +40,5 @@ export function renderToRenderStream<Snapshot extends ValidSnapshot = void>(
     skipNonTrackingRenders,
   });
   const renderResultPromise = Promise.resolve().then(() => render(ui, options));
-  return [stream, renderResultPromise];
+  return { ...stream, renderResultPromise };
 }