@@ -6,7 +6,7 @@ import { applyStackTrace, captureStackTrace } from "./traces.js";
6
6
import type { ProfilerContextValue } from "./context.js" ;
7
7
import { ProfilerContextProvider , useProfilerContext } from "./context.js" ;
8
8
import { disableActWarnings } from "./disableActWarnings.js" ;
9
- import { render } from "@testing-library/react" ;
9
+ import { render as baseRender , RenderOptions } from "@testing-library/react" ;
10
10
11
11
export type ValidSnapshot =
12
12
| void
@@ -80,9 +80,9 @@ export interface RenderStream<Snapshot extends ValidSnapshot>
80
80
extends ProfiledComponentFields < Snapshot > ,
81
81
ProfiledComponentOnlyFields < Snapshot > { }
82
82
83
- export interface RenderStreamWithWrapper < Snapshot extends ValidSnapshot >
83
+ export interface RenderStreamWithRenderFn < Snapshot extends ValidSnapshot >
84
84
extends RenderStream < Snapshot > {
85
- Wrapper : React . FC < { children : React . ReactNode } > ;
85
+ render : typeof baseRender ;
86
86
}
87
87
88
88
export type ProfilerOptions < Snapshot extends ValidSnapshot > = {
@@ -108,7 +108,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
108
108
snapshotDOM = false ,
109
109
initialSnapshot,
110
110
skipNonTrackingRenders,
111
- } : ProfilerOptions < Snapshot > = { } ) : RenderStreamWithWrapper < Snapshot > {
111
+ } : ProfilerOptions < Snapshot > = { } ) : RenderStreamWithRenderFn < Snapshot > {
112
112
let nextRender : Promise < Render < Snapshot > > | undefined ;
113
113
let resolveNextRender : ( ( render : Render < Snapshot > ) => void ) | undefined ;
114
114
let rejectNextRender : ( ( error : unknown ) => void ) | undefined ;
@@ -221,7 +221,27 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
221
221
) ;
222
222
}
223
223
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 (
225
245
{
226
246
replaceSnapshot,
227
247
mergeSnapshot,
@@ -318,7 +338,7 @@ export function createProfiler<Snapshot extends ValidSnapshot = void>({
318
338
return nextRender ;
319
339
} ,
320
340
} satisfies ProfiledComponentFields < Snapshot > ,
321
- { Wrapper }
341
+ { render }
322
342
) ;
323
343
return Profiler ;
324
344
}
0 commit comments