Prevent useContextMenu isOpen from being true if the button ref goes away

src/components/structures/ContextMenu.tsx

@@ -565,8 +565,13 @@ type ContextMenuTuple<T> = [
     (val: boolean) => void,
 ];
 // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-constraint
-export const useContextMenu = <T extends any = HTMLElement>(): ContextMenuTuple<T> => {
-    const button = useRef<T>(null);
+export const useContextMenu = <T extends any = HTMLElement>(inputRef?: RefObject<T>): ContextMenuTuple<T> => {
+    let button = useRef<T>(null);
+    if (inputRef) {
+        // if we are given a ref, use it instead of ours
+        button = inputRef;
+    }
+
     const [isOpen, setIsOpen] = useState(false);
     const open = (ev?: SyntheticEvent) => {
         ev?.preventDefault();
@@ -579,7 +584,7 @@ export const useContextMenu = <T extends any = HTMLElement>(): ContextMenuTuple<
         setIsOpen(false);
     };
 
-    return [isOpen, button, open, close, setIsOpen];
+    return [button.current ? isOpen : false, button, open, close, setIsOpen];
 };
 
 // XXX: Deprecated, used only for dynamic Tooltips. Avoid using at all costs.

src/components/views/location/LocationButton.tsx

@@ -69,6 +69,7 @@ export const LocationButton: React.FC<IProps> = ({ roomId, sender, menuPosition,
             iconClassName="mx_MessageComposer_location"
             onClick={openMenu}
             title={_t("Location")}
+            inputRef={button}
         />
 
         { contextMenu }

src/components/views/rooms/ReadReceiptGroup.tsx

@@ -209,7 +209,8 @@ export function ReadReceiptGroup(
                     onMouseOver={showTooltip}
                     onMouseLeave={hideTooltip}
                     onFocus={showTooltip}
-                    onBlur={hideTooltip}>
+                    onBlur={hideTooltip}
+                >
                     { remText }
                     <span
                         className="mx_ReadReceiptGroup_container"

src/components/views/spaces/SpacePanel.tsx

@@ -204,9 +204,7 @@ const CreateSpaceButton = ({
     isPanelCollapsed,
     setPanelCollapsed,
 }: Pick<IInnerSpacePanelProps, "isPanelCollapsed" | "setPanelCollapsed">) => {
-    // We don't need the handle as we position the menu in a constant location
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    const [menuDisplayed, _handle, openMenu, closeMenu] = useContextMenu<void>();
+    const [menuDisplayed, handle, openMenu, closeMenu] = useContextMenu<HTMLElement>();
 
     useEffect(() => {
         if (!isPanelCollapsed && menuDisplayed) {
@@ -231,13 +229,14 @@ const CreateSpaceButton = ({
         role="treeitem"
     >
         <SpaceButton
-            data-test-id='create-space-button'
+            data-testid='create-space-button'
             className={classNames("mx_SpaceButton_new", {
                 mx_SpaceButton_newCancel: menuDisplayed,
             })}
             label={menuDisplayed ? _t("Cancel") : _t("Create a space")}
             onClick={onNewClick}
             isNarrow={isPanelCollapsed}
+            ref={handle}
         />
 
         { contextMenu }

src/components/views/spaces/SpaceTreeLevel.tsx

@@ -14,7 +14,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-import React, { MouseEvent, ComponentProps, ComponentType, createRef, InputHTMLAttributes, LegacyRef } from "react";
+import React, {
+    MouseEvent,
+    ComponentProps,
+    ComponentType,
+    createRef,
+    InputHTMLAttributes,
+    LegacyRef,
+    forwardRef,
+    RefObject,
+} from "react";
 import classNames from "classnames";
 import { Room, RoomEvent } from "matrix-js-sdk/src/models/room";
 import { DraggableProvidedDragHandleProps } from "react-beautiful-dnd";
@@ -54,7 +63,7 @@ interface IButtonProps extends Omit<ComponentProps<typeof AccessibleTooltipButto
     onClick?(ev?: ButtonEvent): void;
 }
 
-export const SpaceButton: React.FC<IButtonProps> = ({
+export const SpaceButton = forwardRef<HTMLElement, IButtonProps>(({
     space,
     spaceKey,
     className,
@@ -67,9 +76,9 @@ export const SpaceButton: React.FC<IButtonProps> = ({
     children,
     ContextMenuComponent,
     ...props
-}) => {
-    const [menuDisplayed, ref, openMenu, closeMenu] = useContextMenu<HTMLElement>();
-    const [onFocus, isActive, handle] = useRovingTabIndex(ref);
+}, ref: RefObject<HTMLElement>) => {
+    const [menuDisplayed, handle, openMenu, closeMenu] = useContextMenu<HTMLElement>(ref);
+    const [onFocus, isActive] = useRovingTabIndex(handle);
     const tabIndex = isActive ? 0 : -1;
 
     let avatar = <div className="mx_SpaceButton_avatarPlaceholder"><div className="mx_SpaceButton_icon" /></div>;
@@ -150,7 +159,7 @@ export const SpaceButton: React.FC<IButtonProps> = ({
             </div>
         </AccessibleTooltipButton>
     );
-};
+});
 
 interface IItemProps extends InputHTMLAttributes<HTMLLIElement> {
     space: Room;

test/components/views/spaces/SpacePanel-test.tsx

@@ -15,16 +15,13 @@ limitations under the License.
 */
 
 import React from 'react';
-// eslint-disable-next-line deprecate/import
-import { mount } from 'enzyme';
+import { render, screen, fireEvent } from "@testing-library/react";
 import { mocked } from 'jest-mock';
 import { MatrixClient } from 'matrix-js-sdk/src/matrix';
-import { act } from "react-dom/test-utils";
 
 import SpacePanel from '../../../../src/components/views/spaces/SpacePanel';
 import { MatrixClientPeg } from '../../../../src/MatrixClientPeg';
 import { SpaceKey } from '../../../../src/stores/spaces';
-import { findByTestId } from '../../../test-utils';
 import { shouldShowComponent } from '../../../../src/customisations/helpers/UIComponents';
 import { UIComponent } from '../../../../src/settings/UIFeature';
 
@@ -47,10 +44,6 @@ jest.mock('../../../../src/customisations/helpers/UIComponents', () => ({
 }));
 
 describe('<SpacePanel />', () => {
-    const defaultProps = {};
-    const getComponent = (props = {}) =>
-        mount(<SpacePanel {...defaultProps} {...props} />);
-
     const mockClient = {
         getUserId: jest.fn().mockReturnValue('@test:test'),
         isGuest: jest.fn(),
@@ -67,26 +60,21 @@ describe('<SpacePanel />', () => {
 
     describe('create new space button', () => {
         it('renders create space button when UIComponent.CreateSpaces component should be shown', () => {
-            const component = getComponent();
-            expect(findByTestId(component, 'create-space-button').length).toBeTruthy();
+            render(<SpacePanel />);
+            screen.getByTestId("create-space-button");
         });
 
         it('does not render create space button when UIComponent.CreateSpaces component should not be shown', () => {
             mocked(shouldShowComponent).mockReturnValue(false);
-            const component = getComponent();
+            render(<SpacePanel />);
             expect(shouldShowComponent).toHaveBeenCalledWith(UIComponent.CreateSpaces);
-            expect(findByTestId(component, 'create-space-button').length).toBeFalsy();
+            expect(screen.queryByTestId("create-space-button")).toBeFalsy();
         });
 
-        it('opens context menu on create space button click', async () => {
-            const component = getComponent();
-
-            await act(async () => {
-                findByTestId(component, 'create-space-button').at(0).simulate('click');
-                component.setProps({});
-            });
-
-            expect(component.find('SpaceCreateMenu').length).toBeTruthy();
+        it('opens context menu on create space button click', () => {
+            render(<SpacePanel />);
+            fireEvent.click(screen.getByTestId("create-space-button"));
+            screen.getByTestId("create-space-button");
         });
     });
 });