Support returning a Future from a setUp or test body in scheduled_test.

pkg/scheduled_test/lib/scheduled_test.dart

 // Copyright (c) 2013, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 // TODO(nweiz): Add support for calling [schedule] while the schedule is already
 // running.
 // TODO(nweiz): Port the non-Pub-specific scheduled test libraries from Pub.
 /// A package for writing readable tests of asynchronous behavior.
 ///
 /// ## Installing ##
@@ skipping 178 matching lines @@
 /// [pkg]: http://pub.dartlang.org/packages/scheduled_test
 library scheduled_test;
 
 import 'dart:async';
 
 import 'package:stack_trace/stack_trace.dart';
 import 'package:unittest/unittest.dart' as unittest;
 
 import 'src/schedule.dart';
 import 'src/schedule_error.dart';
+import 'src/utils.dart';
 
 export 'package:unittest/unittest.dart' hide
     test, solo_test, group, setUp, tearDown, completes, completion;
 
 export 'src/schedule.dart';
 export 'src/schedule_error.dart';
 export 'src/scheduled_future_matchers.dart';
 export 'src/task.dart';
 
 /// The [Schedule] for the current test. This is used to add new tasks and
 /// inspect the state of the schedule.
 ///
 /// This is `null` when there's no test currently running.
 Schedule get currentSchedule => _currentSchedule;
 Schedule _currentSchedule;
 
 /// The user-provided setUp function. This is set for each test during
 /// `unittest.setUp`.
 Function _setUpFn;
 
-/// Creates a new test case with the given description and body. This has the
-/// same semantics as [unittest.test].
-void test(String description, void body()) =>
+/// Creates a new test case with the given description and body.
+///
+/// This has the same semantics as [unittest.test].
+///
+/// If [body] returns a [Future], that future will automatically be wrapped with
+/// [wrapFuture].
+void test(String description, Future body()) =>
   _test(description, body, unittest.test);
 
 /// Creates a new test case with the given description and body that will be the
-/// only test run in this file. This has the same semantics as
-/// [unittest.solo_test].
-void solo_test(String description, void body()) =>
+/// only test run in this file.
+///
+/// This has the same semantics as [unittest.solo_test].
+///
+/// If [body] returns a [Future], that future will automatically be wrapped with
+/// [wrapFuture].
+void solo_test(String description, Future body()) =>
   _test(description, body, unittest.solo_test);
 
-void _test(String description, void body(), Function testFn) {
+void _test(String description, Future body(), Function testFn) {
+  maybeWrapFuture(future, description) {
+    if (future != null) wrapFuture(future, description);
+  }
+
   unittest.ensureInitialized();
   _ensureSetUpForTopLevel();
   testFn(description, () {
     var completer = new Completer();
 
     Chain.capture(() {
       _currentSchedule = new Schedule();
       return currentSchedule.run(() {
-        if (_setUpFn != null) _setUpFn();
-        body();
+        if (_setUpFn != null) maybeWrapFuture(_setUpFn(), "set up");
+        maybeWrapFuture(body(), "test body");
       }).catchError((error, stackTrace) {
         if (error is ScheduleError) {
           assert(error.schedule.errors.contains(error));
           assert(error.schedule == currentSchedule);
           unittest.registerException(error.schedule.errorString());
         } else {
           unittest.registerException(error, new Chain.forTrace(stackTrace));
         }
       }).then(completer.complete);
     }, onError: (error, chain) => currentSchedule.signalError(error, chain));
@@ skipping 116 matching lines @@
 /// [description] provides an optional description of the future, which is
 /// used when generating error messages.
 Future wrapFuture(Future future, [String description]) {
   if (currentSchedule == null) {
     throw new StateError("Unexpected call to wrapFuture with no current "
         "schedule.");
   }
 
   return currentSchedule.wrapFuture(future, description);
 }