Bug 1647890 - Replace token alias with indicator when the token alias is confirmed or fully typed. r=adw

This patch changes the paramters to setSearchMode. The original patch to introduce search mode passed engineName to setSearchMode instead of the entire engine object. It was suggested in review that the nsISearchEngine be passed so we could run instanceof to distinguish it from a RESULT_TYPE. This patch reverses this and passes engineName instead through a destructured parameter.

In pickResult, we need to enter search mode synchronously based on the information in a result payload. Result payloads can't/shouldn't pass around complex objects like an nsISearchEngine, so we just pass engineName and the alias as strings. Since pickResult is synchronous, we can't use UrlbarSearchUtils to look up the engine based on the engineName. Besides, setSearchMode only uses engineName, so looking up the engine only to just use its name seems like a waste of resources.

This patch also disables autofill in search mode queries. Autofill was interfering with alias replacement. We were already half doing this (https://searchfox.org/mozilla-central/rev/26b13464c2beb26e0d864d561c30e817a85c348a/browser/components/urlbar/UrlbarController.jsm#391) but adding the searchMode check to UrlbarInput._maybeAutofillOnInput should resolve bug 1655473.

There's still one more bug I'm working through where the placeholder disappears after alias replacement. I though I'd get this out to start review regardless since we want to get the three patches discussed in Thursday's meeting out ASAP.

Differential Revision: https://phabricator.services.mozilla.com/D86389

browser/components/urlbar/UrlbarController.jsm

@@ -385,7 +385,7 @@ class UrlbarController {
           this.input.selectionEnd == 0 &&
           !event.shiftKey
         ) {
-          this.input.setSearchMode(null);
+          this.input.setSearchMode({});
           if (this.input.value) {
             this.input.startQuery({
               allowAutofill: false,

browser/components/urlbar/UrlbarInput.jsm

@@ -16,7 +16,6 @@ XPCOMUtils.defineLazyModuleGetters(this, {
   ExtensionSearchHandler: "resource://gre/modules/ExtensionSearchHandler.jsm",
   PrivateBrowsingUtils: "resource://gre/modules/PrivateBrowsingUtils.jsm",
   ReaderMode: "resource://gre/modules/ReaderMode.jsm",
-  SearchEngine: "resource://gre/modules/SearchEngine.jsm",
   Services: "resource://gre/modules/Services.jsm",
   TopSiteAttribution: "resource:///modules/TopSiteAttribution.jsm",
   UrlbarController: "resource:///modules/UrlbarController.jsm",
@@ -424,10 +423,10 @@ class UrlbarInput {
         return;
       }
       if (selectedOneOff && this.view.oneOffsRefresh) {
-        this.view.oneOffSearchButtons.handleSearchCommand(
-          event,
-          selectedOneOff.engine || selectedOneOff.source
-        );
+        this.view.oneOffSearchButtons.handleSearchCommand(event, {
+          engineName: selectedOneOff.engine?.name,
+          source: selectedOneOff.source,
+        });
         return;
       }
     }
@@ -568,7 +567,7 @@ class UrlbarInput {
   handleRevert() {
     this.window.gBrowser.userTypedValue = null;
     this.setURI(null, true);
-    this.setSearchMode(null);
+    this.setSearchMode({});
     if (this.value && this.focused) {
       this.select();
     }
@@ -702,22 +701,29 @@ class UrlbarInput {
       }
       case UrlbarUtils.RESULT_TYPE.SEARCH: {
         if (result.payload.keywordOffer) {
-          // The user confirmed a token alias, so just move the caret
-          // to the end of it. Because there's a trailing space in the value,
-          // the user can directly start typing a query string at that point.
-          this.selectionStart = this.selectionEnd = this.value.length;
-
           this.controller.engagementEvent.record(event, {
             searchString: this._lastSearchString,
             selIndex,
             selType: "keywordoffer",
           });
 
-          // Picking a keyword offer just fills it in the input and doesn't
-          // visit anything.  The user can then type a search string.  Also
-          // start a new search so that the offer appears in the view by itself
-          // to make it even clearer to the user what's going on.
-          this.startQuery();
+          let searchModeParams = this._searchModeForResult(result);
+          if (searchModeParams) {
+            this.setSearchMode(searchModeParams);
+            this.search("");
+          } else {
+            // The user confirmed a token alias, so just move the caret
+            // to the end of it. Because there's a trailing space in the value,
+            // the user can directly start typing a query string at that point.
+            this.selectionStart = this.selectionEnd = this.value.length;
+
+            // Picking a keyword offer just fills it in the input and doesn't
+            // visit anything.  The user can then type a search string.  Also
+            // start a new search so that the offer appears in the view by itself
+            // to make it even clearer to the user what's going on.
+            this.startQuery();
+          }
+
           return;
         }
 
@@ -879,7 +885,7 @@ class UrlbarInput {
    */
   onPrefChanged(changedPref) {
     if (changedPref == "update2" && !UrlbarPrefs.get("update2")) {
-      this.setSearchMode(null);
+      this.setSearchMode({});
     }
   }
 
@@ -1171,23 +1177,19 @@ class UrlbarInput {
    *
    * @param {*} value
    *   A value of one of the following types:
-   *
-   *   nsISearchEngine|SearchEngine
-   *     The search mode source will be set to UrlbarUtils.RESULT_SOURCE.SEARCH
-   *     using the given engine.
-   *   UrlbarUtils.RESULT_SOURCE
-   *     The search mode source will be set to the given source.
-   *   search mode object
-   *     The search mode will be set to this object.  It should be a valid
-   *     search mode object derived from one of the other types above.  This is
-   *     useful for re-entering previous search modes.
-   *   null (or any falsey value)
-   *     Search mode will be exited.
+   * @param {string} engineName
+   *   The name of the search engine to restrict to.
+   * @param {UrlbarUtils.RESULT_SOURCE} source
+   *   A result source to restrict to.
+   * @param {string} [alternateLabel]
+   *   Optional. If provided, this string will be shown in the search mode
+   *   indicator instead of the engine. Does not override a source title.
    */
-  setSearchMode(value) {
+  setSearchMode({ engineName, source, alternateLabel }) {
     if (!UrlbarPrefs.get("update2")) {
       // Exit search mode.
-      value = null;
+      engineName = null;
+      source = null;
     }
 
     this._searchModeIndicatorTitle.textContent = "";
@@ -1195,36 +1197,21 @@ class UrlbarInput {
     this._searchModeIndicatorTitle.removeAttribute("data-l10n-id");
     this._searchModeLabel.removeAttribute("data-l10n-id");
 
-    // First, set this.searchMode based on `value`.
-    if (!value) {
-      this.searchMode = null;
-    } else if (
-      value instanceof Ci.nsISearchEngine ||
-      value instanceof SearchEngine
-    ) {
+    if (engineName) {
       this.searchMode = {
         source: UrlbarUtils.RESULT_SOURCE.SEARCH,
-        engineName: value.name,
+        engineName,
+        alternateLabel,
       };
-    } else if (typeof value == "number") {
-      this.searchMode = { source: value };
-    } else if (typeof value == "object") {
-      this.searchMode = value;
-    } else {
-      Cu.reportError(`Unexpected search mode value: ${value}`);
-      this.searchMode = null;
-    }
-
-    // Now set the indicator and label strings based on this.searchMode.
-    if (this.searchMode?.engineName) {
-      this._searchModeIndicatorTitle.textContent = this.searchMode.engineName;
-      this._searchModeLabel.textContent = this.searchMode.engineName;
-    } else if (this.searchMode?.source) {
-      let sourceName = UrlbarUtils.getResultSourceName(this.searchMode.source);
+      this._searchModeIndicatorTitle.textContent = alternateLabel || engineName;
+      this._searchModeLabel.textContent = alternateLabel || engineName;
+    } else if (source) {
+      let sourceName = UrlbarUtils.getResultSourceName(source);
       if (!sourceName) {
-        Cu.reportError(`Unrecognized source: ${this.searchMode.source}`);
+        Cu.reportError(`Unrecognized source: ${source}`);
         this.searchMode = null;
       } else {
+        this.searchMode = { source };
         let l10nID = `urlbar-search-mode-${sourceName}`;
         this.document.l10n.setAttributes(
           this._searchModeIndicatorTitle,
@@ -1232,8 +1219,8 @@ class UrlbarInput {
         );
         this.document.l10n.setAttributes(this._searchModeLabel, l10nID);
       }
-    } else if (this.searchMode) {
-      Cu.reportError(`Unexpected search mode object: ${this.searchMode}`);
+    } else {
+      // Exit search mode.
       this.searchMode = null;
     }
 
@@ -1261,8 +1248,7 @@ class UrlbarInput {
    */
   searchModeShortcut() {
     if (this.view.oneOffsRefresh) {
-      let defaultEngine = Services.search.defaultEngine;
-      this.setSearchMode(defaultEngine);
+      this.setSearchMode({ engineName: Services.search.defaultEngine.name });
       this.search("");
     } else {
       this.search(UrlbarTokenizer.RESTRICT.SEARCH);
@@ -1459,7 +1445,7 @@ class UrlbarInput {
     let searchMode = this._searchModesByBrowser.get(
       this.window.gBrowser.selectedBrowser
     );
-    this.setSearchMode(searchMode);
+    this.setSearchMode(searchMode || {});
 
     // Switching tabs doesn't always change urlbar focus, so we must try to
     // reopen here too, not just on focus.
@@ -1542,8 +1528,9 @@ class UrlbarInput {
         return result.payload.input;
       case UrlbarUtils.RESULT_TYPE.SEARCH:
         return (
-          (result.payload.keyword ? result.payload.keyword + " " : "") +
-          (result.payload.suggestion || result.payload.query)
+          (result.payload.keyword && !UrlbarPrefs.get("update2")
+            ? result.payload.keyword + " "
+            : "") + (result.payload.suggestion || result.payload.query)
         );
       case UrlbarUtils.RESULT_TYPE.OMNIBOX:
         return result.payload.content;
@@ -1579,7 +1566,7 @@ class UrlbarInput {
    *   Whether autofill should be allowed in the new search.
    */
   _maybeAutofillOnInput(value) {
-    let allowAutofill = this.selectionEnd == value.length;
+    let allowAutofill = this.selectionEnd == value.length && !this.searchMode;
 
     // Determine whether we can autofill the placeholder.  The placeholder is a
     // value that we autofill now, when the search starts and before we wait on
@@ -1988,7 +1975,7 @@ class UrlbarInput {
     // area when the current tab is re-selected.
     browser.focus();
 
-    this.setSearchMode(null);
+    this.setSearchMode({});
 
     if (openUILinkWhere != "current") {
       this.handleRevert();
@@ -2128,6 +2115,37 @@ class UrlbarInput {
     Services.obs.notifyObservers({ result }, "urlbar-user-start-navigation");
   }
 
+  /**
+   * Returns a search mode object if a result should enter search mode when
+   * selected.
+   *
+   * @param {UrlbarResult} result
+   * @returns {object} A search mode object. Null if search mode should not be
+   *   entered. See setSearchMode documentation for details.
+   */
+  _searchModeForResult(result) {
+    if (!UrlbarPrefs.get("update2")) {
+      return null;
+    }
+
+    // If result.originalEngine is set, then the user is Alt+Tabbing through the
+    // one-offs, so the keyword doesn't match the engine.
+    if (
+      result &&
+      result.payload.keywordOffer &&
+      (!result.payload.originalEngine ||
+        result.payload.engine == result.payload.originalEngine)
+    ) {
+      let params = { engineName: result.payload.engine };
+      if (result.payload.keyword && !result.payload.keyword.startsWith("@")) {
+        params.alternateLabel = result.payload.keyword;
+      }
+      return params;
+    }
+
+    return null;
+  }
+
   /**
    * Determines if we should select all the text in the Urlbar based on the
    *  Urlbar state, and whether the selection is empty.
@@ -2222,7 +2240,7 @@ class UrlbarInput {
     }
 
     if (event.target == this._searchModeIndicatorClose && event.button != 2) {
-      this.setSearchMode(null);
+      this.setSearchMode({});
       if (this.view.isOpen) {
         this.startQuery({
           event,
@@ -2357,6 +2375,23 @@ class UrlbarInput {
   }
 
   _on_input(event) {
+    // We enter search mode when space is typed if there is a selected keyword
+    // offer result.
+    let searchModeParams;
+    if (event.data == " ") {
+      let result = this.view.selectedResult;
+      searchModeParams = this._searchModeForResult(result);
+      if (
+        searchModeParams &&
+        this.value.trim() == result.payload.keyword.trim()
+      ) {
+        this.setSearchMode(searchModeParams);
+        this.value = "";
+      } else {
+        searchModeParams = null;
+      }
+    }
+
     let value = this.value;
     this.valueIsTyped = true;
     this._untrimmedValue = value;
@@ -2393,10 +2428,13 @@ class UrlbarInput {
       !this.isPrivate &&
       UrlbarPrefs.get("suggest.topsites") &&
       !this.searchMode;
-    if (!this.view.isOpen || (!value && !canShowTopSites)) {
+    if (
+      !this.view.isOpen ||
+      (!value && !canShowTopSites && !searchModeParams)
+    ) {
       this.view.clear();
     }
-    if (!value && !canShowTopSites) {
+    if (!value && !canShowTopSites && !searchModeParams) {
       this.view.close();
       return;
     }
@@ -2429,7 +2467,7 @@ class UrlbarInput {
     this.startQuery({
       searchString: value,
       allowAutofill,
-      resetSearchState: !!this.searchMode,
+      resetSearchState: false,
       event,
     });
   }

browser/components/urlbar/UrlbarProviderTokenAliasEngines.jsm

@@ -165,7 +165,10 @@ class ProviderTokenAliasEngines extends UrlbarProvider {
           // We found a specific engine. We will add an autofill result.
           let aliasPreservingUserCase =
             token.value + alias.substr(token.value.length);
-          let value = aliasPreservingUserCase + " ";
+          // Don't append a space if update2 is on since selecting this result
+          // will just enter search mode.
+          let value =
+            aliasPreservingUserCase + (UrlbarPrefs.get("update2") ? "" : " ");
           let result = new UrlbarResult(
             UrlbarUtils.RESULT_TYPE.SEARCH,
             UrlbarUtils.RESULT_SOURCE.SEARCH,

browser/components/urlbar/UrlbarProviderUnifiedComplete.jsm

@@ -18,6 +18,7 @@ const { XPCOMUtils } = ChromeUtils.import(
 XPCOMUtils.defineLazyModuleGetters(this, {
   PlacesUtils: "resource://gre/modules/PlacesUtils.jsm",
   Services: "resource://gre/modules/Services.jsm",
+  UrlbarPrefs: "resource:///modules/UrlbarPrefs.jsm",
   UrlbarProvider: "resource:///modules/UrlbarUtils.jsm",
   UrlbarResult: "resource:///modules/UrlbarResult.jsm",
   UrlbarUtils: "resource:///modules/UrlbarUtils.jsm",
@@ -183,7 +184,7 @@ function makeUrlbarResult(tokens, info) {
         if (
           action.params.alias &&
           !action.params.searchQuery.trim() &&
-          action.params.alias.startsWith("@")
+          (UrlbarPrefs.get("update2") || action.params.alias.startsWith("@"))
         ) {
           keywordOffer = info.isHeuristic
             ? UrlbarUtils.KEYWORD_OFFER.HIDE

browser/components/urlbar/UrlbarSearchOneOffs.jsm

@@ -156,20 +156,20 @@ class UrlbarSearchOneOffs extends SearchOneOffs {
    *
    * @param {event} event
    *   The event that triggered the pick.
-   * @param {nsISearchEngine|SearchEngine|UrlbarUtils.RESULT_SOURCE} engineOrSource
-   *   The engine that was picked, or for local search mode sources, the source
-   *   that was picked as a UrlbarUtils.RESULT_SOURCE value.
+   * @param {object} searchMode
+   *   Used by UrlbarInput.setSearchMode to enter search mode. See setSearchMode
+   *   documentation for details.
    * @param {boolean} forceNewTab
    *   True if the search results page should be loaded in a new tab.
    */
-  handleSearchCommand(event, engineOrSource, forceNewTab = false) {
+  handleSearchCommand(event, searchMode, forceNewTab = false) {
     if (!this.view.oneOffsRefresh) {
       let { where, params } = this._whereToOpen(event, forceNewTab);
       this.input.handleCommand(event, where, params);
       return;
     }
 
-    this.input.setSearchMode(engineOrSource);
+    this.input.setSearchMode(searchMode);
     this.selectedButton = null;
     this.input.startQuery({
       allowAutofill: false,
@@ -267,11 +267,13 @@ class UrlbarSearchOneOffs extends SearchOneOffs {
     }
 
     let button = event.originalTarget;
-    let engineOrSource = button.engine || button.source;
-    if (!engineOrSource) {
+    if (!button.engine && !button.source) {
       return;
     }
 
-    this.handleSearchCommand(event, engineOrSource);
+    this.handleSearchCommand(event, {
+      engineName: button.engine?.name,
+      source: button.source,
+    });
   }
 }

browser/components/urlbar/tests/UrlbarTestUtils.jsm

@@ -390,6 +390,12 @@ var UrlbarTestUtils = {
       return;
     }
 
+    // Since alternateLabel is only used in limited circumstances, set it to
+    // undefined here to avoid every consumer having to do this.
+    if (expectedSearchMode.engineName && !expectedSearchMode.alternateLabel) {
+      expectedSearchMode.alternateLabel = undefined;
+    }
+
     this.Assert.deepEqual(
       window.gURLBar.searchMode,
       expectedSearchMode,
@@ -400,7 +406,11 @@ var UrlbarTestUtils = {
     let expectedTextContent = "";
     let expectedL10n = {};
 
-    if (expectedSearchMode.engineName) {
+    if (expectedSearchMode.alternateLabel) {
+      expectedTextContent = expectedSearchMode.alternateLabel;
+    } else if (expectedSearchMode.engineDisplayName) {
+      expectedTextContent = expectedSearchMode.engineDisplayName;
+    } else if (expectedSearchMode.engineName) {
       expectedTextContent = expectedSearchMode.engineName;
     } else if (expectedSearchMode.source) {
       let name = UrlbarUtils.getResultSourceName(expectedSearchMode.source);

browser/components/urlbar/tests/browser/browser.ini

@@ -137,6 +137,10 @@ run-if = e10s
 [browser_revert.js]
 [browser_searchFunction.js]
 [browser_searchMode_suggestions.js]
+support-files =
+  searchSuggestionEngine.xml
+  searchSuggestionEngine.sjs
+[browser_searchMode_alias_replacement.js]
 support-files =
   searchSuggestionEngine.xml
   searchSuggestionEngine.sjs

browser/components/urlbar/tests/browser/browser_oneOffs.js

@@ -471,7 +471,7 @@ add_task(async function oneOffClick() {
         source: UrlbarUtils.RESULT_SOURCE.SEARCH,
         engineName: oneOffs[0].engine.name,
       });
-      window.gURLBar.setSearchMode(null);
+      window.gURLBar.setSearchMode({});
     } else {
       let resultsPromise = BrowserTestUtils.browserLoaded(
         gBrowser.selectedBrowser,
@@ -522,7 +522,7 @@ add_task(async function oneOffReturn() {
         source: UrlbarUtils.RESULT_SOURCE.SEARCH,
         engineName: oneOffs[0].engine.name,
       });
-      window.gURLBar.setSearchMode(null);
+      window.gURLBar.setSearchMode({});
     } else {
       let resultsPromise = BrowserTestUtils.browserLoaded(
         gBrowser.selectedBrowser,
@@ -691,7 +691,7 @@ add_task(async function localOneOffClick() {
     });
   }
 
-  window.gURLBar.setSearchMode(null);
+  window.gURLBar.setSearchMode({});
 
   await hidePopup();
   await SpecialPowers.popPrefEnv();
@@ -756,7 +756,7 @@ add_task(async function localOneOffReturn() {
     });
   }
 
-  window.gURLBar.setSearchMode(null);
+  window.gURLBar.setSearchMode({});
 
   await hidePopup();
   await SpecialPowers.popPrefEnv();

browser/components/urlbar/tests/browser/browser_searchMode_alias_replacement.js

@@ -0,0 +1,136 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+/**
+ * Tests that user-defined aliases are replaced by the search mode indicator.
+ */
+
+const ALIAS = "testalias";
+const TEST_ENGINE_BASENAME = "searchSuggestionEngine.xml";
+
+let defaultEngine, aliasEngine;
+
+add_task(async function setup() {
+  defaultEngine = await SearchTestUtils.promiseNewSearchEngine(
+    getRootDirectory(gTestPath) + TEST_ENGINE_BASENAME
+  );
+  defaultEngine.alias = "@default";
+  let oldDefaultEngine = await Services.search.getDefault();
+  Services.search.setDefault(defaultEngine);
+  aliasEngine = await Services.search.addEngineWithDetails("Test", {
+    alias: ALIAS,
+    template: "http://example.com/?search={searchTerms}",
+  });
+
+  registerCleanupFunction(async function() {
+    await Services.search.removeEngine(aliasEngine);
+    Services.search.setDefault(oldDefaultEngine);
+  });
+
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.urlbar.update2", true]],
+  });
+});
+
+// Tests that a fully typed alias is replaced when space is pressed.
+add_task(async function replaced_on_space() {
+  // Check that a non-fully typed alias is not replaced.
+  await UrlbarTestUtils.promiseAutocompleteResultPopup({
+    window,
+    value: ALIAS.slice(0, -1),
+  });
+
+  let searchPromise = UrlbarTestUtils.promiseSearchComplete(window);
+  EventUtils.synthesizeKey("VK_SPACE");
+  await searchPromise;
+
+  UrlbarTestUtils.assertSearchMode(window, null);
+  Assert.equal(
+    gURLBar.value,
+    ALIAS.slice(0, -1),
+    "The typed value should be unchanged."
+  );
+  await UrlbarTestUtils.promisePopupClose(window, () =>
+    EventUtils.synthesizeKey("KEY_Escape")
+  );
+
+  // Test that the alias is replaced when it is fully typed.
+  await UrlbarTestUtils.promiseAutocompleteResultPopup({
+    window,
+    value: `${ALIAS} `,
+  });
+  let keywordOfferResult = await UrlbarTestUtils.getDetailsOfResultAt(
+    window,
+    0
+  );
+  Assert.equal(
+    keywordOfferResult.searchParams.keyword,
+    ALIAS,
+    "The first result should be a keyword search result with the correct alias."
+  );
+  Assert.equal(
+    keywordOfferResult.searchParams.engine,
+    aliasEngine.name,
+    "The first result should be a keyword search result with the correct engine."
+  );
+
+  searchPromise = UrlbarTestUtils.promiseSearchComplete(window);
+  // Fire an input event simulating typing a space after the ALIAS.
+  UrlbarTestUtils.fireInputEvent(window);
+  await searchPromise;
+
+  UrlbarTestUtils.assertSearchMode(window, {
+    source: UrlbarUtils.RESULT_SOURCE.SEARCH,
+    engineName: aliasEngine.name,
+    alternateLabel: ALIAS,
+  });
+  Assert.ok(!gURLBar.value, "The Urlbar value should be cleared.");
+  gURLBar.setSearchMode({});
+  await UrlbarTestUtils.promisePopupClose(window);
+});
+
+add_task(async function not_replaced_for_alt_tab() {
+  // Test that the alias is replaced when it is fully typed.
+  await UrlbarTestUtils.promiseAutocompleteResultPopup({
+    window,
+    value: `${ALIAS} `,
+  });
+  let keywordOfferResult = await UrlbarTestUtils.getDetailsOfResultAt(
+    window,
+    0
+  );
+  Assert.equal(
+    keywordOfferResult.searchParams.keyword,
+    ALIAS,
+    "The first result should be a keyword search result with the correct alias."
+  );
+  Assert.equal(
+    keywordOfferResult.searchParams.engine,
+    aliasEngine.name,
+    "The first result should be a keyword search result with the correct engine."
+  );
+
+  EventUtils.synthesizeKey("KEY_ArrowDown", { altKey: true, repeat: 2 });
+  keywordOfferResult = await UrlbarTestUtils.waitForAutocompleteResultAt(
+    window,
+    0
+  );
+  Assert.ok(
+    keywordOfferResult.result.payload.originalEngine,
+    "The keyword offer result now has the originalEngine property."
+  );
+  Assert.notEqual(
+    keywordOfferResult.result.payload.engine,
+    keywordOfferResult.result.payload.originalEngine,
+    "engine and originalEngine are different."
+  );
+
+  let searchPromise = UrlbarTestUtils.promiseSearchComplete(window);
+  // Fire an input event simulating typing a space after the ALIAS.
+  UrlbarTestUtils.fireInputEvent(window);
+  await searchPromise;
+
+  UrlbarTestUtils.assertSearchMode(window, null);
+});

browser/components/urlbar/tests/browser/browser_tokenAlias.js

@@ -219,8 +219,9 @@ add_task(async function nonTokenAlias() {
 });
 
 // Clicking on an @ alias offer (an @ alias with an empty search string) in the
-// popup should fill it in the urlbar input.
-add_task(async function clickAndFillAlias() {
+// view should fill it in the urlbar input.
+// This subtest can be removed when update2 is on by default.
+add_task(async function clickAndFillAlias_legacy() {
   // Do a search for "@" to show all the @ aliases.
   gURLBar.search("@");
   await UrlbarTestUtils.promiseSearchComplete(window);
@@ -267,9 +268,50 @@ add_task(async function clickAndFillAlias() {
   );
 });
 
+// Clicking on an @ alias offer (an @ alias with an empty search string) in the
+// view should enter search mode.
+add_task(async function clickAndFillAlias() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.urlbar.update2", true]],
+  });
+
+  // Do a search for "@" to show all the @ aliases.
+  gURLBar.search("@");
+  await UrlbarTestUtils.promiseSearchComplete(window);
+
+  // Find our test engine in the results.  It's probably last, but for
+  // robustness don't assume it is.
+  let testEngineItem;
+  for (let i = 0; !testEngineItem; i++) {
+    let details = await UrlbarTestUtils.getDetailsOfResultAt(window, i);
+    if (details.searchParams && details.searchParams.keyword == ALIAS) {
+      testEngineItem = await UrlbarTestUtils.waitForAutocompleteResultAt(
+        window,
+        i
+      );
+    }
+  }
+
+  // Click it.
+  EventUtils.synthesizeMouseAtCenter(testEngineItem, {});
+
+  UrlbarTestUtils.assertSearchMode(window, {
+    source: UrlbarUtils.RESULT_SOURCE.SEARCH,
+    engineName: testEngineItem.result.payload.engine,
+  });
+
+  gURLBar.setSearchMode({});
+
+  await UrlbarTestUtils.promisePopupClose(window, () =>
+    EventUtils.synthesizeKey("KEY_Escape")
+  );
+  await SpecialPowers.popPrefEnv();
+});
+
 // Pressing enter on an @ alias offer (an @ alias with an empty search string)
-// in the popup should fill it in the urlbar input.
-add_task(async function enterAndFillAlias() {
+// in the view should fill it in the urlbar input.
+// This subtest can be removed when update2 is on by default.
+add_task(async function enterAndFillAlias_legacy() {
   // Do a search for "@" to show all the @ aliases.
   gURLBar.search("@");
   await UrlbarTestUtils.promiseSearchComplete(window);
@@ -315,9 +357,50 @@ add_task(async function enterAndFillAlias() {
   );
 });
 
+// Pressing enter on an @ alias offer (an @ alias with an empty search string)
+// in the view should enter search mode.
+add_task(async function enterAndFillAlias() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.urlbar.update2", true]],
+  });
+
+  // Do a search for "@" to show all the @ aliases.
+  gURLBar.search("@");
+  await UrlbarTestUtils.promiseSearchComplete(window);
+
+  // Find our test engine in the results.  It's probably last, but for
+  // robustness don't assume it is.
+  let details;
+  let index = 0;
+  for (; ; index++) {
+    details = await UrlbarTestUtils.getDetailsOfResultAt(window, index);
+    if (details.searchParams && details.searchParams.keyword == ALIAS) {
+      index++;
+      break;
+    }
+  }
+
+  // Key down to it and press enter.
+  EventUtils.synthesizeKey("KEY_ArrowDown", { repeat: index });
+  EventUtils.synthesizeKey("KEY_Enter");
+
+  UrlbarTestUtils.assertSearchMode(window, {
+    source: UrlbarUtils.RESULT_SOURCE.SEARCH,
+    engineName: details.searchParams.engine,
+  });
+
+  gURLBar.setSearchMode({});
+
+  await UrlbarTestUtils.promisePopupClose(window, () =>
+    EventUtils.synthesizeKey("KEY_Escape")
+  );
+  await SpecialPowers.popPrefEnv();
+});
+
 // Pressing enter on an @ alias autofill should fill it in the urlbar input
 // with a trailing space and move the caret at the end.
-add_task(async function enterAutofillsAlias() {
+// This subtest can be removed when update2 is on by default.
+add_task(async function enterAutofillsAlias_legacy() {
   let expectedString = `${ALIAS} `;
   for (let value of [ALIAS.substring(0, ALIAS.length - 1), ALIAS]) {
     await UrlbarTestUtils.promiseAutocompleteResultPopup({
@@ -344,6 +427,40 @@ add_task(async function enterAutofillsAlias() {
   );
 });
 
+// Pressing enter on an @ alias autofill should fill it in the urlbar input
+// with a trailing space and move the caret at the end.
+add_task(async function enterAutofillsAlias() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.urlbar.update2", true]],
+  });
+  for (let value of [ALIAS.substring(0, ALIAS.length - 1), ALIAS]) {
+    await UrlbarTestUtils.promiseAutocompleteResultPopup({
+      window,
+      value,
+      selectionStart: value.length,
+      selectionEnd: value.length,
+    });
+    let testEngineItem = await UrlbarTestUtils.waitForAutocompleteResultAt(
+      window,
+      0
+    );
+
+    // Press Enter.
+    EventUtils.synthesizeKey("KEY_Enter");
+
+    UrlbarTestUtils.assertSearchMode(window, {
+      source: UrlbarUtils.RESULT_SOURCE.SEARCH,
+      engineName: testEngineItem.result.payload.engine,
+    });
+
+    gURLBar.setSearchMode({});
+  }
+  await UrlbarTestUtils.promisePopupClose(window, () =>
+    EventUtils.synthesizeKey("KEY_Escape")
+  );
+  await SpecialPowers.popPrefEnv();
+});
+
 async function doSimpleTest(revertBetweenSteps) {
   // When autofill is enabled, searching for "@tes" will autofill to "@test",
   // which gets in the way of this test task, so temporarily disable it.

browser/components/urlbar/tests/browser/browser_top_sites.js

@@ -132,7 +132,8 @@ add_task(async function topSitesShown() {
   }
 });
 
-add_task(async function selectSearchTopSite() {
+// This subtest can be removed when update2 is on by default.
+add_task(async function selectSearchTopSite_legacy() {
   await updateTopSites(
     sites => sites && sites[0] && sites[0].searchTopSite,
     true
@@ -174,6 +175,51 @@ add_task(async function selectSearchTopSite() {
   });
 });
 
+add_task(async function selectSearchTopSite() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.urlbar.update2", true]],
+  });
+  await updateTopSites(
+    sites => sites && sites[0] && sites[0].searchTopSite,
+    true
+  );
+  await UrlbarTestUtils.promisePopupOpen(window, () => {
+    if (gURLBar.getAttribute("pageproxystate") == "invalid") {
+      gURLBar.handleRevert();
+    }
+    EventUtils.synthesizeMouseAtCenter(gURLBar.inputField, {});
+  });
+  await UrlbarTestUtils.promiseSearchComplete(window);
+
+  let amazonSearch = await UrlbarTestUtils.waitForAutocompleteResultAt(
+    window,
+    0
+  );
+
+  Assert.equal(
+    amazonSearch.result.type,
+    UrlbarUtils.RESULT_TYPE.SEARCH,
+    "First result should have SEARCH type."
+  );
+
+  Assert.equal(
+    amazonSearch.result.payload.keyword,
+    "@amazon",
+    "First result should have the Amazon keyword."
+  );
+
+  EventUtils.synthesizeMouseAtCenter(amazonSearch, {});
+  UrlbarTestUtils.assertSearchMode(window, {
+    source: UrlbarUtils.RESULT_SOURCE.SEARCH,
+    engineName: amazonSearch.result.payload.engine,
+  });
+  gURLBar.setSearchMode({});
+
+  await UrlbarTestUtils.promisePopupClose(window, () => {
+    gURLBar.blur();
+  });
+});
+
 add_task(async function topSitesBookmarksAndTabs() {
   await addTestVisits();
   let sites = AboutNewTab.getTopSites();