[ToolbarMVC] Move searchQuery logic to LocationBarMediator

Bug: 1147581

Change-Id: Iaa39fc9bb763955476cfade0cf54bcf069147aee
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2533332
Commit-Queue: Patrick Noland <pnoland@chromium.org>
Reviewed-by: Brandon Wylie <wylieb@chromium.org>
Reviewed-by: Filip Gorski <fgorski@chromium.org>
Cr-Commit-Position: refs/heads/master@{#835759}

chrome/android/chrome_test_java_sources.gni

@@ -323,6 +323,7 @@ chrome_test_java_sources = [
   "javatests/src/org/chromium/chrome/browser/omaha/StringSanitizerTest.java",
   "javatests/src/org/chromium/chrome/browser/omaha/UpdateMenuItemHelperTest.java",
   "javatests/src/org/chromium/chrome/browser/omnibox/LocationBarLayoutTest.java",
+  "javatests/src/org/chromium/chrome/browser/omnibox/LocationBarTest.java",
   "javatests/src/org/chromium/chrome/browser/omnibox/OmniboxTest.java",
   "javatests/src/org/chromium/chrome/browser/omnibox/OmniboxUrlEmphasizerTest.java",
   "javatests/src/org/chromium/chrome/browser/omnibox/UrlBarIntegrationTest.java",

chrome/android/java/src/org/chromium/chrome/browser/omnibox/LocationBarCoordinator.java

@@ -33,8 +33,10 @@ import org.chromium.chrome.browser.omnibox.voice.AssistantVoiceSearchService;
 import org.chromium.chrome.browser.omnibox.voice.VoiceRecognitionHandler;
 import org.chromium.chrome.browser.privacy.settings.PrivacyPreferencesManagerImpl;
 import org.chromium.chrome.browser.profiles.Profile;
+import org.chromium.chrome.browser.search_engines.TemplateUrlServiceFactory;
 import org.chromium.chrome.browser.share.ShareDelegate;
 import org.chromium.chrome.browser.tabmodel.IncognitoStateProvider;
+import org.chromium.components.search_engines.TemplateUrlService;
 import org.chromium.ui.KeyboardVisibilityDelegate;
 import org.chromium.ui.base.DeviceFormFactor;
 import org.chromium.ui.base.WindowAndroid;
@@ -68,6 +70,8 @@ public final class LocationBarCoordinator implements LocationBar, NativeInitObse
     private View mAutocompleteAnchorView;
     private LocationBarMediator mLocationBarMediator;
     private View mUrlBar;
+    private final OneshotSupplierImpl<TemplateUrlService> mTemplateUrlServiceSupplier =
+            new OneshotSupplierImpl<>();
     private CallbackController mCallbackController = new CallbackController();
 
     private boolean mNativeInitialized;
@@ -128,7 +132,7 @@ public final class LocationBarCoordinator implements LocationBar, NativeInitObse
         mLocationBarMediator = new LocationBarMediator(mLocationBarLayout, locationBarDataProvider,
                 assistantVoiceSearchSupplier, profileObservableSupplier,
                 PrivacyPreferencesManagerImpl.getInstance(), overrideUrlLoadingDelegate,
-                LocaleManager.getInstance());
+                LocaleManager.getInstance(), mTemplateUrlServiceSupplier);
         mUrlCoordinator =
                 new UrlBarCoordinator((UrlBar) mUrlBar, windowDelegate, actionModeCallback,
                         mCallbackController.makeCancelable(mLocationBarMediator::onUrlFocusChange),
@@ -186,6 +190,7 @@ public final class LocationBarCoordinator implements LocationBar, NativeInitObse
 
     @Override
     public void onFinishNativeInitialization() {
+        mTemplateUrlServiceSupplier.set(TemplateUrlServiceFactory.get());
         mLocationBarMediator.onFinishNativeInitialization();
         mAutocompleteCoordinator.onNativeInitialized();
         mStatusCoordinator.onNativeInitialized();
@@ -411,4 +416,8 @@ public final class LocationBarCoordinator implements LocationBar, NativeInitObse
             VoiceRecognitionHandler voiceRecognitionHandler) {
         mLocationBarMediator.setVoiceRecognitionHandlerForTesting(voiceRecognitionHandler);
     }
+
+    /* package */ LocationBarMediator getMediatorForTesting() {
+        return mLocationBarMediator;
+    }
 }

chrome/android/java/src/org/chromium/chrome/browser/omnibox/LocationBarLayout.java

@@ -308,39 +308,6 @@ public class LocationBarLayout extends FrameLayout implements OnClickListener {
         // When we restore tabs, we focus the selected tab so the URL of the page shows.
     }
 
-    /**
-     * Sets the query string in the omnibox (ensuring the URL bar has focus and triggering
-     * autocomplete for the specified query) as if the user typed it.
-     *
-     * @param query The query to be set in the omnibox.
-     */
-    /* package */ void setSearchQuery(final String query) {
-        if (TextUtils.isEmpty(query)) return;
-
-        if (!mNativeInitialized) {
-            mDeferredNativeRunnables.add(new Runnable() {
-                @Override
-                public void run() {
-                    setSearchQuery(query);
-                }
-            });
-            return;
-        }
-
-        // Ensure the UrlBar has focus before entering text. If the UrlBar is not focused,
-        // autocomplete text will be updated but the visible text will not.
-        setUrlBarFocus(true, null, OmniboxFocusReason.SEARCH_QUERY);
-        setUrlBarText(UrlBarData.forNonUrlText(query), UrlBar.ScrollType.NO_SCROLL,
-                SelectionState.SELECT_ALL);
-        mAutocompleteCoordinator.startAutocompleteForQuery(query);
-        post(new Runnable() {
-            @Override
-            public void run() {
-                getWindowAndroid().getKeyboardDelegate().showKeyboard(mUrlBar);
-            }
-        });
-    }
-
     @Override
     public void onClick(View v) {
         if (v == mDeleteButton) {

chrome/android/java/src/org/chromium/chrome/browser/omnibox/LocationBarMediator.java

@@ -18,6 +18,7 @@ import org.chromium.base.CallbackController;
 import org.chromium.base.CommandLine;
 import org.chromium.base.metrics.RecordUserAction;
 import org.chromium.base.supplier.ObservableSupplier;
+import org.chromium.base.supplier.OneshotSupplier;
 import org.chromium.base.supplier.OneshotSupplierImpl;
 import org.chromium.chrome.browser.AppHooks;
 import org.chromium.chrome.browser.flags.ChromeSwitches;
@@ -36,15 +37,16 @@ import org.chromium.chrome.browser.omnibox.voice.VoiceRecognitionHandler;
 import org.chromium.chrome.browser.preferences.SharedPreferencesManager;
 import org.chromium.chrome.browser.privacy.settings.PrivacyPreferencesManagerImpl;
 import org.chromium.chrome.browser.profiles.Profile;
-import org.chromium.chrome.browser.search_engines.TemplateUrlServiceFactory;
 import org.chromium.chrome.browser.tab.Tab;
 import org.chromium.chrome.browser.util.KeyNavigationUtil;
 import org.chromium.components.embedder_support.util.UrlUtilities;
+import org.chromium.components.search_engines.TemplateUrlService;
 import org.chromium.content_public.browser.LoadUrlParams;
 import org.chromium.content_public.common.ResourceRequestBody;
 import org.chromium.ui.base.PageTransition;
 import org.chromium.ui.base.WindowAndroid;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 
@@ -70,6 +72,8 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
     private CallbackController mCallbackController = new CallbackController();
     private final OverrideUrlLoadingDelegate mOverrideUrlLoadingDelegate;
     private final LocaleManager mLocaleManager;
+    private final List<Runnable> mDeferredNativeRunnables = new ArrayList<>();
+    private final OneshotSupplier<TemplateUrlService> mTemplateUrlServiceSupplier;
 
     private boolean mNativeInitialized;
 
@@ -79,17 +83,19 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
             @NonNull ObservableSupplier<Profile> profileSupplier,
             @NonNull PrivacyPreferencesManagerImpl privacyPreferencesManager,
             @NonNull OverrideUrlLoadingDelegate overrideUrlLoadingDelegate,
-            @NonNull LocaleManager localeManager) {
+            @NonNull LocaleManager localeManager,
+            @NonNull OneshotSupplier<TemplateUrlService> templateUrlServiceSupplier) {
         mLocationBarLayout = locationBarLayout;
         mLocationBarDataProvider = locationBarDataProvider;
         mLocationBarDataProvider.addObserver(this);
         mAssistantVoiceSearchSupplier = assistantVoiceSearchSupplier;
+        mOverrideUrlLoadingDelegate = overrideUrlLoadingDelegate;
+        mLocaleManager = localeManager;
         mVoiceRecognitionHandler = new VoiceRecognitionHandler(this, mAssistantVoiceSearchSupplier);
         mProfileSupplier = profileSupplier;
         mProfileSupplier.addObserver(mCallbackController.makeCancelable(this::setProfile));
         mPrivacyPreferencesManager = privacyPreferencesManager;
-        mOverrideUrlLoadingDelegate = overrideUrlLoadingDelegate;
-        mLocaleManager = localeManager;
+        mTemplateUrlServiceSupplier = templateUrlServiceSupplier;
     }
 
     /**
@@ -116,6 +122,9 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
         mAutocompleteCoordinator = null;
         mUrlCoordinator = null;
         mPrivacyPreferencesManager = null;
+        mVoiceRecognitionHandler = null;
+        mLocationBarDataProvider.removeObserver(this);
+        mDeferredNativeRunnables.clear();
     }
 
     /*package */ void onUrlFocusChange(boolean hasFocus) {
@@ -127,11 +136,16 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
         mOmniboxPrerender = new OmniboxPrerender();
         Context context = mLocationBarLayout.getContext();
         mAssistantVoiceSearchService = new AssistantVoiceSearchService(context,
-                AppHooks.get().getExternalAuthUtils(), TemplateUrlServiceFactory.get(),
+                AppHooks.get().getExternalAuthUtils(), mTemplateUrlServiceSupplier.get(),
                 GSAState.getInstance(context), this, SharedPreferencesManager.getInstance());
         mAssistantVoiceSearchSupplier.set(mAssistantVoiceSearchService);
         mLocationBarLayout.onFinishNativeInitialization();
         setProfile(mProfileSupplier.get());
+
+        for (Runnable deferredRunnable : mDeferredNativeRunnables) {
+            mLocationBarLayout.post(deferredRunnable);
+        }
+        mDeferredNativeRunnables.clear();
     }
 
     /*package */ void setUrlFocusChangeFraction(float fraction) {
@@ -348,12 +362,13 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
     public void performSearchQuery(String query, List<String> searchParams) {
         if (TextUtils.isEmpty(query)) return;
 
-        String queryUrl = TemplateUrlServiceFactory.get().getUrlForSearchQuery(query, searchParams);
+        String queryUrl =
+                mTemplateUrlServiceSupplier.get().getUrlForSearchQuery(query, searchParams);
 
         if (!TextUtils.isEmpty(queryUrl)) {
             loadUrl(queryUrl, PageTransition.GENERATED, 0);
         } else {
-            mLocationBarLayout.setSearchQuery(query);
+            setSearchQuery(query);
         }
     }
 
@@ -408,7 +423,20 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
 
     @Override
     public void setSearchQuery(String query) {
-        mLocationBarLayout.setSearchQuery(query);
+        if (TextUtils.isEmpty(query)) return;
+
+        if (!mNativeInitialized) {
+            mDeferredNativeRunnables.add(() -> setSearchQuery(query));
+            return;
+        }
+
+        // Ensure the UrlBar has focus before entering text. If the UrlBar is not focused,
+        // autocomplete text will be updated but the visible text will not.
+        mLocationBarLayout.setUrlBarFocus(true, null, OmniboxFocusReason.SEARCH_QUERY);
+        mLocationBarLayout.setUrlBarText(UrlBarData.forNonUrlText(query),
+                UrlBar.ScrollType.NO_SCROLL, SelectionState.SELECT_ALL);
+        mAutocompleteCoordinator.startAutocompleteForQuery(query);
+        mUrlCoordinator.setKeyboardVisibility(true, false);
     }
 
     @Override
@@ -438,7 +466,6 @@ class LocationBarMediator implements LocationBarDataProvider.Observer, FakeboxDe
 
     @Override
     public boolean allowKeyboardLearning() {
-        assert mLocationBarDataProvider != null;
         return !mLocationBarDataProvider.isIncognito();
     }
 

chrome/android/javatests/src/org/chromium/chrome/browser/omnibox/LocationBarLayoutTest.java

@@ -263,18 +263,6 @@ public class LocationBarLayoutTest {
         });
     }
 
-    @Test
-    @SmallTest
-    public void testSetSearchQueryFocuses() {
-        final LocationBarLayout locationBar = getLocationBar();
-        final String query = "testing query";
-
-        TestThreadUtils.runOnUiThreadBlocking(() -> {
-            locationBar.setSearchQuery(query);
-            Assert.assertEquals(query, locationBar.mUrlCoordinator.getTextWithoutAutocomplete());
-        });
-    }
-
     /*
      *  Search engine logo tests.
      */

chrome/android/javatests/src/org/chromium/chrome/browser/omnibox/LocationBarTest.java

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.chrome.browser.omnibox;
+
+import static org.mockito.Mockito.doReturn;
+
+import android.content.Intent;
+
+import androidx.test.filters.MediumTest;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import org.chromium.base.CommandLine;
+import org.chromium.base.test.util.CommandLineFlags;
+import org.chromium.base.test.util.Criteria;
+import org.chromium.base.test.util.CriteriaHelper;
+import org.chromium.base.test.util.Matchers;
+import org.chromium.chrome.browser.ChromeTabbedActivity;
+import org.chromium.chrome.browser.app.ChromeActivity;
+import org.chromium.chrome.browser.flags.ChromeSwitches;
+import org.chromium.chrome.browser.lifecycle.InflationObserver;
+import org.chromium.chrome.browser.search_engines.TemplateUrlServiceFactory;
+import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
+import org.chromium.chrome.test.ChromeTabbedActivityTestRule;
+import org.chromium.chrome.test.util.ChromeTabUtils;
+import org.chromium.components.search_engines.TemplateUrlService;
+import org.chromium.content_public.browser.test.util.TestThreadUtils;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Instrumentation tests for the LocationBar component.
+ */
+@RunWith(ChromeJUnit4ClassRunner.class)
+@CommandLineFlags.Add({ChromeSwitches.DISABLE_FIRST_RUN_EXPERIENCE})
+public class LocationBarTest {
+    private static final String TEST_QUERY = "testing query";
+    private static final List<String> TEST_PARAMS = Arrays.asList("foo=bar");
+
+    @Rule
+    public ChromeTabbedActivityTestRule mActivityTestRule = new ChromeTabbedActivityTestRule();
+    @Rule
+    public MockitoRule mMockitoRule = MockitoJUnit.rule();
+
+    @Mock
+    TemplateUrlService mTemplateUrlService;
+
+    ChromeTabbedActivity mActivity;
+    UrlBar mUrlBar;
+    LocationBarMediator mLocationBarMediator;
+    String mSearchUrl;
+
+    @Before
+    public void setUp() throws InterruptedException {
+        TemplateUrlServiceFactory.setInstanceForTesting(mTemplateUrlService);
+    }
+
+    private void startActivityNormally() {
+        mActivityTestRule.startMainActivityOnBlankPage();
+        mActivity = mActivityTestRule.getActivity();
+        doPostActivitySetup(mActivity);
+    }
+
+    private void startActivityWithDeferredNativeInitialization() {
+        CommandLine.getInstance().appendSwitch(ChromeSwitches.DISABLE_NATIVE_INITIALIZATION);
+        Intent intent = new Intent("about:blank");
+        intent.addCategory(Intent.CATEGORY_LAUNCHER);
+        mActivityTestRule.prepareUrlIntent(intent, "about:blank");
+        mActivityTestRule.launchActivity(intent);
+        ChromeActivity chromeActivity = mActivityTestRule.getActivity();
+        if (!chromeActivity.isInitialLayoutInflationComplete()) {
+            AtomicBoolean isInflated = new AtomicBoolean();
+            chromeActivity.getLifecycleDispatcher().register(new InflationObserver() {
+                @Override
+                public void onPreInflationStartup() {}
+
+                @Override
+                public void onPostInflationStartup() {
+                    isInflated.set(true);
+                }
+            });
+            CriteriaHelper.pollUiThread(isInflated::get);
+        }
+        doPostActivitySetup(chromeActivity);
+    }
+
+    private void triggerAndWaitForDeferredNativeInitialization() {
+        CommandLine.getInstance().removeSwitch(ChromeSwitches.DISABLE_NATIVE_INITIALIZATION);
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            mActivityTestRule.getActivity().startDelayedNativeInitializationForTests();
+        });
+        mActivityTestRule.waitForActivityNativeInitializationComplete();
+    }
+
+    private void doPostActivitySetup(ChromeActivity activity) {
+        mUrlBar = activity.findViewById(org.chromium.chrome.R.id.url_bar);
+        // clang-format off
+        mLocationBarMediator = ((LocationBarCoordinator) activity.getToolbarManager()
+            .getToolbarLayoutForTesting()
+            .getLocationBar())
+            .getMediatorForTesting();
+        // clang-format on
+        mSearchUrl = mActivityTestRule.getEmbeddedTestServerRule().getServer().getURL("/search");
+    }
+
+    @Test
+    @MediumTest
+    public void testSetSearchQueryFocusesUrlBar() {
+        startActivityNormally();
+        final String query = "testing query";
+
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            mLocationBarMediator.setSearchQuery(query);
+            Assert.assertEquals(query, mUrlBar.getTextWithoutAutocomplete());
+            Assert.assertTrue(mLocationBarMediator.isUrlBarFocused());
+        });
+
+        CriteriaHelper.pollUiThread(() -> {
+            Criteria.checkThat(
+                    mActivityTestRule.getKeyboardDelegate().isKeyboardShowing(mActivity, mUrlBar),
+                    Matchers.is(true));
+        });
+    }
+
+    @Test
+    @MediumTest
+    public void testSetSearchQueryFocusesUrlBar_preNative() {
+        startActivityWithDeferredNativeInitialization();
+        final String query = "testing query";
+
+        TestThreadUtils.runOnUiThreadBlocking(() -> mLocationBarMediator.setSearchQuery(query));
+
+        triggerAndWaitForDeferredNativeInitialization();
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            Assert.assertEquals(query, mUrlBar.getTextWithoutAutocomplete());
+            Assert.assertTrue(mLocationBarMediator.isUrlBarFocused());
+        });
+
+        CriteriaHelper.pollUiThread(() -> {
+            Criteria.checkThat(
+                    mActivityTestRule.getKeyboardDelegate().isKeyboardShowing(mActivity, mUrlBar),
+                    Matchers.is(true));
+        });
+    }
+
+    @Test
+    @MediumTest
+    public void testPerformSearchQuery() {
+        startActivityNormally();
+        final List<String> params = Arrays.asList("foo=bar");
+        doReturn(mSearchUrl)
+                .when(mTemplateUrlService)
+                .getUrlForSearchQuery(TEST_QUERY, TEST_PARAMS);
+
+        TestThreadUtils.runOnUiThreadBlocking(
+                () -> mLocationBarMediator.performSearchQuery(TEST_QUERY, TEST_PARAMS));
+
+        ChromeTabUtils.waitForTabPageLoaded(
+                mActivityTestRule.getActivity().getActivityTab(), mSearchUrl);
+    }
+
+    @Test
+    @MediumTest
+    public void testPerformSearchQuery_emptyUrl() {
+        startActivityNormally();
+        doReturn("").when(mTemplateUrlService).getUrlForSearchQuery(TEST_QUERY, TEST_PARAMS);
+
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            mLocationBarMediator.performSearchQuery(TEST_QUERY, TEST_PARAMS);
+            Assert.assertEquals(TEST_QUERY, mUrlBar.getTextWithoutAutocomplete());
+        });
+    }
+}

chrome/android/junit/src/org/chromium/chrome/browser/omnibox/LocationBarMediatorTest.java

@@ -10,8 +10,12 @@ import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.notNull;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -20,15 +24,16 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.content.Context;
+import android.text.TextUtils;
 import android.view.View;
 
-import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
+import org.mockito.ArgumentMatcher;
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
@@ -37,6 +42,7 @@ import org.robolectric.annotation.Config;
 
 import org.chromium.base.BuildConfig;
 import org.chromium.base.supplier.ObservableSupplierImpl;
+import org.chromium.base.supplier.OneshotSupplier;
 import org.chromium.base.supplier.OneshotSupplierImpl;
 import org.chromium.base.test.BaseRobolectricTestRunner;
 import org.chromium.base.test.util.JniMocker;
@@ -49,7 +55,6 @@ import org.chromium.chrome.browser.omnibox.voice.AssistantVoiceSearchService;
 import org.chromium.chrome.browser.privacy.settings.PrivacyPreferencesManagerImpl;
 import org.chromium.chrome.browser.profiles.Profile;
 import org.chromium.chrome.browser.profiles.ProfileJni;
-import org.chromium.chrome.browser.search_engines.TemplateUrlServiceFactory;
 import org.chromium.chrome.browser.tab.Tab;
 import org.chromium.chrome.test.util.browser.Features;
 import org.chromium.components.embedder_support.util.UrlConstants;
@@ -57,12 +62,16 @@ import org.chromium.components.search_engines.TemplateUrlService;
 import org.chromium.content_public.browser.LoadUrlParams;
 import org.chromium.ui.base.PageTransition;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
 /** Unit tests for LocationBarMediator. */
 @RunWith(BaseRobolectricTestRunner.class)
 @Config(manifest = Config.NONE)
 @Features.EnableFeatures(ChromeFeatureList.OMNIBOX_ASSISTANT_VOICE_SEARCH)
 public class LocationBarMediatorTest {
-    private static final String TEST_URL = "http://testurl.com";
+    private static final String TEST_URL = "http://www.example.org";
 
     @Rule
     public MockitoRule mMockitoRule = MockitoJUnit.rule();
@@ -82,6 +91,10 @@ public class LocationBarMediatorTest {
     @Mock
     private OneshotSupplierImpl<AssistantVoiceSearchService> mAssistantVoiceSearchSupplier;
     @Mock
+    private OverrideUrlLoadingDelegate mOverrideUrlLoadingDelegate;
+    @Mock
+    private LocaleManager mLocaleManager;
+    @Mock
     private Profile.Natives mProfileNativesJniMock;
     @Mock
     private Tab mTab;
@@ -98,12 +111,12 @@ public class LocationBarMediatorTest {
     @Mock
     private SearchEngineLogoUtils.Delegate mSearchEngineDelegate;
     @Mock
-    private OverrideUrlLoadingDelegate mOverrideUrlLoadingDelegate;
-    @Mock
-    private LocaleManager mLocaleManager;
-    @Mock
     private View mView;
+    @Mock
+    private OneshotSupplier<TemplateUrlService> mTemplateUrlServiceSupplier;
 
+    @Captor
+    private ArgumentCaptor<Runnable> mRunnableCaptor;
     @Captor
     private ArgumentCaptor<LoadUrlParams> mLoadUrlParamsCaptor;
 
@@ -113,12 +126,13 @@ public class LocationBarMediatorTest {
     @Before
     public void setUp() {
         doReturn(mContext).when(mLocationBarLayout).getContext();
-        TemplateUrlServiceFactory.setInstanceForTesting(mTemplateUrlService);
+        doReturn(mTemplateUrlService).when(mTemplateUrlServiceSupplier).get();
         mJniMocker.mock(ProfileJni.TEST_HOOKS, mProfileNativesJniMock);
         mJniMocker.mock(OmniboxPrerenderJni.TEST_HOOKS, mPrerenderJni);
+        SearchEngineLogoUtils.setDelegateForTesting(mSearchEngineDelegate);
         mMediator = new LocationBarMediator(mLocationBarLayout, mLocationBarDataProvider,
                 mAssistantVoiceSearchSupplier, mProfileSupplier, mPrivacyPreferencesManager,
-                mOverrideUrlLoadingDelegate, mLocaleManager);
+                mOverrideUrlLoadingDelegate, mLocaleManager, mTemplateUrlServiceSupplier);
         mMediator.setCoordinators(mUrlCoordinator, mAutocompleteCoordinator, mStatusCoordinator);
         SearchEngineLogoUtils.setDelegateForTesting(mSearchEngineDelegate);
     }
@@ -214,8 +228,8 @@ public class LocationBarMediatorTest {
         mMediator.loadUrl(TEST_URL, PageTransition.TYPED, 0);
 
         verify(mTab).loadUrl(mLoadUrlParamsCaptor.capture());
-        Assert.assertEquals(TEST_URL, mLoadUrlParamsCaptor.getValue().getUrl());
-        Assert.assertEquals(PageTransition.TYPED | PageTransition.FROM_ADDRESS_BAR,
+        assertEquals(TEST_URL, mLoadUrlParamsCaptor.getValue().getUrl());
+        assertEquals(PageTransition.TYPED | PageTransition.FROM_ADDRESS_BAR,
                 mLoadUrlParamsCaptor.getValue().getTransitionType());
     }
 
@@ -267,4 +281,103 @@ public class LocationBarMediatorTest {
         verify(mLocationBarDataProvider, times(2)).getTab();
         verify(mTab, times(1)).getView();
     }
+
+    @Test
+    public void testSetSearchQuery() {
+        String query = "example search";
+        mMediator.onFinishNativeInitialization();
+        mMediator.setSearchQuery(query);
+
+        verify(mLocationBarLayout).setUrlBarFocus(true, null, OmniboxFocusReason.SEARCH_QUERY);
+        verify(mLocationBarLayout)
+                .setUrlBarText(argThat(matchesUrlBarDataForQuery(query)),
+                        eq(UrlBar.ScrollType.NO_SCROLL), eq(SelectionState.SELECT_ALL));
+        verify(mAutocompleteCoordinator).startAutocompleteForQuery(query);
+        verify(mUrlCoordinator).setKeyboardVisibility(true, false);
+    }
+
+    @Test
+    public void testSetSearchQuery_empty() {
+        mMediator.setSearchQuery("");
+        verify(mLocationBarLayout, never()).setUrlBarFocus(anyBoolean(), anyString(), anyInt());
+        verify(mLocationBarLayout, never()).post(any());
+
+        mMediator.onFinishNativeInitialization();
+        verify(mLocationBarLayout, never()).setUrlBarFocus(anyBoolean(), anyString(), anyInt());
+
+        mMediator.setSearchQuery("");
+        verify(mLocationBarLayout, never()).setUrlBarFocus(anyBoolean(), anyString(), anyInt());
+        verify(mLocationBarLayout, never()).post(any());
+    }
+
+    @Test
+    public void testSetSearchQuery_preNative() {
+        String query = "example search";
+        mMediator.setSearchQuery(query);
+        mMediator.onFinishNativeInitialization();
+
+        verify(mLocationBarLayout).post(mRunnableCaptor.capture());
+        mRunnableCaptor.getValue().run();
+
+        verify(mLocationBarLayout).setUrlBarFocus(true, null, OmniboxFocusReason.SEARCH_QUERY);
+        verify(mLocationBarLayout)
+                .setUrlBarText(argThat(matchesUrlBarDataForQuery(query)),
+                        eq(UrlBar.ScrollType.NO_SCROLL), eq(SelectionState.SELECT_ALL));
+        verify(mAutocompleteCoordinator).startAutocompleteForQuery(query);
+        verify(mUrlCoordinator).setKeyboardVisibility(true, false);
+    }
+
+    @Test
+    public void testPerformSearchQuery() {
+        mMediator.onFinishNativeInitialization();
+        String query = "example search";
+        List<String> params = Arrays.asList("param 1", "param 2");
+        doReturn("http://www.search.com")
+                .when(mTemplateUrlService)
+                .getUrlForSearchQuery("example search", params);
+        doReturn(mTab).when(mLocationBarDataProvider).getTab();
+        mMediator.performSearchQuery(query, params);
+
+        verify(mTab).loadUrl(mLoadUrlParamsCaptor.capture());
+        assertEquals("http://www.search.com", mLoadUrlParamsCaptor.getValue().getUrl());
+        assertEquals(PageTransition.GENERATED | PageTransition.FROM_ADDRESS_BAR,
+                mLoadUrlParamsCaptor.getValue().getTransitionType());
+    }
+
+    @Test
+    public void testPerformSearchQuery_empty() {
+        mMediator.performSearchQuery("", Collections.emptyList());
+        verify(mLocationBarLayout, never()).setUrlBarFocus(anyBoolean(), anyString(), anyInt());
+        verify(mLocationBarLayout, never()).post(any());
+
+        mMediator.onFinishNativeInitialization();
+        verify(mLocationBarLayout, never()).setUrlBarFocus(anyBoolean(), anyString(), anyInt());
+
+        mMediator.setSearchQuery("");
+        verify(mLocationBarLayout, never()).setUrlBarFocus(anyBoolean(), anyString(), anyInt());
+        verify(mLocationBarLayout, never()).post(any());
+    }
+
+    @Test
+    public void testPerformSearchQuery_emptyUrl() {
+        String query = "example search";
+        List<String> params = Arrays.asList("param 1", "param 2");
+        mMediator.onFinishNativeInitialization();
+        doReturn("").when(mTemplateUrlService).getUrlForSearchQuery("example search", params);
+        mMediator.performSearchQuery(query, params);
+
+        verify(mLocationBarLayout).setUrlBarFocus(true, null, OmniboxFocusReason.SEARCH_QUERY);
+        verify(mLocationBarLayout)
+                .setUrlBarText(argThat(matchesUrlBarDataForQuery(query)),
+                        eq(UrlBar.ScrollType.NO_SCROLL), eq(SelectionState.SELECT_ALL));
+        verify(mAutocompleteCoordinator).startAutocompleteForQuery(query);
+        verify(mUrlCoordinator).setKeyboardVisibility(true, false);
+    }
+
+    private ArgumentMatcher<UrlBarData> matchesUrlBarDataForQuery(String query) {
+        return actual -> {
+            UrlBarData expected = UrlBarData.forNonUrlText(query);
+            return TextUtils.equals(actual.displayText, expected.displayText);
+        };
+    }
 }