Bug 988844 - do_register_cleanup now accepts asynchronous cleanup functions. r=ted

testing/xpcshell/head.js

@@ -406,9 +406,44 @@ function _execute_test() {
   _load_files(_TAIL_FILES);
 
   // Execute all of our cleanup functions.
-  var func;
-  while ((func = _cleanupFunctions.pop()))
-    func();
+  let reportCleanupError = function(ex) {
+    let stack, filename;
+    if (ex && typeof ex == "object" && "stack" in ex) {
+      stack = ex.stack;
+    } else {
+      stack = Components.stack.caller;
+    }
+    if (stack instanceof Components.interfaces.nsIStackFrame) {
+      filename = stack.filename;
+    } else if (ex.fileName) {
+      filename = ex.fileName;
+    }
+    _log_message_with_stack("test_unexpected_fail",
+                            ex, stack, filename);
+  };
+
+  let func;
+  while ((func = _cleanupFunctions.pop())) {
+    let result;
+    try {
+      result = func();
+    } catch (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);
+      }
+    }
+  }
 
   // Restore idle service to avoid leaks.
   _fakeIdleService.deactivate();

testing/xpcshell/selftest.py

@@ -178,6 +178,49 @@ function run_test() {
 };
 '''
 
+# A test for asynchronous cleanup functions
+ASYNC_CLEANUP = '''
+function run_test() {
+  Components.utils.import("resource://gre/modules/Promise.jsm", this);
+
+  // The list of checkpoints in the order we encounter them.
+  let checkpoints = [];
+
+  // Cleanup tasks, in reverse order
+  do_register_cleanup(function cleanup_checkout() {
+    do_check_eq(checkpoints.join(""), "1234");
+    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(4);
+  });
+
+  do_register_cleanup(function async_cleanup_2() {
+    let deferred = Promise.defer();
+    do_execute_soon(deferred.resolve);
+    return deferred.promise.then(function() {
+      checkpoints.push(3);
+    });
+  });
+
+  do_register_cleanup(function sync_cleanup() {
+    checkpoints.push(2);
+  });
+
+  do_register_cleanup(function async_cleanup() {
+    let deferred = Promise.defer();
+    do_execute_soon(deferred.resolve);
+    return deferred.promise.then(function() {
+      checkpoints.push(1);
+    });
+  });
+
+}
+'''
+
+
 class XPCShellTestsTests(unittest.TestCase):
     """
     Yes, these are unit tests for a unit test harness.
@@ -723,5 +766,17 @@ tail =
         self.assertInLog("test_error.js:4")
         self.assertNotInLog("TEST-PASS")
 
+    def testAsyncCleanup(self):
+        """
+        Check that do_register_cleanup handles nicely cleanup tasks that
+        return a promise
+        """
+        self.writeFile("test_asyncCleanup.js", ASYNC_CLEANUP)
+        self.writeManifest(["test_asyncCleanup.js"])
+        self.assertTestResult(False)
+        self.assertInLog("\"1234\" == \"1234\"")
+        self.assertInLog("At this stage, the test has succeeded")
+        self.assertInLog("Throwing an error to force displaying the log")
+
 if __name__ == "__main__":
     unittest.main()