@@ -6,7 +6,7 @@ import { applyStackTrace, captureStackTrace } from "./traces.js";
66import type { ProfilerContextValue } from "./context.js" ;
77import { ProfilerContextProvider , useProfilerContext } from "./context.js" ;
88import { disableActWarnings } from "./disableActWarnings.js" ;
9- import { render } from "@testing-library/react" ;
9+ import { render as baseRender , RenderOptions } from "@testing-library/react" ;
1010
1111export type ValidSnapshot =
1212 | void
@@ -80,9 +80,9 @@ export interface RenderStream<Snapshot extends ValidSnapshot>
8080 extends ProfiledComponentFields < Snapshot > ,
8181 ProfiledComponentOnlyFields < Snapshot > { }
8282
83- export interface RenderStreamWithWrapper < Snapshot extends ValidSnapshot >
83+ export interface RenderStreamWithRenderFn < Snapshot extends ValidSnapshot >
8484 extends RenderStream < Snapshot > {
85- Wrapper : React . FC < { children : React . ReactNode } > ;
85+ render : typeof baseRender ;
8686}
8787
8888export type ProfilerOptions < Snapshot extends ValidSnapshot > = {
@@ -108,7 +108,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
108108 snapshotDOM = false ,
109109 initialSnapshot,
110110 skipNonTrackingRenders,
111- } : ProfilerOptions < Snapshot > = { } ) : RenderStreamWithWrapper < Snapshot > {
111+ } : ProfilerOptions < Snapshot > = { } ) : RenderStreamWithRenderFn < Snapshot > {
112112 let nextRender : Promise < Render < Snapshot > > | undefined ;
113113 let resolveNextRender : ( ( render : Render < Snapshot > ) => void ) | undefined ;
114114 let rejectNextRender : ( ( error : unknown ) => void ) | undefined ;
@@ -221,7 +221,27 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
221221 ) ;
222222 }
223223
224- const Profiler : RenderStreamWithWrapper < Snapshot > = Object . assign (
224+ const render = ( (
225+ ui : React . ReactNode ,
226+ options ?: RenderOptions < any , any , any >
227+ ) => {
228+ return baseRender ( ui , {
229+ ...options ,
230+ wrapper : ( props ) => {
231+ let elem : React . ReactNode = React . createElement (
232+ Wrapper ,
233+ undefined ,
234+ props . children
235+ ) ;
236+ if ( options ?. wrapper ) {
237+ elem = React . createElement ( options . wrapper , undefined , elem ) ;
238+ }
239+ return elem ;
240+ } ,
241+ } ) ;
242+ } ) as typeof baseRender ;
243+
244+ const Profiler : RenderStreamWithRenderFn < Snapshot > = Object . assign (
225245 {
226246 replaceSnapshot,
227247 mergeSnapshot,
@@ -318,7 +338,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
318338 return nextRender ;
319339 } ,
320340 } satisfies ProfiledComponentFields < Snapshot > ,
321- { Wrapper }
341+ { render }
322342 ) ;
323343 return Profiler ;
324344}
0 commit comments