From e7ea47c3eb31609ee93abdc0f9c670f63a7cbb64 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 | 48 +++++++++++++++++++---------
 types/Environment.js                 |  2 ++
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/packages/jest-runner/src/run_test.js b/packages/jest-runner/src/run_test.js
index a2110f157fa6..3cd824fd8492 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 +
@@ -117,22 +121,36 @@ export default function runTest(
     })
     .then(
       result =>
-        Promise.resolve().then(() => {
-          environment.dispose();
-          if (globalConfig.logHeapUsage) {
-            if (global.gc) {
-              global.gc();
+        Promise.resolve()
+          .then(
+            () =>
+              environment.teardown
+                ? environment.teardown()
+                : Promise.resolve({}),
+          )
+          .then(() => {
+            environment.dispose();
+            if (globalConfig.logHeapUsage) {
+              if (global.gc) {
+                global.gc();
+              }
+              result.memoryUsage = process.memoryUsage().heapUsed;
             }
-            result.memoryUsage = process.memoryUsage().heapUsed;
-          }
 
-          // Delay the resolution to allow log messages to be output.
-          return new Promise(resolve => setImmediate(() => resolve(result)));
-        }),
+            // 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;
-        }),
+        Promise.resolve()
+          .then(
+            () =>
+              environment.teardown
+                ? environment.teardown()
+                : Promise.resolve({}),
+          )
+          .then(() => {
+            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;