fix(dialog): restoring focus too early

src/lib/dialog/dialog-container.ts

@@ -109,9 +109,7 @@ export class MdDialogContainer extends BasePortalHost {
     return this._portalHost.attachTemplatePortal(portal);
   }
 
-  /**
-   * Moves the focus inside the focus trap.
-   */
+  /** Moves the focus inside the focus trap. */
   private _trapFocus() {
     if (!this._focusTrap) {
       this._focusTrap = this._focusTrapFactory.create(this._elementRef.nativeElement);
@@ -123,45 +121,36 @@ export class MdDialogContainer extends BasePortalHost {
     this._focusTrap.focusFirstTabbableElementWhenReady();
   }
 
-  /**
-   * Saves a reference to the element that was focused before the dialog was opened.
-   */
+  /** Restores focus to the element that was focused before the dialog opened. */
+  private _restoreFocus() {
+    const toFocus = this._elementFocusedBeforeDialogWasOpened;
+
+    // We need the extra check, because IE can set the `activeElement` to null in some cases.
+    if (toFocus && 'focus' in toFocus) {
+      toFocus.focus();
+    }
+
+    if (this._focusTrap) {
+      this._focusTrap.destroy();
+    }
+  }
+
+  /** Saves a reference to the element that was focused before the dialog was opened. */
   private _savePreviouslyFocusedElement() {
     if (this._document) {
       this._elementFocusedBeforeDialogWasOpened = this._document.activeElement as HTMLElement;
     }
   }
 
-  /**
-   * Callback, invoked whenever an animation on the host completes.
-   * @docs-private
-   */
+  /** Callback, invoked whenever an animation on the host completes. */
   _onAnimationDone(event: AnimationEvent) {
     this._onAnimationStateChange.emit(event);
 
     if (event.toState === 'enter') {
       this._trapFocus();
     } else if (event.toState === 'exit') {
+      this._restoreFocus();
       this._onAnimationStateChange.complete();
     }
   }
-
-  /**
-   * Kicks off the leave animation and restores focus to the previously-focused element.
-   * @docs-private
-   */
-  _exit(): void {
-    // We need the extra check, because IE can set the `activeElement` to null in some cases.
-    let toFocus = this._elementFocusedBeforeDialogWasOpened;
-
-    if (toFocus && 'focus' in toFocus) {
-      toFocus.focus();
-    }
-
-    if (this._focusTrap) {
-      this._focusTrap.destroy();
-    }
-
-    this._state = 'exit';
-  }
 }

src/lib/dialog/dialog-ref.ts

@@ -42,7 +42,7 @@ export class MdDialogRef<T> {
    */
   close(dialogResult?: any): void {
     this._result = dialogResult;
-    this._containerInstance._exit();
+    this._containerInstance._state = 'exit';
     this._overlayRef.detachBackdrop(); // Transition the backdrop in parallel with the dialog.
   }
 

src/lib/dialog/dialog.spec.ts

@@ -480,15 +480,9 @@ describe('MdDialog', () => {
   });
 
   describe('focus management', () => {
-
     // When testing focus, all of the elements must be in the DOM.
-    beforeEach(() => {
-      document.body.appendChild(overlayContainerElement);
-    });
-
-    afterEach(() => {
-      document.body.removeChild(overlayContainerElement);
-    });
+    beforeEach(() => document.body.appendChild(overlayContainerElement));
+    afterEach(() => document.body.removeChild(overlayContainerElement));
 
     it('should focus the first tabbable element of the dialog on open', fakeAsync(() => {
       dialog.open(PizzaMsg, {
@@ -515,17 +509,19 @@ describe('MdDialog', () => {
 
       viewContainerFixture.detectChanges();
       flushMicrotasks();
-
       expect(document.activeElement.id)
           .not.toBe('dialog-trigger', 'Expected the focus to change when dialog was opened.');
 
       dialogRef.close();
+      expect(document.activeElement.id).not.toBe('dialog-trigger',
+          'Expcted the focus not to have changed before the animation finishes.');
+
       tick(500);
       viewContainerFixture.detectChanges();
       flushMicrotasks();
 
-      expect(document.activeElement.id)
-          .toBe('dialog-trigger', 'Expected that the trigger was refocused after dialog close');
+      expect(document.activeElement.id).toBe('dialog-trigger',
+          'Expected that the trigger was refocused after the dialog is closed.');
 
       document.body.removeChild(button);
     }));
@@ -552,7 +548,7 @@ describe('MdDialog', () => {
       flushMicrotasks();
 
       expect(document.activeElement.id).toBe('input-to-be-focused',
-          'Expected that the trigger was refocused after dialog close');
+          'Expected that the trigger was refocused after the dialog is closed.');
 
       document.body.removeChild(button);
       document.body.removeChild(input);