From ee1a5a7aa3c31c00c4d5bcd3285a1302e7bd5a43 Mon Sep 17 00:00:00 2001
From: Kristiyan Kostadinov <crisbeto@users.noreply.github.com>
Date: Tue, 18 Apr 2017 00:16:02 +0200
Subject: [PATCH] fix(dialog): fire afterClosed callback after all dialog
 actions are done (#3892)

* Fires the afterClosed callback once the dialog is removed and the previously-focused element has been refocused. Until now the order was reversed, which prevents people from being able to refocus a different element.
* Cleaned up the `MdDialogContainer` logic to simplify it and to remove the need to subscribe to zone events.
---
 src/demo-app/dialog/dialog-demo.html |  2 +-
 src/lib/dialog/dialog-container.ts   | 58 +++++++++-------------------
 src/lib/dialog/dialog-ref.ts         | 24 ++++++------
 src/lib/dialog/dialog.spec.ts        | 51 +++++++++++++-----------
 4 files changed, 61 insertions(+), 74 deletions(-)

diff --git a/src/demo-app/dialog/dialog-demo.html b/src/demo-app/dialog/dialog-demo.html
index 45e7ba5925c..8d3f4890819 100644
--- a/src/demo-app/dialog/dialog-demo.html
+++ b/src/demo-app/dialog/dialog-demo.html
@@ -1,6 +1,6 @@
 <h1>Dialog demo</h1>
 
-<button md-raised-button color="primary" (click)="openJazz()" [disabled]="dialogRef">
+<button md-raised-button color="primary" (click)="openJazz()">
   Open dialog
 </button>
 <button md-raised-button color="accent" (click)="openContentElement()">
diff --git a/src/lib/dialog/dialog-container.ts b/src/lib/dialog/dialog-container.ts
index 1416d36a7bb..9e0c95aeefd 100644
--- a/src/lib/dialog/dialog-container.ts
+++ b/src/lib/dialog/dialog-container.ts
@@ -3,8 +3,6 @@ import {
   ComponentRef,
   ViewChild,
   ViewEncapsulation,
-  NgZone,
-  OnDestroy,
   Renderer,
   ElementRef,
   EventEmitter,
@@ -21,11 +19,6 @@ import {BasePortalHost, ComponentPortal, PortalHostDirective, TemplatePortal} fr
 import {MdDialogConfig} from './dialog-config';
 import {MdDialogContentAlreadyAttachedError} from './dialog-errors';
 import {FocusTrapFactory, FocusTrap} from '../core/a11y/focus-trap';
-import 'rxjs/add/operator/first';
-
-
-/** Possible states for the dialog container animation. */
-export type MdDialogContainerAnimationState = 'void' | 'enter' | 'exit' | 'exit-start';
 
 
 /**
@@ -54,7 +47,7 @@ export type MdDialogContainerAnimationState = 'void' | 'enter' | 'exit' | 'exit-
     '(@slideDialog.done)': '_onAnimationDone($event)',
   },
 })
-export class MdDialogContainer extends BasePortalHost implements OnDestroy {
+export class MdDialogContainer extends BasePortalHost {
   /** The portal host inside of this container into which the dialog content will be loaded. */
   @ViewChild(PortalHostDirective) _portalHost: PortalHostDirective;
 
@@ -68,13 +61,12 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
   dialogConfig: MdDialogConfig;
 
   /** State of the dialog animation. */
-  _state: MdDialogContainerAnimationState = 'enter';
+  _state: 'void' | 'enter' | 'exit' = 'enter';
 
   /** Emits the current animation state whenever it changes. */
-  _onAnimationStateChange = new EventEmitter<MdDialogContainerAnimationState>();
+  _onAnimationStateChange = new EventEmitter<AnimationEvent>();
 
   constructor(
-    private _ngZone: NgZone,
     private _renderer: Renderer,
     private _elementRef: ElementRef,
     private _focusTrapFactory: FocusTrapFactory) {
@@ -108,7 +100,6 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
 
   /**
    * Moves the focus inside the focus trap.
-   * @private
    */
   private _trapFocus() {
     if (!this._focusTrap) {
@@ -122,47 +113,36 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
     this._focusTrap.focusFirstTabbableElementWhenReady();
   }
 
-  /**
-   * Kicks off the leave animation.
-   * @docs-private
-   */
-  _exit(): void {
-    this._state = 'exit';
-    this._onAnimationStateChange.emit('exit-start');
-  }
-
   /**
    * Callback, invoked whenever an animation on the host completes.
    * @docs-private
    */
   _onAnimationDone(event: AnimationEvent) {
+    this._onAnimationStateChange.emit(event);
+
     if (event.toState === 'enter') {
       this._trapFocus();
+    } else if (event.toState === 'exit') {
+      this._onAnimationStateChange.complete();
     }
-
-    this._onAnimationStateChange.emit(event.toState as MdDialogContainerAnimationState);
   }
 
-  ngOnDestroy() {
-    // When the dialog is destroyed, return focus to the element that originally had it before
-    // the dialog was opened. Wait for the DOM to finish settling before changing the focus so
-    // that it doesn't end up back on the <body>. Also note that we need the extra check, because
-    // IE can set the `activeElement` to null in some cases.
-    let toFocus = this._elementFocusedBeforeDialogWasOpened as HTMLElement;
+  /**
+   * 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;
 
-    // We shouldn't use `this` inside of the NgZone subscription, because it causes a memory leak.
-    let animationStream = this._onAnimationStateChange;
-
-    this._ngZone.onMicrotaskEmpty.first().subscribe(() => {
-      if (toFocus && 'focus' in toFocus) {
-        toFocus.focus();
-      }
-
-      animationStream.complete();
-    });
+    if (toFocus && 'focus' in toFocus) {
+      toFocus.focus();
+    }
 
     if (this._focusTrap) {
       this._focusTrap.destroy();
     }
+
+    this._state = 'exit';
   }
 }
diff --git a/src/lib/dialog/dialog-ref.ts b/src/lib/dialog/dialog-ref.ts
index 429dc914c38..c7fbcdbdb14 100644
--- a/src/lib/dialog/dialog-ref.ts
+++ b/src/lib/dialog/dialog-ref.ts
@@ -1,8 +1,10 @@
 import {OverlayRef, GlobalPositionStrategy} from '../core';
+import {AnimationEvent} from '@angular/animations';
 import {DialogPosition} from './dialog-config';
 import {Observable} from 'rxjs/Observable';
 import {Subject} from 'rxjs/Subject';
-import {MdDialogContainer, MdDialogContainerAnimationState} from './dialog-container';
+import {MdDialogContainer} from './dialog-container';
+import 'rxjs/add/operator/filter';
 
 
 // TODO(jelbourn): resizing
@@ -23,17 +25,14 @@ export class MdDialogRef<T> {
   private _result: any;
 
   constructor(private _overlayRef: OverlayRef, public _containerInstance: MdDialogContainer) {
-    _containerInstance._onAnimationStateChange.subscribe(
-      (state: MdDialogContainerAnimationState) => {
-        if (state === 'exit-start') {
-          // Transition the backdrop in parallel with the dialog.
-          this._overlayRef.detachBackdrop();
-        } else if (state === 'exit') {
-          this._overlayRef.dispose();
-          this._afterClosed.next(this._result);
-          this._afterClosed.complete();
-          this.componentInstance = null;
-        }
+    _containerInstance._onAnimationStateChange
+      .filter((event: AnimationEvent) => event.toState === 'exit')
+      .subscribe(() => {
+        this._overlayRef.dispose();
+        this.componentInstance = null;
+      }, null, () => {
+        this._afterClosed.next(this._result);
+        this._afterClosed.complete();
       });
   }
 
@@ -44,6 +43,7 @@ export class MdDialogRef<T> {
   close(dialogResult?: any): void {
     this._result = dialogResult;
     this._containerInstance._exit();
+    this._overlayRef.detachBackdrop(); // Transition the backdrop in parallel with the dialog.
   }
 
   /**
diff --git a/src/lib/dialog/dialog.spec.ts b/src/lib/dialog/dialog.spec.ts
index d9833fe3b53..187c37908e4 100644
--- a/src/lib/dialog/dialog.spec.ts
+++ b/src/lib/dialog/dialog.spec.ts
@@ -334,28 +334,6 @@ describe('MdDialog', () => {
     expect(dialogContainer._state).toBe('exit');
   });
 
-  it('should emit an event with the proper animation state', async(() => {
-    let dialogRef = dialog.open(PizzaMsg, { viewContainerRef: testViewContainerRef });
-    let dialogContainer: MdDialogContainer =
-        viewContainerFixture.debugElement.query(By.directive(MdDialogContainer)).componentInstance;
-    let spy = jasmine.createSpy('animation state callback');
-
-    dialogContainer._onAnimationStateChange.subscribe(spy);
-    viewContainerFixture.detectChanges();
-
-    viewContainerFixture.whenStable().then(() => {
-      expect(spy).toHaveBeenCalledWith('enter');
-
-      dialogRef.close();
-      viewContainerFixture.detectChanges();
-      expect(spy).toHaveBeenCalledWith('exit-start');
-
-      viewContainerFixture.whenStable().then(() => {
-        expect(spy).toHaveBeenCalledWith('exit');
-      });
-    });
-  }));
-
   describe('passing in data', () => {
     it('should be able to pass in data', () => {
       let config = {
@@ -469,6 +447,35 @@ describe('MdDialog', () => {
 
       document.body.removeChild(button);
     }));
+
+    it('should allow the consumer to shift focus in afterClosed', fakeAsync(() => {
+      // Create a element that has focus before the dialog is opened.
+      let button = document.createElement('button');
+      let input = document.createElement('input');
+
+      button.id = 'dialog-trigger';
+      input.id = 'input-to-be-focused';
+
+      document.body.appendChild(button);
+      document.body.appendChild(input);
+      button.focus();
+
+      let dialogRef = dialog.open(PizzaMsg, { viewContainerRef: testViewContainerRef });
+
+      dialogRef.afterClosed().subscribe(() => input.focus());
+
+      dialogRef.close();
+      tick(500);
+      viewContainerFixture.detectChanges();
+      flushMicrotasks();
+
+      expect(document.activeElement.id).toBe('input-to-be-focused',
+          'Expected that the trigger was refocused after dialog close');
+
+      document.body.removeChild(button);
+      document.body.removeChild(input);
+    }));
+
   });
 
   describe('dialog content elements', () => {