Bug 1332295 - do_register_cleanup should support generators and async functions. r=ted

MozReview-Commit-ID: BPCwPlWQ8G0

testing/xpcshell/head.js

@@ -23,6 +23,7 @@ _register_modules_protocol_handler();
 
 var _Promise = Components.utils.import("resource://gre/modules/Promise.jsm", {}).Promise;
 var _PromiseTestUtils = Components.utils.import("resource://testing-common/PromiseTestUtils.jsm", {}).PromiseTestUtils;
+var _Task = Components.utils.import("resource://gre/modules/Task.jsm", {}).Task;
 Components.utils.importGlobalProperties(["XMLHttpRequest"]);
 
 // Support a common assertion library, Assert.jsm.
@@ -594,27 +595,22 @@ function _execute_test() {
                       });
   };
 
-  let func;
+  let complete = _cleanupFunctions.length == 0;
-  while ((func = _cleanupFunctions.pop())) {
+  _Task.spawn(function*() {
-    let result;
+    for (let func of _cleanupFunctions.reverse()) {
-    try {
+      try {
-      result = func();
+        yield func();
-    } catch (ex) {
+      } catch (ex) {
-      reportCleanupError(ex);
+        reportCleanupError(ex);
-      continue;
-    }
-    if (result && typeof result == "object"
-        && "then" in result && typeof result.then == "function") {
-      // This is a promise, wait until it is satisfied before proceeding
-      let complete = false;
-      let promise = result.then(null, reportCleanupError);
-      promise = promise.then(() => complete = true);
-      let thr = Components.classes["@mozilla.org/thread-manager;1"]
-                  .getService().currentThread;
-      while (!complete) {
-        thr.processNextEvent(true);
       }
     }
+    _cleanupFunctions = [];
+  }.bind(this)).catch(reportCleanupError)
+               .then(() => complete = true);
+  let thr = Components.classes["@mozilla.org/thread-manager;1"]
+                      .getService().currentThread;
+  while (!complete) {
+    thr.processNextEvent(true);
   }
 
   // Restore idle service to avoid leaks.
@@ -1512,7 +1508,6 @@ function add_task(funcOrProperties, func) {
 add_task.only = _add_only.bind(undefined, add_task);
 add_task.skip = _add_skip.bind(undefined, add_task);
 
-var _Task = Components.utils.import("resource://gre/modules/Task.jsm", {}).Task;
 _Task.Debugging.maintainStack = true;
 
 

testing/xpcshell/selftest.py

@@ -301,12 +301,22 @@ function run_test() {
 
   // Cleanup tasks, in reverse order
   do_register_cleanup(function cleanup_checkout() {
-    do_check_eq(checkpoints.join(""), "1234");
+    do_check_eq(checkpoints.join(""), "123456");
     do_print("At this stage, the test has succeeded");
     do_throw("Throwing an error to force displaying the log");
   });
 
   do_register_cleanup(function sync_cleanup_2() {
+    checkpoints.push(6);
+  });
+
+  do_register_cleanup(async function async_cleanup_4() {
+    await undefined;
+    checkpoints.push(5);
+  });
+
+  do_register_cleanup(function* async_cleanup_3() {
+    yield undefined;
     checkpoints.push(4);
   });
 
@@ -1184,13 +1194,12 @@ add_test({
 
     def testAsyncCleanup(self):
         """
-        Check that do_register_cleanup handles nicely cleanup tasks that
+        Check that do_register_cleanup handles nicely async cleanup tasks
-        return a promise
         """
         self.writeFile("test_asyncCleanup.js", ASYNC_CLEANUP)
         self.writeManifest(["test_asyncCleanup.js"])
         self.assertTestResult(False)
-        self.assertInLog("\"1234\" == \"1234\"")
+        self.assertInLog("\"123456\" == \"123456\"")
         self.assertInLog("At this stage, the test has succeeded")
         self.assertInLog("Throwing an error to force displaying the log")