From 2986cde88cb06c5d769ca07da349298f6e1f5773 Mon Sep 17 00:00:00 2001
From: Lenz Weber-Tronic <lorenz.weber-tronic@apollographql.com>
Date: Tue, 8 Oct 2024 18:07:23 +0200
Subject: [PATCH] move `render` into `createProfiler`, remove `Wrapper` prop

---
 src/index.ts                      |  6 +++++-
 src/profile/profile.tsx           | 32 +++++++++++++++++++++++++------
 src/renderHookToSnapshotStream.ts | 20 +++----------------
 src/renderToRenderStream.ts       | 25 ++++--------------------
 4 files changed, 38 insertions(+), 45 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index a7dc54a94..97b8cac55 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,4 +1,8 @@
-export type { NextRenderOptions, RenderStream } from "./profile/profile.js";
+export type {
+  NextRenderOptions,
+  RenderStream,
+  RenderStreamWithRenderFn,
+} from "./profile/profile.js";
 export {
   createProfiler,
   useTrackRenders,
diff --git a/src/profile/profile.tsx b/src/profile/profile.tsx
index 2d8ccfd08..a84d7f2b5 100644
--- a/src/profile/profile.tsx
+++ b/src/profile/profile.tsx
@@ -6,7 +6,7 @@ import { applyStackTrace, captureStackTrace } from "./traces.js";
 import type { ProfilerContextValue } from "./context.js";
 import { ProfilerContextProvider, useProfilerContext } from "./context.js";
 import { disableActWarnings } from "./disableActWarnings.js";
-import { render } from "@testing-library/react";
+import { render as baseRender, RenderOptions } from "@testing-library/react";
 
 export type ValidSnapshot =
   | void
@@ -80,9 +80,9 @@ export interface RenderStream<Snapshot extends ValidSnapshot>
   extends ProfiledComponentFields<Snapshot>,
     ProfiledComponentOnlyFields<Snapshot> {}
 
-export interface RenderStreamWithWrapper<Snapshot extends ValidSnapshot>
+export interface RenderStreamWithRenderFn<Snapshot extends ValidSnapshot>
   extends RenderStream<Snapshot> {
-  Wrapper: React.FC<{ children: React.ReactNode }>;
+  render: typeof baseRender;
 }
 
 export type ProfilerOptions<Snapshot extends ValidSnapshot> = {
@@ -108,7 +108,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
   snapshotDOM = false,
   initialSnapshot,
   skipNonTrackingRenders,
-}: ProfilerOptions<Snapshot> = {}): RenderStreamWithWrapper<Snapshot> {
+}: ProfilerOptions<Snapshot> = {}): RenderStreamWithRenderFn<Snapshot> {
   let nextRender: Promise<Render<Snapshot>> | undefined;
   let resolveNextRender: ((render: Render<Snapshot>) => void) | undefined;
   let rejectNextRender: ((error: unknown) => void) | undefined;
@@ -221,7 +221,27 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
     );
   }
 
-  const Profiler: RenderStreamWithWrapper<Snapshot> = Object.assign(
+  const render = ((
+    ui: React.ReactNode,
+    options?: RenderOptions<any, any, any>
+  ) => {
+    return baseRender(ui, {
+      ...options,
+      wrapper: (props) => {
+        let elem: React.ReactNode = React.createElement(
+          Wrapper,
+          undefined,
+          props.children
+        );
+        if (options?.wrapper) {
+          elem = React.createElement(options.wrapper, undefined, elem);
+        }
+        return elem;
+      },
+    });
+  }) as typeof baseRender;
+
+  const Profiler: RenderStreamWithRenderFn<Snapshot> = Object.assign(
     {
       replaceSnapshot,
       mergeSnapshot,
@@ -318,7 +338,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
         return nextRender;
       },
     } satisfies ProfiledComponentFields<Snapshot>,
-    { Wrapper }
+    { render }
   );
   return Profiler;
 }
diff --git a/src/renderHookToSnapshotStream.ts b/src/renderHookToSnapshotStream.ts
index 8864c3c60..65123e598 100644
--- a/src/renderHookToSnapshotStream.ts
+++ b/src/renderHookToSnapshotStream.ts
@@ -1,8 +1,7 @@
-import { render, RenderHookOptions } from "@testing-library/react";
+import { RenderHookOptions } from "@testing-library/react";
 import {
   createProfiler,
   ProfiledComponentFields,
-  RenderStream,
   ValidSnapshot,
 } from "./profile/profile.js";
 import { Render } from "./profile/Render.js";
@@ -47,7 +46,7 @@ export function renderHookToSnapshotStream<
     unmount: () => void;
   },
 ] {
-  const { Wrapper, ...stream } = createProfiler<ReturnValue>();
+  const { render, ...stream } = createProfiler<ReturnValue>();
 
   const ProfiledHook: React.FC<Props> = (props) => {
     stream.replaceSnapshot(renderCallback(props));
@@ -56,20 +55,7 @@ export function renderHookToSnapshotStream<
 
   const { rerender: baseRerender, unmount } = render(
     createElement(ProfiledHook, initialProps),
-    {
-      ...options,
-      wrapper(props) {
-        let elem: React.ReactNode = createElement(
-          Wrapper,
-          undefined,
-          props.children
-        );
-        if (options.wrapper) {
-          elem = createElement(options.wrapper, undefined, elem);
-        }
-        return elem;
-      },
-    }
+    options
   );
 
   function rerender(rerenderCallbackProps: Props) {
diff --git a/src/renderToRenderStream.ts b/src/renderToRenderStream.ts
index 359f237a0..268c11042 100644
--- a/src/renderToRenderStream.ts
+++ b/src/renderToRenderStream.ts
@@ -1,5 +1,4 @@
 import {
-  render,
   type RenderOptions as BaseOptions,
   type RenderResult as BaseResult,
 } from "@testing-library/react";
@@ -10,7 +9,6 @@ import {
   ProfilerOptions,
   ValidSnapshot,
 } from "./profile/profile.js";
-import { createElement } from "react";
 
 type RenderOptions<Snapshot extends ValidSnapshot = void> = BaseOptions &
   ProfilerOptions<Snapshot>;
@@ -18,7 +16,7 @@ type RenderOptions<Snapshot extends ValidSnapshot = void> = BaseOptions &
 type RenderResult<Snapshot extends ValidSnapshot = void> = [
   Stream: ProfiledComponentFields<Snapshot> &
     ProfiledComponentOnlyFields<Snapshot>,
-  resultPromise: Promise<BaseResult>,
+  renderResultPromise: Promise<BaseResult>,
 ];
 
 /**
@@ -35,27 +33,12 @@ export function renderToRenderStream<Snapshot extends ValidSnapshot = void>(
     ...options
   }: RenderOptions<Snapshot> = {}
 ): RenderResult<Snapshot> {
-  const { Wrapper, ...stream } = createProfiler<Snapshot>({
+  const { render, ...stream } = createProfiler<Snapshot>({
     onRender,
     snapshotDOM,
     initialSnapshot,
     skipNonTrackingRenders,
   });
-  const result = Promise.resolve().then(() =>
-    render(ui, {
-      ...options,
-      wrapper(props) {
-        let elem: React.ReactNode = createElement(
-          Wrapper,
-          undefined,
-          props.children
-        );
-        if (options.wrapper) {
-          elem = createElement(options.wrapper, undefined, elem);
-        }
-        return elem;
-      },
-    })
-  );
-  return [stream, result];
+  const renderResultPromise = Promise.resolve().then(() => render(ui, options));
+  return [stream, renderResultPromise];
 }