From f20c680611c3101a6cbba33d7925ec871d46e582 Mon Sep 17 00:00:00 2001
From: Anton Dosov <anton.dosov@elastic.co>
Date: Wed, 7 Feb 2024 12:28:28 +0100
Subject: [PATCH] [Discover] fix using the data grid full screen mode causes
 the top nav buttons to disappear (#176324)

## Summary

fix https://github.com/elastic/kibana/issues/172716

The header_action_menu didn't reset the state correctly when it was
unmounted. This caused an state bug when remounting and as the result
https://github.com/elastic/kibana/issues/172716 the action menu didn't
render
---
 .../src/ui/header/header_action_menu.test.tsx | 20 +++++++++++++++++++
 .../src/ui/header/header_action_menu.tsx      | 11 +++++-----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.test.tsx b/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.test.tsx
index 9a7b4fa5a8b14..359439f3e4cdf 100644
--- a/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.test.tsx
+++ b/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.test.tsx
@@ -147,4 +147,24 @@ describe('HeaderActionMenu', () => {
     expect(unmounts.FOO).toHaveBeenCalledTimes(1);
     expect(unmounts.BAR).not.toHaveBeenCalled();
   });
+
+  it('calls mount point `unmount` when unmounts', () => {
+    const TestComponent = () => {
+      const mounter = useHeaderActionMenuMounter(menuMount$);
+      return <HeaderActionMenu mounter={mounter} />;
+    };
+    component = mount(<TestComponent />);
+
+    act(() => {
+      menuMount$.next(createMountPoint('FOO'));
+    });
+    refresh();
+
+    expect(Object.keys(unmounts)).toEqual(['FOO']);
+    expect(unmounts.FOO).not.toHaveBeenCalled();
+
+    component.unmount();
+
+    expect(unmounts.FOO).toHaveBeenCalledTimes(1);
+  });
 });
diff --git a/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.tsx b/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.tsx
index ec2a06266a0fe..91d2077b28ea9 100644
--- a/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.tsx
+++ b/packages/core/chrome/core-chrome-browser-internal/src/ui/header/header_action_menu.tsx
@@ -39,11 +39,6 @@ export const HeaderActionMenu: FC<HeaderActionMenuProps> = ({ mounter }) => {
   const unmountRef = useRef<UnmountCallback | null>(null);
 
   useLayoutEffect(() => {
-    if (unmountRef.current) {
-      unmountRef.current();
-      unmountRef.current = null;
-    }
-
     if (mounter.mount && elementRef.current) {
       try {
         unmountRef.current = mounter.mount(elementRef.current);
@@ -53,6 +48,12 @@ export const HeaderActionMenu: FC<HeaderActionMenuProps> = ({ mounter }) => {
         console.error(e);
       }
     }
+    return () => {
+      if (unmountRef.current) {
+        unmountRef.current();
+        unmountRef.current = null;
+      }
+    };
   }, [mounter]);
 
   return <div data-test-subj="headerAppActionMenu" ref={elementRef} />;