From 7b0d0e249d40f20523c81fd742a8d83bb7090362 Mon Sep 17 00:00:00 2001
From: Simen Bekkhus <sbekkhus91@gmail.com>
Date: Sun, 13 May 2018 19:06:08 +0200
Subject: [PATCH] spy on timer functions

---
 packages/jest-environment-jsdom/src/index.js |  6 ++-
 packages/jest-environment-node/src/index.js  |  6 ++-
 packages/jest-util/src/fake_timers.js        | 43 ++++++++++++++++++--
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/packages/jest-environment-jsdom/src/index.js b/packages/jest-environment-jsdom/src/index.js
index ec6e9ac2a9d7..8b7b710dde77 100644
--- a/packages/jest-environment-jsdom/src/index.js
+++ b/packages/jest-environment-jsdom/src/index.js
@@ -72,7 +72,11 @@ class JSDOMEnvironment {
 
     this.moduleMocker = new mock.ModuleMocker(global);
 
-    this.fakeTimers = new FakeTimers({config, global});
+    this.fakeTimers = new FakeTimers({
+      config,
+      global,
+      moduleMocker: this.moduleMocker,
+    });
   }
 
   setup(): Promise<void> {
diff --git a/packages/jest-environment-node/src/index.js b/packages/jest-environment-node/src/index.js
index f7d87f3cd631..fef3495a15ae 100644
--- a/packages/jest-environment-node/src/index.js
+++ b/packages/jest-environment-node/src/index.js
@@ -43,7 +43,11 @@ class NodeEnvironment {
     installCommonGlobals(global, config.globals);
     this.moduleMocker = new mock.ModuleMocker(global);
 
-    this.fakeTimers = new FakeTimers({config, global});
+    this.fakeTimers = new FakeTimers({
+      config,
+      global,
+      moduleMocker: this.moduleMocker,
+    });
   }
 
   setup(): Promise<void> {
diff --git a/packages/jest-util/src/fake_timers.js b/packages/jest-util/src/fake_timers.js
index f94d0070192d..02d4492be1b2 100644
--- a/packages/jest-util/src/fake_timers.js
+++ b/packages/jest-util/src/fake_timers.js
@@ -9,33 +9,50 @@
 
 import type {ProjectConfig} from 'types/Config';
 import type {Global} from 'types/Global';
+import type {ModuleMocker} from 'types/Mock';
 
 import {withGlobal as lolexWithGlobal} from 'lolex';
 import {formatStackTrace} from 'jest-message-util';
 
+const functionsToSpyOn: Array<string> = [
+  'clearImmediate',
+  'clearInterval',
+  'clearTimeout',
+  'nextTick',
+  'setImmediate',
+  'setInterval',
+  'setTimeout',
+];
+
 export default class FakeTimers {
   _clock: LolexClock;
   _config: ProjectConfig;
   _fakingTime: boolean;
   _global: Global;
-  _maxLoops: number;
   _lolex: LolexExports;
+  _maxLoops: number;
+  _moduleMocker: ModuleMocker;
+  _timerMocks: Array<JestMockFn>;
 
   constructor({
     global,
+    moduleMocker,
     config,
     maxLoops,
   }: {
     global: Global,
+    moduleMocker: ModuleMocker,
     config: ProjectConfig,
     maxLoops?: number,
   }) {
     this._global = global;
     this._config = config;
     this._maxLoops = maxLoops || 100000;
+    this._moduleMocker = moduleMocker;
 
     this._fakingTime = false;
     this._lolex = lolexWithGlobal(global);
+    this._timerMocks = [];
   }
 
   clearAllTimers() {
@@ -81,20 +98,38 @@ export default class FakeTimers {
 
   useRealTimers() {
     if (this._fakingTime) {
+      this._timerMocks.forEach(func => {
+        func.mockRestore();
+      });
       this._clock.uninstall();
       this._fakingTime = false;
+      this._timerMocks = [];
     }
   }
 
   useFakeTimers() {
     if (!this._fakingTime) {
-      // This creates stubs based on the host environment, we want it to be based
-      // on `this._global`. See https://github.com/sinonjs/lolex/issues/146
+      const toFake = Object.keys(this._lolex.timers);
       this._clock = this._lolex.install({
         loopLimit: this._maxLoops,
         target: this._global,
-        toFake: Object.keys(this._lolex.timers),
+        toFake,
       });
+
+      this._timerMocks = toFake
+        .filter(func => functionsToSpyOn.includes(func))
+        .map(func => {
+          if (func === 'hrtime' || func === 'nextTick') {
+            return this._moduleMocker.spyOn(this._global.process, func);
+          }
+
+          if (func === 'performance') {
+            return this._moduleMocker.spyOn(this._global.performance, 'now');
+          }
+
+          return this._moduleMocker.spyOn(this._global, func);
+        });
+
       this._fakingTime = true;
     }
   }