1 files changed, 233 insertions, 0 deletions
diff --git a/browser/components/search/test/browser/head.js b/browser/components/search/test/browser/head.js
new file mode 100644
index 0000000000..97615b0dad
--- /dev/null
+++ b/browser/components/search/test/browser/head.js
@@ -0,0 +1,233 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+XPCOMUtils.defineLazyModuleGetters(this, {
+  ADLINK_CHECK_TIMEOUT_MS: "resource:///actors/SearchSERPTelemetryChild.jsm",
+  AddonTestUtils: "resource://testing-common/AddonTestUtils.jsm",
+  CustomizableUITestUtils:
+    "resource://testing-common/CustomizableUITestUtils.jsm",
+  FormHistoryTestUtils: "resource://testing-common/FormHistoryTestUtils.jsm",
+  PrivateBrowsingUtils: "resource://gre/modules/PrivateBrowsingUtils.jsm",
+  SearchTestUtils: "resource://testing-common/SearchTestUtils.jsm",
+  SearchUtils: "resource://gre/modules/SearchUtils.jsm",
+  TelemetryTestUtils: "resource://testing-common/TelemetryTestUtils.jsm",
+  UrlbarSearchUtils: "resource:///modules/UrlbarSearchUtils.jsm",
+});
+
+let gCUITestUtils = new CustomizableUITestUtils(window);
+
+AddonTestUtils.initMochitest(this);
+SearchTestUtils.init(this);
+
+/**
+ * Recursively compare two objects and check that every property of expectedObj has the same value
+ * on actualObj.
+ *
+ * @param {object} expectedObj
+ *        The expected object to find.
+ * @param {object} actualObj
+ *        The object to inspect.
+ * @param {string} name
+ *        The name of the engine, used for test detail logging.
+ */
+function isSubObjectOf(expectedObj, actualObj, name) {
+  for (let prop in expectedObj) {
+    if (typeof expectedObj[prop] == "function") {
+      continue;
+    }
+    if (expectedObj[prop] instanceof Object) {
+      is(
+        actualObj[prop].length,
+        expectedObj[prop].length,
+        name + "[" + prop + "]"
+      );
+      isSubObjectOf(
+        expectedObj[prop],
+        actualObj[prop],
+        name + "[" + prop + "]"
+      );
+    } else {
+      is(actualObj[prop], expectedObj[prop], name + "[" + prop + "]");
+    }
+  }
+}
+
+function getLocale() {
+  return Services.locale.requestedLocale || undefined;
+}
+
+function promiseEvent(aTarget, aEventName, aPreventDefault) {
+  function cancelEvent(event) {
+    if (aPreventDefault) {
+      event.preventDefault();
+    }
+
+    return true;
+  }
+
+  return BrowserTestUtils.waitForEvent(aTarget, aEventName, false, cancelEvent);
+}
+
+/**
+ * Adds a new search engine to the search service and confirms it completes.
+ *
+ * @param {string} basename  The file to load that contains the search engine
+ *                           details.
+ * @param {object} [options] Options for the test:
+ *   - {String} [iconURL]       The icon to use for the search engine.
+ *   - {Boolean} [setAsCurrent] Whether to set the new engine to be the
+ *                              current engine or not.
+ *   - {Boolean} [setAsCurrentPrivate] Whether to set the new engine to be the
+ *                              current private browsing mode engine or not.
+ *                              Defaults to false.
+ *   - {String} [testPath]      Used to override the current test path if this
+ *                              file is used from a different directory.
+ * @returns {Promise} The promise is resolved once the engine is added, or
+ *                    rejected if the addition failed.
+ */
+async function promiseNewEngine(basename, options = {}) {
+  // Default the setAsCurrent option to true.
+  let setAsCurrent =
+    options.setAsCurrent == undefined ? true : options.setAsCurrent;
+  info("Waiting for engine to be added: " + basename);
+  let url = getRootDirectory(options.testPath || gTestPath) + basename;
+  let engine = await Services.search.addOpenSearchEngine(
+    url,
+    options.iconURL || ""
+  );
+  info("Search engine added: " + basename);
+  const current = await Services.search.getDefault();
+  if (setAsCurrent) {
+    await Services.search.setDefault(engine);
+  }
+  const currentPrivate = await Services.search.getDefaultPrivate();
+  if (options.setAsCurrentPrivate) {
+    await Services.search.setDefaultPrivate(engine);
+  }
+  registerCleanupFunction(async () => {
+    if (setAsCurrent) {
+      await Services.search.setDefault(current);
+    }
+    if (options.setAsCurrentPrivate) {
+      await Services.search.setDefaultPrivate(currentPrivate);
+    }
+    await Services.search.removeEngine(engine);
+    info("Search engine removed: " + basename);
+  });
+  return engine;
+}
+
+// Get an array of the one-off buttons.
+function getOneOffs() {
+  let oneOffs = [];
+  let searchPopup = document.getElementById("PopupSearchAutoComplete");
+  let oneOffsContainer = searchPopup.searchOneOffsContainer;
+  let oneOff = oneOffsContainer.querySelector(".search-panel-one-offs");
+  for (oneOff = oneOff.firstChild; oneOff; oneOff = oneOff.nextSibling) {
+    if (oneOff.nodeType == Node.ELEMENT_NODE) {
+      oneOffs.push(oneOff);
+    }
+  }
+  return oneOffs;
+}
+
+async function typeInSearchField(browser, text, fieldName) {
+  await SpecialPowers.spawn(browser, [[fieldName, text]], async function([
+    contentFieldName,
+    contentText,
+  ]) {
+    // Put the focus on the search box.
+    let searchInput = content.document.getElementById(contentFieldName);
+    searchInput.focus();
+    searchInput.value = contentText;
+  });
+}
+
+XPCOMUtils.defineLazyGetter(this, "searchCounts", () => {
+  return Services.telemetry.getKeyedHistogramById("SEARCH_COUNTS");
+});
+
+XPCOMUtils.defineLazyGetter(this, "SEARCH_AD_CLICK_SCALARS", () => {
+  const sources = [
+    ...BrowserSearchTelemetry.KNOWN_SEARCH_SOURCES.values(),
+    "unknown",
+  ];
+  return [
+    "browser.search.with_ads",
+    "browser.search.ad_clicks",
+    ...sources.map(v => `browser.search.withads.${v}`),
+    ...sources.map(v => `browser.search.adclicks.${v}`),
+  ];
+});
+
+// Ad links are processed after a small delay. We need to allow tests to wait
+// for that before checking telemetry, otherwise the received values may be
+// too small in some cases.
+function promiseWaitForAdLinkCheck() {
+  return new Promise(resolve =>
+    /* eslint-disable-next-line mozilla/no-arbitrary-setTimeout */
+    setTimeout(resolve, ADLINK_CHECK_TIMEOUT_MS)
+  );
+}
+
+async function assertSearchSourcesTelemetry(
+  expectedHistograms,
+  expectedScalars
+) {
+  let histSnapshot = {};
+  let scalars = {};
+
+  await TestUtils.waitForCondition(() => {
+    histSnapshot = searchCounts.snapshot();
+    return (
+      Object.getOwnPropertyNames(histSnapshot).length ==
+      Object.getOwnPropertyNames(expectedHistograms).length
+    );
+  }, "should have the correct number of histograms");
+
+  if (Object.entries(expectedScalars).length) {
+    await TestUtils.waitForCondition(() => {
+      scalars =
+        Services.telemetry.getSnapshotForKeyedScalars("main", false).parent ||
+        {};
+      return Object.getOwnPropertyNames(expectedScalars).every(
+        scalar => scalar in scalars
+      );
+    }, "should have the expected keyed scalars");
+  }
+
+  Assert.equal(
+    Object.getOwnPropertyNames(histSnapshot).length,
+    Object.getOwnPropertyNames(expectedHistograms).length,
+    "Should only have one key"
+  );
+
+  for (let [key, value] of Object.entries(expectedHistograms)) {
+    Assert.ok(
+      key in histSnapshot,
+      `Histogram should have the expected key: ${key}`
+    );
+    Assert.equal(
+      histSnapshot[key].sum,
+      value,
+      `Should have counted the correct number of visits for ${key}`
+    );
+  }
+
+  for (let [name, value] of Object.entries(expectedScalars)) {
+    Assert.ok(name in scalars, `Scalar ${name} should have been added.`);
+    Assert.deepEqual(
+      scalars[name],
+      value,
+      `Should have counted the correct number of visits for ${name}`
+    );
+  }
+
+  for (let name of SEARCH_AD_CLICK_SCALARS) {
+    Assert.equal(
+      name in scalars,
+      name in expectedScalars,
+      `Should have matched ${name} in scalars and expectedScalars`
+    );
+  }
+}