From c9989c97afdfeba6447b63101d7a8bc52b4bc38b Mon Sep 17 00:00:00 2001
From: Fumihiro Xue <xfumihiro@gmail.com>
Date: Tue, 19 Sep 2017 09:54:12 +0800
Subject: [PATCH] Add Async Test Environment APIs

---
 packages/jest-runner/src/run_test.js | 39 +++++++++++++++-------------
 types/Environment.js                 |  2 ++
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/packages/jest-runner/src/run_test.js b/packages/jest-runner/src/run_test.js
index a2110f157fa6..e407d2fbaeee 100644
--- a/packages/jest-runner/src/run_test.js
+++ b/packages/jest-runner/src/run_test.js
@@ -100,7 +100,11 @@ export default function runTest(
     mapCoverage: globalConfig.mapCoverage,
   });
   const start = Date.now();
-  return testFramework(globalConfig, config, environment, runtime, path)
+  return (environment.setup ? environment.setup() : Promise.resolve({}))
+    .then(testSetup => {
+      setGlobal(environment.global, 'setup', testSetup);
+      return testFramework(globalConfig, config, environment, runtime, path);
+    })
     .then((result: TestResult) => {
       const testCount =
         result.numPassingTests +
@@ -116,23 +120,22 @@ export default function runTest(
       return result;
     })
     .then(
-      result =>
-        Promise.resolve().then(() => {
-          environment.dispose();
-          if (globalConfig.logHeapUsage) {
-            if (global.gc) {
-              global.gc();
-            }
-            result.memoryUsage = process.memoryUsage().heapUsed;
+      async result => {
+        if (environment.teardown) await environment.teardown();
+        environment.dispose();
+        if (globalConfig.logHeapUsage) {
+          if (global.gc) {
+            global.gc();
           }
-
-          // Delay the resolution to allow log messages to be output.
-          return new Promise(resolve => setImmediate(() => resolve(result)));
-        }),
-      err =>
-        Promise.resolve().then(() => {
-          environment.dispose();
-          throw err;
-        }),
+          result.memoryUsage = process.memoryUsage().heapUsed;
+        }
+        // Delay the resolution to allow log messages to be output.
+        return new Promise(resolve => setImmediate(() => resolve(result)));
+      },
+      async err => {
+        if (environment.teardown) await environment.teardown();
+        environment.dispose();
+        throw err;
+      },
     );
 }
diff --git a/types/Environment.js b/types/Environment.js
index cd129a962779..295d6381c1bb 100644
--- a/types/Environment.js
+++ b/types/Environment.js
@@ -31,6 +31,8 @@ declare class $JestEnvironment {
   },
   testFilePath: string,
   moduleMocker: ModuleMocker,
+  setup(): Promise<Global>,
+  teardown(): Promise<void>,
 }
 
 export type Environment = $JestEnvironment;