From 1618234d538cd36e8c167a495820147577369866 Mon Sep 17 00:00:00 2001
From: Matt Loring <mattloring@google.com>
Date: Tue, 13 Jun 2017 17:54:58 -0700
Subject: [PATCH] Makes WrappedPromise an ES6 class

This allow for correct dispatch of method calls made on ES6 classes that
subclass promise.

Fixes: https://github.com/othiym23/async-listener/issues/112
---
 es6-wrapped-promise.js      | 37 ++++++++++++++++++++++++++++++++
 index.js                    | 42 +++++++++++++++++++++----------------
 test/native-promises.tap.js | 18 ++++++++++++++++
 test/promise-subclass.js    | 11 ++++++++++
 4 files changed, 90 insertions(+), 18 deletions(-)
 create mode 100644 es6-wrapped-promise.js
 create mode 100644 test/promise-subclass.js

diff --git a/es6-wrapped-promise.js b/es6-wrapped-promise.js
new file mode 100644
index 0000000..577b33d
--- /dev/null
+++ b/es6-wrapped-promise.js
@@ -0,0 +1,37 @@
+'use strict';
+
+module.exports = (Promise, ensureAslWrapper) => {
+  // Updates to this class should also be applied to the the ES3 version
+  // in index.js.
+  return class WrappedPromise extends Promise {
+    constructor(executor) {
+      var context, args;
+      super(wrappedExecutor);
+      var promise = this;
+
+      try {
+        executor.apply(context, args);
+      } catch (err) {
+        args[1](err);
+      }
+
+      return promise;
+      function wrappedExecutor(resolve, reject) {
+        context = this;
+        args = [wrappedResolve, wrappedReject];
+
+        // These wrappers create a function that can be passed a function and an argument to
+        // call as a continuation from the resolve or reject.
+        function wrappedResolve(val) {
+          ensureAslWrapper(promise, false);
+          return resolve(val);
+        }
+
+        function wrappedReject(val) {
+          ensureAslWrapper(promise, false);
+          return reject(val);
+        }
+      }
+    }
+  }
+};
diff --git a/index.js b/index.js
index 741d8eb..9086aa4 100644
--- a/index.js
+++ b/index.js
@@ -10,6 +10,7 @@ var shimmer      = require('shimmer')
   , util         = require('util')
   ;
 
+var v6plus = semver.gte(process.version, '6.0.0');
 var v7plus = semver.gte(process.version, '7.0.0');
 
 var net = require('net');
@@ -391,6 +392,8 @@ if (instrumentPromise) {
 function wrapPromise() {
   var Promise = global.Promise;
 
+  // Updates to this class should also be applied to the the ES6 version
+  // in es6-wrapped-promise.js.
   function wrappedPromise(executor) {
     if (!(this instanceof wrappedPromise)) {
       return Promise(executor);
@@ -407,7 +410,7 @@ function wrapPromise() {
     try {
       executor.apply(context, args);
     } catch (err) {
-      args[1](err)
+      args[1](err);
     }
 
     return promise;
@@ -438,23 +441,26 @@ function wrapPromise() {
     wrap(Promise.prototype, 'chain', wrapThen);
   }
 
-  var PromiseFunctions = [
-    'all',
-    'race',
-    'reject',
-    'resolve',
-    'accept',  // Node.js <v7 only
-    'defer'    // Node.js <v7 only
-  ];
-
-  PromiseFunctions.forEach(function(key) {
-    // don't break `in` by creating a key for undefined entries
-    if (typeof Promise[key] === 'function') {
-      wrappedPromise[key] = Promise[key];
-    }
-  });
-
-  global.Promise = wrappedPromise;
+  if (v6plus) {
+    global.Promise = require('./es6-wrapped-promise.js')(Promise, ensureAslWrapper);
+  } else {
+    var PromiseFunctions = [
+      'all',
+      'race',
+      'reject',
+      'resolve',
+      'accept',  // Node.js <v7 only
+      'defer'    // Node.js <v7 only
+    ];
+
+    PromiseFunctions.forEach(function(key) {
+      // don't break `in` by creating a key for undefined entries
+      if (typeof Promise[key] === 'function') {
+        wrappedPromise[key] = Promise[key];
+      }
+    });
+    global.Promise = wrappedPromise
+  }
 
   function ensureAslWrapper(promise, overwrite) {
     if (!promise.__asl_wrapper || overwrite) {
diff --git a/test/native-promises.tap.js b/test/native-promises.tap.js
index b9f071b..1424961 100644
--- a/test/native-promises.tap.js
+++ b/test/native-promises.tap.js
@@ -2198,3 +2198,21 @@ if (typeof Promise.prototype.chain === 'function') {
     }
   });
 }
+
+test('subclasses', function(t) {
+  if (nodeVersion[0] < 6) {
+    // class syntax is not supported before node 6.
+    t.end();
+    return;
+  }
+
+  // SubclassedPromise does 2 asserts.
+  t.plan(3);
+
+  var SubclassedPromise = require('./promise-subclass.js')(t);
+
+  var s = SubclassedPromise.resolve(42).then(function(val) {
+    t.strictEqual(val, 42);
+    t.end();
+  });
+});
diff --git a/test/promise-subclass.js b/test/promise-subclass.js
new file mode 100644
index 0000000..fbf2507
--- /dev/null
+++ b/test/promise-subclass.js
@@ -0,0 +1,11 @@
+'use strict';
+
+module.exports = (tap) => {
+  return class SubclassedPromise extends Promise {
+    then(onSuccess, onReject) {
+      tap.type(onSuccess, 'function');
+      tap.type(onReject, 'undefined');
+      return Promise.prototype.then.call(this, onSuccess, onReject);
+    }
+  };
+};