@@ -400,7 +400,29 @@ function createWrapper(
400400 return { wrapperState, wrapper } ;
401401}
402402
403- export type RenderConnectionsOptions = RenderOptions & TestConnectionsOptions ;
403+ /**
404+ * Returns a new {@link RenderResult} with the {@link RenderResult.container} replaced by the container inserted by the context menu provider.
405+ */
406+ function unwrapContextMenuContainer ( result : RenderResult ) {
407+ const { container, ...rest } = result ;
408+ const { firstChild } = container ;
409+ if (
410+ firstChild instanceof HTMLElement &&
411+ firstChild . getAttribute ( 'data-testid' ) === 'context-menu-children-container'
412+ ) {
413+ return { container : firstChild , ...rest } ;
414+ } else {
415+ return { container, ...rest } ;
416+ }
417+ }
418+
419+ export type RenderConnectionsOptions = RenderOptions &
420+ TestConnectionsOptions & {
421+ /**
422+ * Whether to include the context menu container and menu in the container of the returned result.
423+ */
424+ includeContextMenu ?: boolean ;
425+ } ;
404426
405427export type RenderWithConnectionsResult = ReturnType <
406428 typeof createWrapper
@@ -415,6 +437,7 @@ function renderWithConnections(
415437 baseElement,
416438 queries,
417439 hydrate,
440+ includeContextMenu = false ,
418441 ...connectionsOptions
419442 } : RenderConnectionsOptions = { }
420443) : RenderWithConnectionsResult {
@@ -443,7 +466,10 @@ function renderWithConnections(
443466 true ,
444467 'Expected initial connections to load before rendering rest of the tested UI, but it did not happen'
445468 ) ;
446- return { ...wrapperState , ...result } ;
469+ return {
470+ ...wrapperState ,
471+ ...( includeContextMenu ? result : unwrapContextMenuContainer ( result ) ) ,
472+ } ;
447473}
448474
449475export type RenderHookConnectionsOptions < HookProps > = Omit <
0 commit comments