Make it idempotent to capturePromise (#400)

* Make it idempotent to capturePromise

Add uncapturePromise function in order to cancel patch in tests.

fix #398

* Move Move cleanup after the assertion

To ensure it's executed after the test will execute.

* Change setup process for process_p.test.js

To assure variables reference to original `then` and `catch` function.

Author: fossamagna

packages/core/lib/patchers/promise_p.js

@@ -9,23 +9,29 @@
 
 const contextUtils = require('../context_utils');
 
+const originalThen = Symbol('original then');
+const originalCatch = Symbol('original catch');
+
 function patchPromise(Promise) {
   const then = Promise.prototype.then;
-  Promise.prototype.then = function(onFulfilled, onRejected) {
-    if (contextUtils.isAutomaticMode()
-      && tryGetCurrentSegment()
-    ) {
-      const ns = contextUtils.getNamespace();
-
-      onFulfilled = onFulfilled && ns.bind(onFulfilled);
-      onRejected = onRejected && ns.bind(onRejected);
-    }
-
-    return then.call(this, onFulfilled, onRejected);
-  };
+  if (!then[originalThen]) {
+    Promise.prototype.then = function(onFulfilled, onRejected) {
+      if (contextUtils.isAutomaticMode()
+        && tryGetCurrentSegment()
+      ) {
+        const ns = contextUtils.getNamespace();
+  
+        onFulfilled = onFulfilled && ns.bind(onFulfilled);
+        onRejected = onRejected && ns.bind(onRejected);
+      }
+  
+      return then.call(this, onFulfilled, onRejected);
+    };
+    Promise.prototype.then[originalThen] = then;
+  }
 
   const origCatch = Promise.prototype.catch;
-  if (origCatch) {
+  if (origCatch && !origCatch[originalCatch]) {
     Promise.prototype.catch = function (onRejected) {
       if (contextUtils.isAutomaticMode()
         && tryGetCurrentSegment()
@@ -37,6 +43,18 @@ function patchPromise(Promise) {
 
       return origCatch.call(this, onRejected);
     };
+    Promise.prototype.catch[originalCatch] = origCatch;
+  }
+}
+
+function unpatchPromise(Promise) {
+  const then = Promise.prototype.then;
+  if (then[originalThen]) {
+    Promise.prototype.then = then[originalThen];
+  }
+  const origCatch = Promise.prototype.catch;
+  if (origCatch && origCatch[originalCatch]) {
+    Promise.prototype.catch = origCatch[originalCatch];
   }
 }
 
@@ -52,6 +70,11 @@ function capturePromise() {
   patchPromise(Promise);
 }
 
+function uncapturePromise() {
+  unpatchPromise(Promise);
+}
+
 capturePromise.patchThirdPartyPromise = patchPromise;
 
 module.exports.capturePromise = capturePromise;
+module.exports.uncapturePromise = uncapturePromise;

packages/core/test/integration/segment_maintained_across_shared_promise.test.js

@@ -37,6 +37,8 @@ var server = http
         // setTimeout so the assertion isn't caught by the promise
         setTimeout(function() {
           assert.equal(segment.id, retrievedSegment.id);
+          // Cancel the patch because it doesn't affect other tests
+          require('../../lib/patchers/promise_p').uncapturePromise();
         });
       });
     });

packages/core/test/unit/patchers/promise_p.test.js

@@ -1,12 +1,11 @@
 const expect = require('chai').expect;
 const sinon = require('sinon');
 
-const { capturePromise } = require('../../../lib/patchers/promise_p');
+const { capturePromise, uncapturePromise } = require('../../../lib/patchers/promise_p');
 const contextUtils = require('../../../lib/context_utils');
 
-const origThen = Promise.prototype.then;
-const origCatch = Promise.prototype.catch;
-
+let origThen;
+let origCatch;
 let getNamespaceStub;
 let isAutomaticModeStub;
 let getSegmentStub;
@@ -26,14 +25,16 @@ beforeEach(() => {
 
 beforeEach(() => {
   sinon.restore();
-  Promise.prototype.then = origThen;
-  Promise.prototype.catch = origCatch;
+  uncapturePromise();
+  origThen = Promise.prototype.then;
+  origCatch = Promise.prototype.catch;
 });
 
 after(() => {
   sinon.restore();
-  Promise.prototype.then = origThen;
-  Promise.prototype.catch = origCatch;
+  uncapturePromise();
+  origThen = Promise.prototype.then;
+  origCatch = Promise.prototype.catch;
 });
 
 function verifyBindings (segment, isAutomaticMode, onFulfilled, onRejected) {
@@ -125,6 +126,20 @@ describe('capturePromise', () => {
     });
     createTests(true);
   });
+
+  it('should be idempotent', async () => {
+    capturePromise();
+    const thenFirst = Promise.prototype.then;
+    const catchFirst = Promise.prototype.catch;
+    capturePromise();
+    const thenSecond = Promise.prototype.then;
+    const catchSecond = Promise.prototype.catch;
+    capturePromise();
+    expect(Promise.prototype.then).to.equal(thenFirst);
+    expect(Promise.prototype.catch).to.equal(catchFirst);
+    expect(Promise.prototype.then).to.equal(thenSecond);
+    expect(Promise.prototype.catch).to.equal(catchSecond);
+  });
 });
 
 describe('capturePromise.patchThirdPartyPromise', () => {