From 3f4050f7b4f5a2ad2c2420d74a1fbc135e6b5d3e Mon Sep 17 00:00:00 2001
From: Sammy Jelin <sjelin@gmail.com>
Date: Tue, 17 Jan 2017 17:44:46 -0800
Subject: [PATCH] chore(frameworks): allow frameworks to call `afterEach`.
 (#3906)

Follow up to https://github.com/angular/protractor/pull/3893.

Part of #3893
---
 lib/runner.ts | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/lib/runner.ts b/lib/runner.ts
index b680c0598..983e17338 100644
--- a/lib/runner.ts
+++ b/lib/runner.ts
@@ -33,6 +33,8 @@ export class Runner extends EventEmitter {
   driverprovider_: DriverProvider;
   o: any;
   plugins_: Plugins;
+  restartPromise: q.Promise<any>;
+  frameworkUsesAfterEach: boolean;
 
   constructor(config: Config) {
     super();
@@ -86,6 +88,26 @@ export class Runner extends EventEmitter {
     });
   }
 
+  /**
+   * Called after each test finishes.
+   *
+   * Responsible for `restartBrowserBetweenTests`
+   *
+   * @public
+   * @return {q.Promise} A promise that will resolve when the work here is done
+   */
+  afterEach(): q.Promise<void> {
+    let ret: q.Promise<void>;
+    this.frameworkUsesAfterEach = true;
+    if (this.config_.restartBrowserBetweenTests) {
+      // TODO(sjelin): remove the `|| q()` once `restart()` returns a promise
+      this.restartPromise = this.restartPromise || protractor.browser.restart() || q();
+      ret = this.restartPromise;
+      this.restartPromise = undefined;
+    }
+    return ret || q();
+  }
+
   /**
    * Grab driver provider based on type
    * @private
@@ -322,8 +344,12 @@ export class Runner extends EventEmitter {
           }
 
           if (this.config_.restartBrowserBetweenTests) {
+            // TODO(sjelin): replace with warnings once `afterEach` support is required
             let restartDriver = () => {
-              browser_.restart();
+              if (!this.frameworkUsesAfterEach) {
+                // TODO(sjelin): remove the `|| q()` once `restart()` returns a promise
+                this.restartPromise = browser_.restart() || q();
+              }
             };
             this.on('testPass', restartDriver);
             this.on('testFail', restartDriver);