feat(test): can handle non zone aware task in promise (#1014)

Author: JiaLiPassion

Diff for: gulpfile.js

@@ -200,6 +200,14 @@ gulp.task('build/zone-patch-socket-io.min.js', ['compile-esm'], function(cb) {
   return generateScript('./lib/extra/socket-io.ts', 'zone-patch-socket-io.min.js', true, cb);
 });
 
+gulp.task('build/zone-patch-promise-testing.js', ['compile-esm'], function(cb) {
+  return generateScript('./lib/testing/promise-testing.ts', 'zone-patch-promise-test.js', false, cb);
+});
+
+gulp.task('build/zone-patch-promise-testing.min.js', ['compile-esm'], function(cb) {
+  return generateScript('./lib/testing/promise-testing.ts', 'zone-patch-promise-test.min.js', true, cb);
+});
+
 gulp.task('build/bluebird.js', ['compile-esm'], function(cb) {
     return generateScript('./lib/extra/bluebird.ts', 'zone-bluebird.js', false, cb);
 });
@@ -323,6 +331,8 @@ gulp.task('build', [
   'build/zone-patch-user-media.min.js',
   'build/zone-patch-socket-io.js',
   'build/zone-patch-socket-io.min.js',
+  'build/zone-patch-promise-testing.js',
+  'build/zone-patch-promise-testing.min.js',
   'build/zone-mix.js',
   'build/bluebird.js',
   'build/bluebird.min.js',

Diff for: karma-dist.conf.js

@@ -20,5 +20,6 @@ module.exports = function (config) {
   config.files.push('dist/sync-test.js');
   config.files.push('dist/task-tracking.js');
   config.files.push('dist/wtf.js');
+  config.files.push('dist/zone-patch-promise-test.js');
   config.files.push('build/test/main.js');
 };

Diff for: lib/common/promise.ts

@@ -268,7 +268,7 @@ Zone.__load_patch('ZoneAwarePromise', (global: any, Zone: ZoneType, api: _ZonePr
         // if error occurs, should always return this error
         resolvePromise(chainPromise, false, error);
       }
-    });
+    }, chainPromise as TaskData);
   }
 
   const ZONE_AWARE_PROMISE_TO_STRING = 'function ZoneAwarePromise() { [native code] }';

Diff for: lib/testing/promise-testing.ts

@@ -0,0 +1,68 @@
+/**
+ * @license
+ * Copyright Google Inc. All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.io/license
+ */
+
+ /**
+  * Promise for async/fakeAsync zoneSpec test
+  * can support async operation which not supported by zone.js
+  * such as
+  * it ('test jsonp in AsyncZone', async() => {
+  *   new Promise(res => {
+  *     jsonp(url, (data) => {
+  *       // success callback  
+  *       res(data); 
+  *     });
+  *   }).then((jsonpResult) => {
+  *     // get jsonp result. 
+  *
+  *     // user will expect AsyncZoneSpec wait for
+  *     // then, but because jsonp is not zone aware
+  *     // AsyncZone will finish before then is called.
+  *   });
+  * });
+  */
+Zone.__load_patch('promisefortest', (global: any, Zone: ZoneType, api: _ZonePrivate) => {
+  const symbolState: string = api.symbol('state');
+  const UNRESOLVED: null = null;
+  const symbolParentUnresolved = api.symbol('parentUnresolved');
+  
+  // patch Promise.prototype.then to keep an internal
+  // number for tracking unresolved chained promise
+  // we will decrease this number when the parent promise
+  // being resolved/rejected and chained promise was
+  // scheduled as a microTask.
+  // so we can know such kind of chained promise still
+  // not resolved in AsyncTestZone
+  (Promise as any)[api.symbol('patchPromiseForTest')] = function patchPromiseForTest() {
+    let oriThen = (Promise as any)[Zone.__symbol__('ZonePromiseThen')];
+    if (oriThen) {
+      return;
+    }
+    oriThen = (Promise as any)[Zone.__symbol__('ZonePromiseThen')] = Promise.prototype.then;
+    Promise.prototype.then = function() {
+      const chained = oriThen.apply(this, arguments); 
+      if (this[symbolState] === UNRESOLVED) {
+        // parent promise is unresolved.
+        const asyncTestZoneSpec = Zone.current.get('AsyncTestZoneSpec');
+        if (asyncTestZoneSpec) {
+          asyncTestZoneSpec.unresolvedChainedPromiseCount ++;
+          chained[symbolParentUnresolved] = true;
+        }
+      } 
+      return chained;
+    };
+  };
+
+  (Promise as any)[api.symbol('unPatchPromiseForTest')] = function unpatchPromiseForTest() {
+    // restore origin then
+    const oriThen = (Promise as any)[Zone.__symbol__('ZonePromiseThen')];
+    if (oriThen) {
+      Promise.prototype.then = oriThen;
+      (Promise as any)[Zone.__symbol__('ZonePromiseThen')] = undefined;
+    }
+  };
+});

Diff for: lib/testing/zone-testing.ts

@@ -12,4 +12,5 @@ import '../zone-spec/proxy';
 import '../zone-spec/sync-test';
 import '../jasmine/jasmine';
 import '../zone-spec/async-test';
-import '../zone-spec/fake-async-test';
+import '../zone-spec/fake-async-test';
+import './promise-testing';

Diff for: lib/zone-spec/async-test.ts

@@ -7,21 +7,27 @@
  */
 
 class AsyncTestZoneSpec implements ZoneSpec {
+  static symbolParentUnresolved = Zone.__symbol__('parentUnresolved');
+
   _finishCallback: Function;
   _failCallback: Function;
   _pendingMicroTasks: boolean = false;
   _pendingMacroTasks: boolean = false;
   _alreadyErrored: boolean = false;
   runZone = Zone.current;
+  unresolvedChainedPromiseCount = 0;
 
   constructor(finishCallback: Function, failCallback: Function, namePrefix: string) {
     this._finishCallback = finishCallback;
     this._failCallback = failCallback;
     this.name = 'asyncTestZone for ' + namePrefix;
+    this.properties = {
+      'AsyncTestZoneSpec': this 
+    };
   }
 
   _finishCallbackIfDone() {
-    if (!(this._pendingMicroTasks || this._pendingMacroTasks)) {
+    if (!(this._pendingMicroTasks || this._pendingMacroTasks || this.unresolvedChainedPromiseCount !== 0)) {
       // We do this because we would like to catch unhandled rejected promises.
       this.runZone.run(() => {
         setTimeout(() => {
@@ -33,19 +39,48 @@ class AsyncTestZoneSpec implements ZoneSpec {
     }
   }
 
+  patchPromiseForTest() {
+    const patchPromiseForTest = (Promise as any)[Zone.__symbol__('patchPromiseForTest')];
+    if (patchPromiseForTest) {
+      patchPromiseForTest();
+    }
+  }
+
+  unPatchPromiseForTest() {
+    const unPatchPromiseForTest = (Promise as any)[Zone.__symbol__('unPatchPromiseForTest')];
+    if (unPatchPromiseForTest) {
+      unPatchPromiseForTest();
+    }
+  }
+
   // ZoneSpec implementation below.
 
   name: string;
 
+  properties: {[key: string]: any};
+
+  onScheduleTask(delegate: ZoneDelegate, current: Zone, target: Zone, task: Task): Task {
+    if (task.type === 'microTask' && task.data && task.data instanceof Promise) {
+      // check whether the promise is a chained promise 
+      if ((task.data as any)[AsyncTestZoneSpec.symbolParentUnresolved] === true) {
+        // chained promise is being scheduled
+        this.unresolvedChainedPromiseCount --;
+      }
+    } 
+    return delegate.scheduleTask(target, task);
+  }
+
   // Note - we need to use onInvoke at the moment to call finish when a test is
   // fully synchronous. TODO(juliemr): remove this when the logic for
   // onHasTask changes and it calls whenever the task queues are dirty.
   onInvoke(
       parentZoneDelegate: ZoneDelegate, currentZone: Zone, targetZone: Zone, delegate: Function,
       applyThis: any, applyArgs: any[], source: string): any {
     try {
+      this.patchPromiseForTest();
       return parentZoneDelegate.invoke(targetZone, delegate, applyThis, applyArgs, source);
     } finally {
+      this.unPatchPromiseForTest();
       this._finishCallbackIfDone();
     }
   }

Diff for: test/browser-zone-setup.ts

@@ -19,3 +19,4 @@ import '../lib/zone-spec/sync-test';
 import '../lib/zone-spec/task-tracking';
 import '../lib/zone-spec/wtf';
 import '../lib/extra/cordova';
+import '../lib/testing/promise-testing';

Diff for: test/node_entry_point.ts

@@ -24,6 +24,7 @@ import '../lib/zone-spec/task-tracking';
 import '../lib/zone-spec/wtf';
 import '../lib/rxjs/rxjs';
 
+import '../lib/testing/promise-testing';
 // Setup test environment
 import './test-env-setup-jasmine';
 

Diff for: test/zone-spec/async-test.spec.ts

@@ -316,4 +316,31 @@ describe('AsyncTestZoneSpec', function() {
     });
 
   });
+
+  describe('non zone aware async task in promise should be detected', () => {
+    it('should be able to detect non zone aware async task in promise', (done) => {
+      let finished = false;
+
+      const testZoneSpec = new AsyncTestZoneSpec(
+          () => {
+            expect(finished).toBe(true);
+            done();
+          },
+          () => {
+            done.fail('async zone called failCallback unexpectedly');
+          },
+          'name');
+
+      const atz = Zone.current.fork(testZoneSpec);
+
+      atz.run(() => {
+        new Promise((res, rej) => {
+          const g: any = typeof window === 'undefined' ? global : window;
+          g[Zone.__symbol__('setTimeout')](res, 100);
+        }).then(() => {
+          finished = true;
+        });
+      });
+    });
+  });
 });