From bfa185305d28fdedd2705b22d5318d0b29299df9 Mon Sep 17 00:00:00 2001
From: Kristiyan Kostadinov <crisbeto@users.noreply.github.com>
Date: Mon, 9 Sep 2019 16:12:49 +0200
Subject: [PATCH] fix(menu): restore focus immediately when menu is closed
 (#16960)

Currently in some cases we restore focus immediately and in others we wait for the exit animation to finish (when using `matMenuContent`) because the focus restoration logic is coupled to the menu cleanup. This means that if multiple animations overlap, we could end up restoring focus to the wrong element. These changes switch to restoring focus immediately in all cases.

Fixes #16954.
---
 .../mdc-menu/menu.spec.ts                     | 19 +++++++++++++++++++
 src/material/menu/menu-trigger.ts             | 17 +++++++----------
 src/material/menu/menu.spec.ts                | 19 +++++++++++++++++++
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/src/material-experimental/mdc-menu/menu.spec.ts b/src/material-experimental/mdc-menu/menu.spec.ts
index ec2ca52cc503..70fbd4828e2c 100644
--- a/src/material-experimental/mdc-menu/menu.spec.ts
+++ b/src/material-experimental/mdc-menu/menu.spec.ts
@@ -179,6 +179,25 @@ describe('MatMenu', () => {
     expect(document.activeElement).not.toBe(triggerEl);
   }));
 
+  it('should restore focus to the trigger immediately once the menu is closed', () => {
+    const fixture = createComponent(SimpleMenu, [], [FakeIcon]);
+    fixture.detectChanges();
+    const triggerEl = fixture.componentInstance.triggerEl.nativeElement;
+
+    // A click without a mousedown before it is considered a keyboard open.
+    triggerEl.click();
+    fixture.detectChanges();
+
+    expect(overlayContainerElement.querySelector('.mat-mdc-menu-panel')).toBeTruthy();
+
+    fixture.componentInstance.trigger.closeMenu();
+    fixture.detectChanges();
+    // Note: don't add a `tick` here since we're testing
+    // that focus is restored before the animation is done.
+
+    expect(document.activeElement).toBe(triggerEl);
+  });
+
   it('should be able to set a custom class on the backdrop', fakeAsync(() => {
     const fixture = createComponent(SimpleMenu, [], [FakeIcon]);
 
diff --git a/src/material/menu/menu-trigger.ts b/src/material/menu/menu-trigger.ts
index 8d73e70828f9..e40dc3a2b951 100644
--- a/src/material/menu/menu-trigger.ts
+++ b/src/material/menu/menu-trigger.ts
@@ -298,18 +298,20 @@ export class MatMenuTrigger implements AfterContentInit, OnDestroy {
           .subscribe({
             next: () => menu.lazyContent!.detach(),
             // No matter whether the content got re-attached, reset the menu.
-            complete: () => this._resetMenu()
+            complete: () => this._setIsMenuOpen(false)
           });
       } else {
-        this._resetMenu();
+        this._setIsMenuOpen(false);
       }
     } else {
-      this._resetMenu();
+      this._setIsMenuOpen(false);
 
       if (menu.lazyContent) {
         menu.lazyContent.detach();
       }
     }
+
+    this._restoreFocus();
   }
 
   /**
@@ -339,13 +341,8 @@ export class MatMenuTrigger implements AfterContentInit, OnDestroy {
     }
   }
 
-  /**
-   * This method resets the menu when it's closed, most importantly restoring
-   * focus to the menu trigger if the menu was opened via the keyboard.
-   */
-  private _resetMenu(): void {
-    this._setIsMenuOpen(false);
-
+  /** Restores focus to the element that was focused before the menu was open. */
+  private _restoreFocus() {
     // We should reset focus if the user is navigating using a keyboard or
     // if we have a top-level trigger which might cause focus to be lost
     // when clicking on the backdrop.
diff --git a/src/material/menu/menu.spec.ts b/src/material/menu/menu.spec.ts
index 4c975790508e..67c4bb1908dd 100644
--- a/src/material/menu/menu.spec.ts
+++ b/src/material/menu/menu.spec.ts
@@ -181,6 +181,25 @@ describe('MatMenu', () => {
     expect(document.activeElement).not.toBe(triggerEl);
   }));
 
+  it('should restore focus to the trigger immediately once the menu is closed', () => {
+    const fixture = createComponent(SimpleMenu, [], [FakeIcon]);
+    fixture.detectChanges();
+    const triggerEl = fixture.componentInstance.triggerEl.nativeElement;
+
+    // A click without a mousedown before it is considered a keyboard open.
+    triggerEl.click();
+    fixture.detectChanges();
+
+    expect(overlayContainerElement.querySelector('.mat-menu-panel')).toBeTruthy();
+
+    fixture.componentInstance.trigger.closeMenu();
+    fixture.detectChanges();
+    // Note: don't add a `tick` here since we're testing
+    // that focus is restored before the animation is done.
+
+    expect(document.activeElement).toBe(triggerEl);
+  });
+
   it('should be able to set a custom class on the backdrop', fakeAsync(() => {
     const fixture = createComponent(SimpleMenu, [], [FakeIcon]);