Bug 1262368 - Fix browser_UITour_heartbeat.js to work on e10s. r=MattN
authorAlessio Placitelli <alessio.placitelli@gmail.com>
Fri, 08 Apr 2016 23:50:00 +0200
changeset 316333 a003e040f00771f18c84b691a58c09a0c139f953
parent 316313 1709e72eedc79a1953aa8c2712f6c9c604542172
child 316334 8ff3c4cb70c382a99c2c46558242fce8cc971dd7
push id9480
push userjlund@mozilla.com
push dateMon, 25 Apr 2016 17:12:58 +0000
treeherdermozilla-aurora@0d6a91c76a9e [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersMattN
bugs1262368
milestone48.0a1
Bug 1262368 - Fix browser_UITour_heartbeat.js to work on e10s. r=MattN MozReview-Commit-ID: 6DtCPx9wa1r
browser/components/uitour/test/browser.ini
browser/components/uitour/test/browser_UITour_heartbeat.js
browser/components/uitour/test/head.js
--- a/browser/components/uitour/test/browser.ini
+++ b/browser/components/uitour/test/browser.ini
@@ -26,17 +26,16 @@ skip-if = os == "linux" # Intermittent f
 [browser_UITour3.js]
 skip-if = os == "linux" # Linux: Bug 986760, Bug 989101.
 [browser_UITour_availableTargets.js]
 [browser_UITour_annotation_size_attributes.js]
 [browser_UITour_defaultBrowser.js]
 [browser_UITour_detach_tab.js]
 [browser_UITour_forceReaderMode.js]
 [browser_UITour_heartbeat.js]
-skip-if = e10s # Bug 1240747 - UITour.jsm not e10s friendly.
 [browser_UITour_loop.js]
 skip-if = true # Bug 1225832 - New Loop architecture is not compatible with test.
 [browser_UITour_loop_panel.js]
 [browser_UITour_modalDialog.js]
 skip-if = os != "mac" # modal dialog disabling only working on OS X.
 [browser_UITour_observe.js]
 [browser_UITour_panel_close_annotation.js]
 skip-if = true # Disabled due to frequent failures, bugs 1026310 and 1032137
--- a/browser/components/uitour/test/browser_UITour_heartbeat.js
+++ b/browser/components/uitour/test/browser_UITour_heartbeat.js
@@ -2,24 +2,16 @@
  * http://creativecommons.org/publicdomain/zero/1.0/ */
 
 "use strict";
 
 var gTestTab;
 var gContentAPI;
 var gContentWindow;
 
-function test() {
-  UITourTest();
-  requestLongerTimeout(2);
-  registerCleanupFunction(() => {
-    Services.prefs.clearUserPref("browser.uitour.surveyDuration");
-  });
-}
-
 function getHeartbeatNotification(aId, aChromeWindow = window) {
   let notificationBox = aChromeWindow.document.getElementById("high-priority-global-notificationbox");
   // UITour.jsm prefixes the notification box ID with "heartbeat-" to prevent collisions.
   return notificationBox.getNotificationWithValue("heartbeat-" + aId);
 }
 
 /**
  * Simulate a click on a rating element in the Heartbeat notification.
@@ -97,524 +89,622 @@ function checkTelemetry(aPayload, aFlowI
       ok(Number.isInteger(ts) && ts > 0, "Timestamp '" + field + "' must be a natural number");
     }
     extraKeys.delete(field);
   }
 
   is(extraKeys.size, 0, "No unexpected fields in the Telemetry payload");
 }
 
-var tests = [
-  /**
-   * Check that the "stars" heartbeat UI correctly shows and closes.
-   */
-  function test_heartbeat_stars_show(done) {
-    let flowId = "ui-ratefirefox-" + Math.random();
-    let engagementURL = "http://example.com";
+/**
+ * Waits for an UITour notification dispatched through |UITour.notify|. This should be
+ * done with |gContentAPI.observe|. Unfortunately, in e10s, |gContentAPI.observe| doesn't
+ * allow for multiple calls to the same callback, allowing to catch just the first
+ * notification.
+ *
+ * @param aEventName
+ *        The notification name to wait for.
+ * @return {Promise} Resolved with the data that comes with the event.
+ */
+function promiseWaitHeartbeatNotification(aEventName) {
+  return ContentTask.spawn(gTestTab.linkedBrowser, { aEventName },
+      function({ aEventName }) {
+        return new Promise(resolve => {
+          addEventListener("mozUITourNotification", function listener(event) {
+            if (event.detail.event !== aEventName) {
+              return;
+            }
+            removeEventListener("mozUITourNotification", listener, false);
+            resolve(event.detail.params);
+          }, false);
+        });
+      });
+}
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          cleanUpNotification(flowId);
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          done();
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received");
-          checkTelemetry(aData, flowId, ["offeredTS", "closedTS"]);
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+/**
+ * Waits for UITour notifications dispatched through |UITour.notify|. This works like
+ * |promiseWaitHeartbeatNotification|, but waits for all the passed notifications to
+ * be received before resolving. If it receives an unaccounted notification, it rejects.
+ *
+ * @param events
+ *        An array of expected notification names to wait for.
+ * @return {Promise} Resolved with the data that comes with the event. Rejects with the
+ *         name of an undesired notification if received.
+ */
+function promiseWaitExpectedNotifications(events) {
+  return ContentTask.spawn(gTestTab.linkedBrowser, { events },
+      function({ events }) {
+        let stillToReceive = events;
+        return new Promise((res, rej) => {
+          addEventListener("mozUITourNotification", function listener(event) {
+            if (stillToReceive.includes(event.detail.event)) {
+              // Filter out the received event.
+              stillToReceive = stillToReceive.filter(x => x !== event.detail.event);
+            } else {
+              removeEventListener("mozUITourNotification", listener, false);
+              rej(event.detail.event);
+            }
+            // We still need to catch some notifications. Don't do anything.
+            if (stillToReceive.length > 0) {
+              return;
+            }
+            // We don't need to listen for other notifications. Resolve the promise.
+            removeEventListener("mozUITourNotification", listener, false);
+            res();
+          }, false);
+        });
+      });
+}
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, engagementURL);
-  },
+function validateTimestamp(eventName, timestamp) {
+  info("'" + eventName + "' notification received (timestamp " + timestamp.toString() + ").");
+  ok(Number.isFinite(timestamp), "Timestamp must be a number.");
+}
 
-  /**
-   * Test that the heartbeat UI correctly works with null engagement URL.
-   */
-  function test_heartbeat_null_engagementURL(done) {
-    let flowId = "ui-ratefirefox-" + Math.random();
-    let originalTabCount = gBrowser.tabs.length;
+add_task(function* test_setup(){
+  yield setup_UITourTest();
+  requestLongerTimeout(2);
+  registerCleanupFunction(() => {
+    Services.prefs.clearUserPref("browser.uitour.surveyDuration");
+  });
+});
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
-          simulateVote(flowId, 2);
-          break;
-        }
-        case "Heartbeat:Voted": {
-          info("'Heartbeat:Voted' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          is(gBrowser.tabs.length, originalTabCount, "No engagement tab should be opened.");
-          done();
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received.");
-          checkTelemetry(aData, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
-          is(aData.score, 2, "Checking Telemetry payload.score");
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+/**
+ * Check that the "stars" heartbeat UI correctly shows and closes.
+ */
+add_UITour_task(function* test_heartbeat_stars_show() {
+  let flowId = "ui-ratefirefox-" + Math.random();
+  let engagementURL = "http://example.com";
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(
+    ["Heartbeat:NotificationOffered", "Heartbeat:NotificationClosed", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, engagementURL);
+
+  // Validate the returned timestamp.
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Close the heartbeat notification.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
+  cleanUpNotification(flowId);
+
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received");
+  checkTelemetry(data, flowId, ["offeredTS", "closedTS"]);
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, null);
-  },
+/**
+ * Test that the heartbeat UI correctly works with null engagement URL.
+ */
+add_UITour_task(function* test_heartbeat_null_engagementURL() {
+  let flowId = "ui-ratefirefox-" + Math.random();
+  let originalTabCount = gBrowser.tabs.length;
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:Voted", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, null);
 
-   /**
-   * Test that the heartbeat UI correctly works with an invalid, but non null, engagement URL.
-   */
-  function test_heartbeat_invalid_engagement_URL(done) {
-    let flowId = "ui-ratefirefox-" + Math.random();
-    let originalTabCount = gBrowser.tabs.length;
-    let invalidEngagementURL = "invalidEngagement";
+  // Validate the returned timestamp.
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Wait an the Voted, Closed and Telemetry Sent events. They are fired together, so
+  // wait for them here.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let votedPromise = promiseWaitHeartbeatNotification("Heartbeat:Voted");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
+
+  // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
+  simulateVote(flowId, 2);
+  data = yield votedPromise;
+  validateTimestamp('Heartbeat:Voted', data.timestamp);
+
+  // Validate the closing timestamp.
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+  is(gBrowser.tabs.length, originalTabCount, "No engagement tab should be opened.");
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
-          simulateVote(flowId, 2);
-          break;
-        }
-        case "Heartbeat:Voted": {
-          info("'Heartbeat:Voted' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          is(gBrowser.tabs.length, originalTabCount, "No engagement tab should be opened.");
-          done();
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received.");
-          checkTelemetry(aData, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
-          is(aData.score, 2, "Checking Telemetry payload.score");
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+  // Validate the data we send out.
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
+  is(data.score, 2, "Checking Telemetry payload.score");
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
+
+/**
+ * Test that the heartbeat UI correctly works with an invalid, but non null, engagement URL.
+ */
+add_UITour_task(function* test_heartbeat_invalid_engagement_URL() {
+  let flowId = "ui-ratefirefox-" + Math.random();
+  let originalTabCount = gBrowser.tabs.length;
+  let invalidEngagementURL = "invalidEngagement";
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:Voted", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, invalidEngagementURL);
+
+  // Validate the returned timestamp.
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Wait an the Voted, Closed and Telemetry Sent events. They are fired together, so
+  // wait for them here.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let votedPromise = promiseWaitHeartbeatNotification("Heartbeat:Voted");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, invalidEngagementURL);
-  },
+  // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
+  simulateVote(flowId, 2);
+  data = yield votedPromise;
+  validateTimestamp('Heartbeat:Voted', data.timestamp);
+
+  // Validate the closing timestamp.
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+  is(gBrowser.tabs.length, originalTabCount, "No engagement tab should be opened.");
+
+  // Validate the data we send out.
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
+  is(data.score, 2, "Checking Telemetry payload.score");
 
-  /**
-   * Test that the score is correctly reported.
-   */
-  function test_heartbeat_stars_vote(done) {
-    const expectedScore = 4;
-    let flowId = "ui-ratefirefox-" + Math.random();
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
+
+/**
+ * Test that the score is correctly reported.
+ */
+add_UITour_task(function* test_heartbeat_stars_vote() {
+  const expectedScore = 4;
+  let originalTabCount = gBrowser.tabs.length;
+  let flowId = "ui-ratefirefox-" + Math.random();
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:Voted", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, null);
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
-          simulateVote(flowId, expectedScore);
-          break;
-        }
-        case "Heartbeat:Voted": {
-          info("'Heartbeat:Voted' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          is(aData.score, expectedScore, "Should report a score of " + expectedScore);
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          done();
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received.");
-          checkTelemetry(aData, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
-          is(aData.score, expectedScore, "Checking Telemetry payload.score");
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+  // Validate the returned timestamp.
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Wait an the Voted, Closed and Telemetry Sent events. They are fired together, so
+  // wait for them here.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let votedPromise = promiseWaitHeartbeatNotification("Heartbeat:Voted");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, null);
-  },
+  // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
+  simulateVote(flowId, expectedScore);
+  data = yield votedPromise;
+  validateTimestamp('Heartbeat:Voted', data.timestamp);
+  is(data.score, expectedScore, "Should report a score of " + expectedScore);
 
-  /**
-   * Test that the engagement page is correctly opened when voting.
-   */
-  function test_heartbeat_engagement_tab(done) {
-    let engagementURL = "http://example.com";
-    let flowId = "ui-ratefirefox-" + Math.random();
-    let originalTabCount = gBrowser.tabs.length;
-    const expectedTabCount = originalTabCount + 1;
-    let heartbeatVoteSeen = false;
+  // Validate the closing timestamp and vote.
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+  is(gBrowser.tabs.length, originalTabCount, "No engagement tab should be opened.");
+
+  // Validate the data we send out.
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
+  is(data.score, expectedScore, "Checking Telemetry payload.score");
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
-          simulateVote(flowId, 1);
-          break;
-        }
-        case "Heartbeat:Voted": {
-          info("'Heartbeat:Voted' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          heartbeatVoteSeen = true;
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          ok(heartbeatVoteSeen, "Heartbeat vote should have been received");
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          is(gBrowser.tabs.length, expectedTabCount, "Engagement URL should open in a new tab.");
-          gBrowser.removeCurrentTab();
-          done();
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received.");
-          checkTelemetry(aData, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
-          is(aData.score, 1, "Checking Telemetry payload.score");
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+/**
+ * Test that the engagement page is correctly opened when voting.
+ */
+add_UITour_task(function* test_heartbeat_engagement_tab() {
+  let engagementURL = "http://example.com";
+  let flowId = "ui-ratefirefox-" + Math.random();
+  let originalTabCount = gBrowser.tabs.length;
+  const expectedTabCount = originalTabCount + 1;
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:Voted", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, engagementURL);
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, engagementURL);
-  },
+  // Validate the returned timestamp.
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Wait an the Voted, Closed and Telemetry Sent events. They are fired together, so
+  // wait for them here.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let votedPromise = promiseWaitHeartbeatNotification("Heartbeat:Voted");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
 
-  /**
-   * Test that the engagement button opens the engagement URL.
-   */
-  function test_heartbeat_engagement_button(done) {
-    let engagementURL = "http://example.com";
-    let flowId = "ui-engagewithfirefox-" + Math.random();
-    let originalTabCount = gBrowser.tabs.length;
-    const expectedTabCount = originalTabCount + 1;
-    let heartbeatEngagedSeen = false;
+  // The UI was just shown. We can simulate a click on a rating element (i.e., "star").
+  simulateVote(flowId, 1);
+  data = yield votedPromise;
+  validateTimestamp('Heartbeat:Voted', data.timestamp);
+
+  // Validate the closing timestamp, vote and make sure the engagement page was opened.
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+  is(gBrowser.tabs.length, expectedTabCount, "Engagement URL should open in a new tab.");
+  gBrowser.removeCurrentTab();
+
+  // Validate the data we send out.
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "votedTS", "closedTS", "score"]);
+  is(data.score, 1, "Checking Telemetry payload.score");
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          let notification = getHeartbeatNotification(flowId);
-          is(notification.querySelectorAll(".star-x").length, 0, "No stars should be present");
-          // The UI was just shown. We can simulate a click on the engagement button.
-          let engagementButton = notification.querySelector(".notification-button");
-          is(engagementButton.label, "Engage Me", "Check engagement button text");
-          engagementButton.doCommand();
-          break;
-        }
-        case "Heartbeat:Engaged": {
-          info("'Heartbeat:Engaged' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          heartbeatEngagedSeen = true;
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(heartbeatEngagedSeen, "Heartbeat:Engaged should have been received");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          is(gBrowser.tabs.length, expectedTabCount, "Engagement URL should open in a new tab.");
-          gBrowser.removeCurrentTab();
-          executeSoon(done);
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received.");
-          checkTelemetry(aData, flowId, ["offeredTS", "engagedTS", "closedTS"]);
-          break;
-        }
-        default: {
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-        }
-      }
-    });
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
+
+/**
+ * Test that the engagement button opens the engagement URL.
+ */
+add_UITour_task(function* test_heartbeat_engagement_button() {
+  let engagementURL = "http://example.com";
+  let flowId = "ui-engagewithfirefox-" + Math.random();
+  let originalTabCount = gBrowser.tabs.length;
+  const expectedTabCount = originalTabCount + 1;
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:Engaged", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL, null, null, {
+    engagementButtonLabel: "Engage Me",
+  });
+
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Wait an the Engaged, Closed and Telemetry Sent events. They are fired together, so
+  // wait for them here.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let engagedPromise = promiseWaitHeartbeatNotification("Heartbeat:Engaged");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
 
-    gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL, null, null, {
-      engagementButtonLabel: "Engage Me",
-    });
-  },
+  // Simulate user engagement.
+  let notification = getHeartbeatNotification(flowId);
+  is(notification.querySelectorAll(".star-x").length, 0, "No stars should be present");
+  // The UI was just shown. We can simulate a click on the engagement button.
+  let engagementButton = notification.querySelector(".notification-button");
+  is(engagementButton.label, "Engage Me", "Check engagement button text");
+  engagementButton.doCommand();
+
+  data = yield engagedPromise;
+  validateTimestamp('Heartbeat:Engaged', data.timestamp);
+
+  // Validate the closing timestamp, vote and make sure the engagement page was opened.
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+  is(gBrowser.tabs.length, expectedTabCount, "Engagement URL should open in a new tab.");
+  gBrowser.removeCurrentTab();
 
-  /**
-   * Test that the learn more link is displayed and that the page is correctly opened when
-   * clicking on it.
-   */
-  function test_heartbeat_learnmore(done) {
-    let dummyURL = "http://example.com";
-    let flowId = "ui-ratefirefox-" + Math.random();
-    let originalTabCount = gBrowser.tabs.length;
-    const expectedTabCount = originalTabCount + 1;
+  // Validate the data we send out.
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "engagedTS", "closedTS"]);
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
+
+/**
+ * Test that the learn more link is displayed and that the page is correctly opened when
+ * clicking on it.
+ */
+add_UITour_task(function* test_heartbeat_learnmore() {
+  let dummyURL = "http://example.com";
+  let flowId = "ui-ratefirefox-" + Math.random();
+  let originalTabCount = gBrowser.tabs.length;
+  const expectedTabCount = originalTabCount + 1;
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          // The UI was just shown. Simulate a click on the learn more link.
-          clickLearnMore(flowId);
-          break;
-        }
-        case "Heartbeat:LearnMore": {
-          info("'Heartbeat:LearnMore' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          cleanUpNotification(flowId);
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          ok(Number.isFinite(aData.timestamp), "Timestamp must be a number.");
-          is(gBrowser.tabs.length, expectedTabCount, "Learn more URL should open in a new tab.");
-          gBrowser.removeCurrentTab();
-          done();
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received.");
-          checkTelemetry(aData, flowId, ["offeredTS", "learnMoreTS", "closedTS"]);
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:LearnMore", "Heartbeat:TelemetrySent"]);
+
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, dummyURL,
+                            "What is this?", dummyURL);
+
+  let data = yield shownPromise;
+  validateTimestamp('Heartbeat:Offered', data.timestamp);
+
+  // Wait an the LearnMore, Closed and Telemetry Sent events. They are fired together, so
+  // wait for them here.
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let learnMorePromise = promiseWaitHeartbeatNotification("Heartbeat:LearnMore");
+  let pingSentPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, dummyURL,
-                              "What is this?", dummyURL);
-  },
+  // The UI was just shown. Simulate a click on the learn more link.
+  clickLearnMore(flowId);
+
+  data = yield learnMorePromise;
+  validateTimestamp('Heartbeat:LearnMore', data.timestamp);
+  cleanUpNotification(flowId);
 
-  taskify(function* test_invalidEngagementButtonLabel() {
-    let engagementURL = "http://example.com";
-    let flowId = "invalidEngagementButtonLabel-" + Math.random();
+  // The notification was closed.
+  data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+  is(gBrowser.tabs.length, expectedTabCount, "Learn more URL should open in a new tab.");
+  gBrowser.removeCurrentTab();
+
+  // Validate the data we send out.
+  data = yield pingSentPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "learnMoreTS", "closedTS"]);
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
 
-    let eventPromise = promisePageEvent();
+add_UITour_task(function* test_invalidEngagementButtonLabel() {
+  let engagementURL = "http://example.com";
+  let flowId = "invalidEngagementButtonLabel-" + Math.random();
 
-    gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
-                              null, null, {
-                                engagementButtonLabel: 42,
-                              });
+  let eventPromise = promisePageEvent();
 
-    yield eventPromise;
-    ok(!isTourBrowser(gBrowser.selectedBrowser),
-       "Invalid engagementButtonLabel should prevent init");
-
-  }),
+  gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
+                            null, null, {
+                              engagementButtonLabel: 42,
+                            });
 
-  taskify(function* test_privateWindowsOnly_noneOpen() {
-    let engagementURL = "http://example.com";
-    let flowId = "privateWindowsOnly_noneOpen-" + Math.random();
+  yield eventPromise;
+  ok(!isTourBrowser(gBrowser.selectedBrowser),
+     "Invalid engagementButtonLabel should prevent init");
 
-    let eventPromise = promisePageEvent();
+})
 
-    gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
-                              null, null, {
-                                engagementButtonLabel: "Yes!",
-                                privateWindowsOnly: true,
-                              });
+add_UITour_task(function* test_privateWindowsOnly_noneOpen() {
+  let engagementURL = "http://example.com";
+  let flowId = "privateWindowsOnly_noneOpen-" + Math.random();
+
+  let eventPromise = promisePageEvent();
 
-    yield eventPromise;
-    ok(!isTourBrowser(gBrowser.selectedBrowser),
-       "If there are no private windows opened, tour init should be prevented");
-  }),
+  gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
+                            null, null, {
+                              engagementButtonLabel: "Yes!",
+                              privateWindowsOnly: true,
+                            });
 
-  taskify(function* test_privateWindowsOnly_notMostRecent() {
-    let engagementURL = "http://example.com";
-    let flowId = "notMostRecent-" + Math.random();
-
-    let privateWin = yield BrowserTestUtils.openNewBrowserWindow({ private: true });
-    let mostRecentWin = yield BrowserTestUtils.openNewBrowserWindow();
+  yield eventPromise;
+  ok(!isTourBrowser(gBrowser.selectedBrowser),
+     "If there are no private windows opened, tour init should be prevented");
+})
 
-    let eventPromise = promisePageEvent();
+add_UITour_task(function* test_privateWindowsOnly_notMostRecent() {
+  let engagementURL = "http://example.com";
+  let flowId = "notMostRecent-" + Math.random();
 
-    gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
-                              null, null, {
-                                engagementButtonLabel: "Yes!",
-                                privateWindowsOnly: true,
-                              });
+  let privateWin = yield BrowserTestUtils.openNewBrowserWindow({ private: true });
+  let mostRecentWin = yield BrowserTestUtils.openNewBrowserWindow();
 
-    yield eventPromise;
-    is(getHeartbeatNotification(flowId, window), null,
-       "Heartbeat shouldn't appear in the default window");
-    is(!!getHeartbeatNotification(flowId, privateWin), true,
-       "Heartbeat should appear in the most recent private window");
-    is(getHeartbeatNotification(flowId, mostRecentWin), null,
-       "Heartbeat shouldn't appear in the most recent non-private window");
+  let eventPromise = promisePageEvent();
+
+  gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
+                            null, null, {
+                              engagementButtonLabel: "Yes!",
+                              privateWindowsOnly: true,
+                            });
 
-    yield BrowserTestUtils.closeWindow(mostRecentWin);
-    yield BrowserTestUtils.closeWindow(privateWin);
-  }),
-
-  taskify(function* test_privateWindowsOnly() {
-    let engagementURL = "http://example.com";
-    let learnMoreURL = "http://example.org/learnmore/";
-    let flowId = "ui-privateWindowsOnly-" + Math.random();
+  yield eventPromise;
+  is(getHeartbeatNotification(flowId, window), null,
+     "Heartbeat shouldn't appear in the default window");
+  is(!!getHeartbeatNotification(flowId, privateWin), true,
+     "Heartbeat should appear in the most recent private window");
+  is(getHeartbeatNotification(flowId, mostRecentWin), null,
+     "Heartbeat shouldn't appear in the most recent non-private window");
 
-    let privateWin = yield BrowserTestUtils.openNewBrowserWindow({ private: true });
+  yield BrowserTestUtils.closeWindow(mostRecentWin);
+  yield BrowserTestUtils.closeWindow(privateWin);
+})
 
-    yield new Promise((resolve) => {
-      gContentAPI.observe(function(aEventName, aData) {
-        info(aEventName + " notification received: " + JSON.stringify(aData, null, 2));
-        ok(false, "No heartbeat notifications should arrive for privateWindowsOnly");
-      }, resolve);
-    });
+add_UITour_task(function* test_privateWindowsOnly() {
+  let engagementURL = "http://example.com";
+  let learnMoreURL = "http://example.org/learnmore/";
+  let flowId = "ui-privateWindowsOnly-" + Math.random();
 
-    gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
-                              "Learn More", learnMoreURL, {
-                                engagementButtonLabel: "Yes!",
-                                privateWindowsOnly: true,
-                              });
+  let privateWin = yield BrowserTestUtils.openNewBrowserWindow({ private: true });
 
-    yield promisePageEvent();
-
-    ok(isTourBrowser(gBrowser.selectedBrowser), "UITour should have been init for the browser");
+  yield new Promise((resolve) => {
+    gContentAPI.observe(function(aEventName, aData) {
+      info(aEventName + " notification received: " + JSON.stringify(aData, null, 2));
+      ok(false, "No heartbeat notifications should arrive for privateWindowsOnly");
+    }, resolve);
+  });
 
-    let notification = getHeartbeatNotification(flowId, privateWin);
-
-    is(notification.querySelectorAll(".star-x").length, 0, "No stars should be present");
+  gContentAPI.showHeartbeat("Do you want to engage with us?", "Thank you!", flowId, engagementURL,
+                            "Learn More", learnMoreURL, {
+                              engagementButtonLabel: "Yes!",
+                              privateWindowsOnly: true,
+                            });
 
-    info("Test the learn more link.");
-    let learnMoreLink = notification.querySelector(".text-link");
-    is(learnMoreLink.value, "Learn More", "Check learn more label");
-    let learnMoreTabPromise = BrowserTestUtils.waitForNewTab(privateWin.gBrowser, null);
-    learnMoreLink.click();
-    let learnMoreTab = yield learnMoreTabPromise;
-    is(learnMoreTab.linkedBrowser.currentURI.host, "example.org", "Check learn more site opened");
-    ok(PrivateBrowsingUtils.isBrowserPrivate(learnMoreTab.linkedBrowser), "Ensure the learn more tab is private");
-    yield BrowserTestUtils.removeTab(learnMoreTab);
+  yield promisePageEvent();
+
+  ok(isTourBrowser(gBrowser.selectedBrowser), "UITour should have been init for the browser");
+
+  let notification = getHeartbeatNotification(flowId, privateWin);
+
+  is(notification.querySelectorAll(".star-x").length, 0, "No stars should be present");
 
-    info("Test the engagement button's new tab.");
-    let engagementButton = notification.querySelector(".notification-button");
-    is(engagementButton.label, "Yes!", "Check engagement button text");
-    let engagementTabPromise = BrowserTestUtils.waitForNewTab(privateWin.gBrowser, null);
-    engagementButton.doCommand();
-    let engagementTab = yield engagementTabPromise;
-    is(engagementTab.linkedBrowser.currentURI.host, "example.com", "Check enagement site opened");
-    ok(PrivateBrowsingUtils.isBrowserPrivate(engagementTab.linkedBrowser), "Ensure the engagement tab is private");
-    yield BrowserTestUtils.removeTab(engagementTab);
+  info("Test the learn more link.");
+  let learnMoreLink = notification.querySelector(".text-link");
+  is(learnMoreLink.value, "Learn More", "Check learn more label");
+  let learnMoreTabPromise = BrowserTestUtils.waitForNewTab(privateWin.gBrowser, null);
+  learnMoreLink.click();
+  let learnMoreTab = yield learnMoreTabPromise;
+  is(learnMoreTab.linkedBrowser.currentURI.host, "example.org", "Check learn more site opened");
+  ok(PrivateBrowsingUtils.isBrowserPrivate(learnMoreTab.linkedBrowser), "Ensure the learn more tab is private");
+  yield BrowserTestUtils.removeTab(learnMoreTab);
 
-    yield BrowserTestUtils.closeWindow(privateWin);
-  }),
+  info("Test the engagement button's new tab.");
+  let engagementButton = notification.querySelector(".notification-button");
+  is(engagementButton.label, "Yes!", "Check engagement button text");
+  let engagementTabPromise = BrowserTestUtils.waitForNewTab(privateWin.gBrowser, null);
+  engagementButton.doCommand();
+  let engagementTab = yield engagementTabPromise;
+  is(engagementTab.linkedBrowser.currentURI.host, "example.com", "Check enagement site opened");
+  ok(PrivateBrowsingUtils.isBrowserPrivate(engagementTab.linkedBrowser), "Ensure the engagement tab is private");
+  yield BrowserTestUtils.removeTab(engagementTab);
 
-  /**
-   * Test that the survey closes itself after a while and submits Telemetry
-   */
-  taskify(function* test_telemetry_surveyExpired() {
-    let flowId = "survey-expired-" + Math.random();
-    let engagementURL = "http://example.com";
-    let surveyDuration = 1; // 1 second (pref is in seconds)
-    Services.prefs.setIntPref("browser.uitour.surveyDuration", surveyDuration);
+  yield BrowserTestUtils.closeWindow(privateWin);
+})
 
-    let telemetryPromise = new Promise((resolve, reject) => {
-        gContentAPI.observe(function (aEventName, aData) {
-          switch (aEventName) {
-            case "Heartbeat:NotificationOffered":
-              info("'Heartbeat:NotificationOffered' notification received");
-              break;
-            case "Heartbeat:SurveyExpired":
-              info("'Heartbeat:SurveyExpired' notification received");
-              ok(true, "Survey should end on its own after a time out");
-            case "Heartbeat:NotificationClosed":
-              info("'Heartbeat:NotificationClosed' notification received");
-              break;
-            case "Heartbeat:TelemetrySent": {
-              info("'Heartbeat:TelemetrySent' notification received");
-              checkTelemetry(aData, flowId, ["offeredTS", "expiredTS", "closedTS"]);
-              resolve();
-              break;
-            }
-            default:
-              // not expecting other states for this test
-              ok(false, "Unexpected notification received: " + aEventName);
-              reject();
-          }
-        });
-    });
+/**
+ * Test that the survey closes itself after a while and submits Telemetry
+ */
+add_UITour_task(function* test_telemetry_surveyExpired() {
+  let flowId = "survey-expired-" + Math.random();
+  let engagementURL = "http://example.com";
+  let surveyDuration = 1; // 1 second (pref is in seconds)
+  Services.prefs.setIntPref("browser.uitour.surveyDuration", surveyDuration);
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(["Heartbeat:NotificationOffered",
+    "Heartbeat:NotificationClosed", "Heartbeat:SurveyExpired", "Heartbeat:TelemetrySent"]);
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, engagementURL);
-    yield telemetryPromise;
-    Services.prefs.clearUserPref("browser.uitour.surveyDuration");
-  }),
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!", flowId, engagementURL);
+
+  let expiredPromise = promiseWaitHeartbeatNotification("Heartbeat:SurveyExpired");
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let pingPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
+
+  yield Promise.all([shownPromise, expiredPromise, closedPromise]);
+  // Validate the ping data.
+  let data = yield pingPromise;
+  checkTelemetry(data, flowId, ["offeredTS", "expiredTS", "closedTS"]);
+
+  Services.prefs.clearUserPref("browser.uitour.surveyDuration");
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
 
-  /**
-   * Check that certain whitelisted experiment parameters get reflected in the
-   * Telemetry ping
-   */
-  function test_telemetry_params(done) {
-    let flowId = "telemetry-params-" + Math.random();
-    let engagementURL = "http://example.com";
-    let extraParams = {
-      "surveyId": "foo",
-      "surveyVersion": 1.5,
-      "testing": true,
-      "notWhitelisted": 123,
-    };
-    let expectedFields = ["surveyId", "surveyVersion", "testing"];
+/**
+ * Check that certain whitelisted experiment parameters get reflected in the
+ * Telemetry ping
+ */
+add_UITour_task(function* test_telemetry_params() {
+  let flowId = "telemetry-params-" + Math.random();
+  let engagementURL = "http://example.com";
+  let extraParams = {
+    "surveyId": "foo",
+    "surveyVersion": 1.5,
+    "testing": true,
+    "notWhitelisted": 123,
+  };
+  let expectedFields = ["surveyId", "surveyVersion", "testing"];
+
+  // We need to call |gContentAPI.observe| at least once to set a valid |notificationListener|
+  // in UITour-lib.js, otherwise no message will get propagated.
+  gContentAPI.observe(() => {});
+
+  let receivedExpectedPromise = promiseWaitExpectedNotifications(
+    ["Heartbeat:NotificationOffered", "Heartbeat:NotificationClosed", "Heartbeat:TelemetrySent"]);
 
-    gContentAPI.observe(function (aEventName, aData) {
-      switch (aEventName) {
-        case "Heartbeat:NotificationOffered": {
-          info("'Heartbeat:Offered' notification received (timestamp " + aData.timestamp.toString() + ").");
-          cleanUpNotification(flowId);
-          break;
-        }
-        case "Heartbeat:NotificationClosed": {
-          info("'Heartbeat:NotificationClosed' notification received (timestamp " + aData.timestamp.toString() + ").");
-          break;
-        }
-        case "Heartbeat:TelemetrySent": {
-          info("'Heartbeat:TelemetrySent' notification received");
-          checkTelemetry(aData, flowId, ["offeredTS", "closedTS"].concat(expectedFields));
-          for (let param of expectedFields) {
-            is(aData[param], extraParams[param],
-               "Whitelisted experiment configs should be copied into Telemetry pings");
-          }
-          done();
-          break;
-        }
-        default:
-          // We are not expecting other states for this test.
-          ok(false, "Unexpected notification received: " + aEventName);
-      }
-    });
+  // Show the Heartbeat notification and wait for it to be displayed.
+  let shownPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationOffered");
+  gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!",
+                            flowId, engagementURL, null, null, extraParams);
+  yield shownPromise;
+
+  let closedPromise = promiseWaitHeartbeatNotification("Heartbeat:NotificationClosed");
+  let pingPromise = promiseWaitHeartbeatNotification("Heartbeat:TelemetrySent");
+  cleanUpNotification(flowId);
 
-    gContentAPI.showHeartbeat("How would you rate Firefox?", "Thank you!",
-                              flowId, engagementURL, null, null, extraParams);
-  },
-];
+  // The notification was closed.
+  let data = yield closedPromise;
+  validateTimestamp('Heartbeat:NotificationClosed', data.timestamp);
+
+  // Validate the data we send out.
+  data = yield pingPromise;
+  info("'Heartbeat:TelemetrySent' notification received.");
+  checkTelemetry(data, flowId, ["offeredTS", "closedTS"].concat(expectedFields));
+  for (let param of expectedFields) {
+    is(data[param], extraParams[param],
+       "Whitelisted experiment configs should be copied into Telemetry pings");
+  }
+
+  // This rejects whenever an unexpected notification is received.
+  yield receivedExpectedPromise;
+})
--- a/browser/components/uitour/test/head.js
+++ b/browser/components/uitour/test/head.js
@@ -283,16 +283,19 @@ function loadUITourTestPage(callback, ho
               // forwarded functions. We'll construct a function on the content-side
               // that forwards all its arguments to a message, and we'll listen for
               // those messages on our side and call the corresponding function with
               // the arguments we got from the content side.
               if (typeof arg == "function") {
                 callbackMap.set(index, arg);
                 fnIndices.push(index);
                 let handler = function(msg) {
+                  // Please note that this handler assumes that the callback is used only once.
+                  // That means that a single gContentAPI.observer() call can't be used to observe
+                  // multiple events.
                   browser.messageManager.removeMessageListener(proxyFunctionName + index, handler);
                   callbackMap.get(index).apply(null, msg.data);
                 };
                 browser.messageManager.addMessageListener(proxyFunctionName + index, handler);
                 return "";
               }
               return arg;
             });