vikerman · Apr 26, 2016
diff --git a/modules/angular2/src/testing/async.dart b/modules/angular2/src/testing/async.dart
@@ -0,0 +1,4 @@
+// This symbol is not used on the Dart side. This exists just as a stub.
+async(Function fn) {
+  throw 'async() test wrapper not available for Dart.';
+}
diff --git a/modules/angular2/src/testing/async.ts b/modules/angular2/src/testing/async.ts
@@ -0,0 +1,25 @@
+/**
+ * Wraps a test function in an asynchronous test zone. The test will automatically
+ * complete when all asynchronous calls within this zone are done. Can be used
+ * to wrap an {@link inject} call.
+ *
+ * Example:
+ *
+ * ```
+ * it('...', async(inject([AClass], (object) => {
+ *   object.doSomething.then(() => {
+ *     expect(...);
+ *   })
+ * });
+ * ```
+ */
+export function async(fn: Function): Function {
+  return () => {
+    return new Promise<void>((finishCallback, failCallback) => {
+      var AsyncTestZoneSpec = Zone['AsyncTestZoneSpec'];
+      var testZoneSpec = new AsyncTestZoneSpec(finishCallback, failCallback, 'test');
+      var testZone = Zone.current.fork(testZoneSpec);
+      return testZone.run(fn);
+    });
+  }
+}
diff --git a/modules/angular2/src/testing/async_test_completer.ts b/modules/angular2/src/testing/async_test_completer.ts
@@ -0,0 +1,13 @@
+import {PromiseCompleter} from 'angular2/src/facade/promise';
+
+/**
+ * Injectable completer that allows signaling completion of an asynchronous test. Used internally.
+ */
+export class AsyncTestCompleter {
+  private _completer = new PromiseCompleter<any>();
+  done(value?: any) { this._completer.resolve(value); }
+
+  fail(error?: any, stackTrace?: string) { this._completer.reject(error, stackTrace); }
+
+  get promise(): Promise<any> { return this._completer.promise; }
+}
diff --git a/modules/angular2/src/testing/fake_async.dart b/modules/angular2/src/testing/fake_async.dart
@@ -4,8 +4,6 @@ import 'dart:async' show runZoned, ZoneSpecification;
 import 'package:quiver/testing/async.dart' as quiver;
 import 'package:angular2/src/facade/exceptions.dart' show BaseException;
 
-import 'test_injector.dart' show getTestInjector, FunctionWithParamTokens;
-
 const _u = const Object();
 
 quiver.FakeAsync _fakeAsync = null;
@@ -22,24 +20,11 @@ quiver.FakeAsync _fakeAsync = null;
  *
  * Returns a `Function` that wraps [fn].
  */
-Function fakeAsync(dynamic /* Function | FunctionWithParamTokens */ fn) {
+Function fakeAsync(Function fn) {
   if (_fakeAsync != null) {
     throw 'fakeAsync() calls can not be nested';
   }
 
-  Function innerFn = null;
-  if (fn is FunctionWithParamTokens) {
-    if (fn.isAsync) {
-      throw 'Cannot wrap async test with fakeAsync';
-    }
-    innerFn = () { getTestInjector().execute(fn); };
-  } else if (fn is Function) {
-    innerFn = fn;
-  } else {
-    throw 'fakeAsync can wrap only test functions but got object of type ' +
-      fn.runtimeType.toString();
-  }
-
   return ([a0 = _u,
            a1 = _u,
            a2 = _u,
@@ -58,7 +43,7 @@ Function fakeAsync(dynamic /* Function | FunctionWithParamTokens */ fn) {
           List args = [a0, a1, a2, a3, a4, a5, a6, a7, a8, a9]
               .takeWhile((a) => a != _u)
               .toList();
-          var res = Function.apply(innerFn, args);
+          var res = Function.apply(fn, args);
           _fakeAsync.flushMicrotasks();
 
           if (async.periodicTimerCount > 0) {

diff --git a/modules/angular2/src/testing/fake_async.ts b/modules/angular2/src/testing/fake_async.ts
@@ -1,5 +1,5 @@
 import {BaseException} from 'angular2/src/facade/exceptions';
-import {getTestInjector, FunctionWithParamTokens} from './test_injector';
+import {getTestInjector} from './test_injector';
 
 let _FakeAsyncTestZoneSpecType = Zone['FakeAsyncTestZoneSpec'];
 
@@ -19,28 +19,17 @@ let _FakeAsyncTestZoneSpecType = Zone['FakeAsyncTestZoneSpec'];
  * @param fn
  * @returns {Function} The function wrapped to be executed in the fakeAsync zone
  */
-export function fakeAsync(fn: Function | FunctionWithParamTokens): Function {
+export function fakeAsync(fn: Function): Function {
   if (Zone.current.get('FakeAsyncTestZoneSpec') != null) {
     throw new BaseException('fakeAsync() calls can not be nested');
   }
 
   let fakeAsyncTestZoneSpec = new _FakeAsyncTestZoneSpecType();
   let fakeAsyncZone = Zone.current.fork(fakeAsyncTestZoneSpec);
 
-  let innerTestFn: Function = null;
-
-  if (fn instanceof FunctionWithParamTokens) {
-    if (fn.isAsync) {
-      throw new BaseException('Cannot wrap async test with fakeAsync');
-    }
-    innerTestFn = () => { getTestInjector().execute(fn as FunctionWithParamTokens); };
-  } else {
-    innerTestFn = fn;
-  }
-
   return function(...args) {
     let res = fakeAsyncZone.run(() => {
-      let res = innerTestFn(...args);
+      let res = fn(...args);
       flushMicrotasks();
       return res;
     });

diff --git a/modules/angular2/src/testing/test_injector.ts b/modules/angular2/src/testing/test_injector.ts
@@ -3,6 +3,11 @@ import {BaseException, ExceptionHandler} from 'angular2/src/facade/exceptions';
 import {ListWrapper} from 'angular2/src/facade/collection';
 import {FunctionWrapper, isPresent, Type} from 'angular2/src/facade/lang';
 
+import {async} from './async';
+import {AsyncTestCompleter} from './async_test_completer';
+
+export {async} from './async';
+
 export class TestInjector {
   private _instantiated: boolean = false;
 
@@ -35,15 +40,19 @@ export class TestInjector {
     return this._injector;
   }
 
-  execute(fn: FunctionWithParamTokens): any {
-    var additionalProviders = fn.additionalProviders();
-    if (additionalProviders.length > 0) {
-      this.addProviders(additionalProviders);
+  get(token: any) {
+    if (!this._instantiated) {
+      this.createInjector();
     }
+    return this._injector.get(token);
+  }
+
+  execute(tokens: any[], fn: Function): any {
     if (!this._instantiated) {
       this.createInjector();
     }
-    return fn.execute(this._injector);
+    var params = tokens.map(t => this._injector.get(t));
+    return FunctionWrapper.apply(fn, params);
   }
 }
 
@@ -117,22 +126,47 @@ export function resetBaseTestProviders() {
  *
  * @param {Array} tokens
  * @param {Function} fn
- * @return {FunctionWithParamTokens}
+ * @return {Function}
  */
-export function inject(tokens: any[], fn: Function): FunctionWithParamTokens {
-  return new FunctionWithParamTokens(tokens, fn, false);
+export function inject(tokens: any[], fn: Function): Function {
+  let testInjector = getTestInjector();
+  if (tokens.indexOf(AsyncTestCompleter) >= 0) {
+    // Return an async test method that returns a Promise if AsyncTestCompleter is one of the
+    // injected tokens.
+    return () => {
+      let completer: AsyncTestCompleter = testInjector.get(AsyncTestCompleter);
+      testInjector.execute(tokens, fn);
+      return completer.promise;
+    }
+  } else {
+    // Return a synchronous test method with the injected tokens.
+    return () => { return getTestInjector().execute(tokens, fn); };
+  }
 }
 
 export class InjectSetupWrapper {
   constructor(private _providers: () => any) {}
 
-  inject(tokens: any[], fn: Function): FunctionWithParamTokens {
-    return new FunctionWithParamTokens(tokens, fn, false, this._providers);
+  private _addProviders() {
+    var additionalProviders = this._providers();
+    if (additionalProviders.length > 0) {
+      getTestInjector().addProviders(additionalProviders);
+    }
+  }
+
+  inject(tokens: any[], fn: Function): Function {
+    return () => {
+      this._addProviders();
+      return inject(tokens, fn)();
+    }
   }
 
   /** @Deprecated {use async(withProviders().inject())} */
-  injectAsync(tokens: any[], fn: Function): FunctionWithParamTokens {
-    return new FunctionWithParamTokens(tokens, fn, true, this._providers);
+  injectAsync(tokens: any[], fn: Function): Function {
+    return () => {
+      this._addProviders();
+      return injectAsync(tokens, fn)();
+    }
   }
 }
 
@@ -158,53 +192,8 @@ export function withProviders(providers: () => any) {
  *
  * @param {Array} tokens
  * @param {Function} fn
- * @return {FunctionWithParamTokens}
+ * @return {Function}
  */
-export function injectAsync(tokens: any[], fn: Function): FunctionWithParamTokens {
-  return new FunctionWithParamTokens(tokens, fn, true);
-}
-
-/**
- * Wraps a test function in an asynchronous test zone. The test will automatically
- * complete when all asynchronous calls within this zone are done. Can be used
- * to wrap an {@link inject} call.
- *
- * Example:
- *
- * ```
- * it('...', async(inject([AClass], (object) => {
- *   object.doSomething.then(() => {
- *     expect(...);
- *   })
- * });
- * ```
- */
-export function async(fn: Function | FunctionWithParamTokens): FunctionWithParamTokens {
-  if (fn instanceof FunctionWithParamTokens) {
-    fn.isAsync = true;
-    return fn;
-  } else if (fn instanceof Function) {
-    return new FunctionWithParamTokens([], fn, true);
-  } else {
-    throw new BaseException('argument to async must be a function or inject(<Function>)');
-  }
-}
-
-function emptyArray(): Array<any> {
-  return [];
-}
-
-export class FunctionWithParamTokens {
-  constructor(private _tokens: any[], public fn: Function, public isAsync: boolean,
-              public additionalProviders: () => any = emptyArray) {}
-
-  /**
-   * Returns the value of the executed function.
-   */
-  execute(injector: ReflectiveInjector): any {
-    var params = this._tokens.map(t => injector.get(t));
-    return FunctionWrapper.apply(this.fn, params);
-  }
-
-  hasToken(token: any): boolean { return this._tokens.indexOf(token) > -1; }
+export function injectAsync(tokens: any[], fn: Function): Function {
+  return async(inject(tokens, fn));
 }