// Copyright 2019 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/navigation_predictor/search_engine_preconnector.h"

#include "base/run_loop.h"
#include "base/task/single_thread_task_runner.h"
#include "base/test/gmock_callback_support.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/simple_test_tick_clock.h"
#include "chrome/browser/battery/battery_saver.h"
#include "chrome/browser/navigation_predictor/navigation_predictor_keyed_service.h"
#include "chrome/browser/navigation_predictor/navigation_predictor_keyed_service_factory.h"
#include "chrome/browser/navigation_predictor/search_engine_preconnector_keyed_service_factory.h"
#include "chrome/browser/predictors/loading_predictor.h"
#include "chrome/browser/predictors/loading_predictor_factory.h"
#include "chrome/browser/search_engines/template_url_service_factory.h"
#include "chrome/browser/subresource_filter/subresource_filter_browser_test_harness.h"
#include "chrome/browser/ui/browser_window/public/browser_window_interface.h"
#include "chrome/browser/ui/tabs/tab_strip_model.h"
#include "chrome/test/base/in_process_browser_test.h"
#include "chrome/test/base/search_test_utils.h"
#include "chrome/test/base/ui_test_utils.h"
#include "components/search_engines/template_url_service.h"
#include "components/ukm/test_ukm_recorder.h"
#include "content/public/browser/preconnect_manager.h"
#include "content/public/browser/storage_partition.h"
#include "content/public/test/browser_test.h"
#include "content/public/test/browser_test_utils.h"
#include "net/base/features.h"
#include "net/dns/mock_host_resolver.h"
#include "net/net_buildflags.h"
#include "net/test/embedded_test_server/embedded_test_server.h"
#include "services/metrics/public/cpp/ukm_builders.h"
#include "url/url_constants.h"

#if BUILDFLAG(ENABLE_DEVICE_BOUND_SESSIONS)
#include "chrome/browser/navigation_predictor/navigation_predictor_features.h"
#include "chrome/common/pref_names.h"
#include "components/prefs/pref_service.h"
#include "components/search_engines/template_url_service.h"
#include "services/network/test/mock_device_bound_session_manager.h"
#endif  // BUILDFLAG(ENABLE_DEVICE_BOUND_SESSIONS)

namespace {

class SearchEnginePreconnectorBrowserTest
    : public subresource_filter::SubresourceFilterBrowserTest,
      public content::PreconnectManager::Observer {
 public:
  static constexpr char kFakeSearch[] = "https://www.fakesearch.com/";
  static constexpr char kGoogleSearch[] = "https://www.google.com/";

  SearchEnginePreconnectorBrowserTest() = default;
  ~SearchEnginePreconnectorBrowserTest() override = default;

  void SetUp() override {
    https_server_ = std::make_unique<net::EmbeddedTestServer>(
        net::EmbeddedTestServer::TYPE_HTTPS);
    https_server_->ServeFilesFromSourceDirectory(
        "chrome/test/data/navigation_predictor");
    ASSERT_TRUE(https_server_->Start());

    preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()] = 0;
    preresolve_counts_[GURL(kGoogleSearch)] = 0;
    preresolve_counts_[GURL(kFakeSearch)] = 0;

    subresource_filter::SubresourceFilterBrowserTest::SetUp();
  }

  SearchEnginePreconnector* GetSearchEnginePreconnector() {
    if (PreconnectFromKeyedServiceEnabled()) {
      return SearchEnginePreconnectorKeyedServiceFactory::GetForProfile(
          browser()->GetProfile());
    }

    NavigationPredictorKeyedService* navigation_predictor_keyed_service =
        NavigationPredictorKeyedServiceFactory::GetForProfile(
            browser()->GetProfile());
    EXPECT_TRUE(navigation_predictor_keyed_service);

    return navigation_predictor_keyed_service->search_engine_preconnector();
  }

  void SetUpOnMainThread() override {
    subresource_filter::SubresourceFilterBrowserTest::SetUpOnMainThread();
    host_resolver()->ClearRules();

    // Get notified for Loading predictor's preconnect observer.
    auto* loading_predictor =
        predictors::LoadingPredictorFactory::GetForProfile(
            browser()->GetProfile());
    ASSERT_TRUE(loading_predictor);
    loading_predictor->preconnect_manager()->SetObserverForTesting(this);

    // Also get notified for the SearchEnginePreconnect's preconnect observer
    SearchEnginePreconnector* preconnector = GetSearchEnginePreconnector();
    ASSERT_TRUE(preconnector);
    preconnector->GetPreconnectManager().SetObserverForTesting(this);
  }

  const GURL GetTestURL(const char* file) const {
    return https_server_->GetURL(file);
  }

  void OnPreresolveFinished(
      const GURL& url,
      const net::NetworkAnonymizationKey& network_anonymization_key,
      mojo::PendingRemote<network::mojom::ConnectionChangeObserverClient>&
          observer,
      bool success) override {
    // Take the observer so that we can manually send mojo message.
    if (observer.is_valid() && !remote_.is_bound()) {
      remote_.Bind(std::move(observer));
    }

    const GURL origin = url.DeprecatedGetOriginAsURL();
    if (!preresolve_counts_.contains(origin)) {
      return;
    }

    // Only assert the positive case: the test URL must preconnect. Don't
    // assert the search host fails to resolve, because a leaked
    // "*"->127.0.0.1 rule in the shared network service can't be cleared
    // from the test, so the search host may resolve to 127.0.0.1 here too.
    if (origin == GetTestURL("/").DeprecatedGetOriginAsURL()) {
      EXPECT_TRUE(success);
    }

    ++preresolve_counts_[origin];
    if (run_loops_[origin])
      run_loops_[origin]->Quit();
  }

  void WaitForPreresolveCountForURL(const GURL& url, int expected_count) {
    const GURL origin = url.DeprecatedGetOriginAsURL();
    EXPECT_TRUE(preresolve_counts_.contains(origin));
    while (preresolve_counts_[origin] < expected_count) {
      run_loops_[origin] = std::make_unique<base::RunLoop>();
      run_loops_[origin]->Run();
      run_loops_[origin].reset();
    }
  }

  void WaitForDelay(base::TimeDelta delay) {
    base::RunLoop run_loop;
    base::SingleThreadTaskRunner::GetCurrentDefault()->PostDelayedTask(
        FROM_HERE, run_loop.QuitClosure(), delay);
    run_loop.Run();
  }

  virtual bool PreconnectFromKeyedServiceEnabled() const;

 protected:
  std::map<GURL, int> preresolve_counts_;
  base::test::ScopedFeatureList feature_list_;
  std::map<GURL, std::unique_ptr<base::RunLoop>> run_loops_;

  mojo::Remote<network::mojom::ConnectionChangeObserverClient> remote_;

 private:
  std::unique_ptr<net::EmbeddedTestServer> https_server_;
};

bool SearchEnginePreconnectorBrowserTest::PreconnectFromKeyedServiceEnabled()
    const {
  return SearchEnginePreconnector::ShouldBeEnabledAsKeyedService();
}

// static
constexpr char SearchEnginePreconnectorBrowserTest::kFakeSearch[];
constexpr char SearchEnginePreconnectorBrowserTest::kGoogleSearch[];

class SearchEnginePreconnectorNoDelaysBrowserTest
    : public SearchEnginePreconnectorBrowserTest,
      public testing::WithParamInterface<std::tuple<bool, bool>> {
 public:
  SearchEnginePreconnectorNoDelaysBrowserTest() {
    std::vector<base::test::FeatureRefAndParams> enabled_features{
        {features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
        {net::features::kSearchEnginePreconnectInterval,
         {{"preconnect_interval", "0"}}}};

    std::vector<base::test::FeatureRef> disabled_features;

    if (PreconnectFromKeyedServiceEnabled()) {
      enabled_features.push_back(
          {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
    } else {
      disabled_features.emplace_back(features::kPreconnectFromKeyedService);
    }

    if (SearchEnginePreconnect2Enabled()) {
      enabled_features.push_back({net::features::kSearchEnginePreconnect2, {}});
    } else {
      disabled_features.emplace_back(net::features::kSearchEnginePreconnect2);
    }

    feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                disabled_features);
  }

  bool PreconnectFromKeyedServiceEnabled() const override {
    return std::get<0>(GetParam());
  }
  bool SearchEnginePreconnect2Enabled() const {
    return std::get<1>(GetParam());
  }

  ~SearchEnginePreconnectorNoDelaysBrowserTest() override = default;
};

INSTANTIATE_TEST_SUITE_P(All,
                         SearchEnginePreconnectorNoDelaysBrowserTest,
                         ::testing::Combine(::testing::Bool(),
                                            ::testing::Bool()));

// Test routinely flakes on the Mac10.11 Tests bot (https://crbug.com/40726926).
IN_PROC_BROWSER_TEST_P(SearchEnginePreconnectorNoDelaysBrowserTest,
                       DISABLED_PreconnectSearch) {
  // Put the fake search URL to be preconnected in foreground.
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);
  // Verifies that the default search is preconnected.
  constexpr char16_t kShortName[] = u"test";
  constexpr char kSearchURL[] = "/anchors_different_area.html?q={searchTerms}";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  // Check default URL is being preconnected and test URL is not.
  const GURL kDefaultUrl(kGoogleSearch);
  WaitForPreresolveCountForURL(kDefaultUrl, 2);
  EXPECT_EQ(2, preresolve_counts_[kDefaultUrl]);
  EXPECT_EQ(0, preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);

  TemplateURLData data;
  data.SetShortName(kShortName);
  data.SetKeyword(data.short_name());
  data.SetURL(GetTestURL(kSearchURL).spec());
  data.preconnect_to_search_url = true;

  TemplateURL* template_url = model->Add(std::make_unique<TemplateURL>(data));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  // Put the fake search URL to be preconnected in foreground.
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  // After switching search providers, the test URL should now start being
  // preconnected.
  WaitForPreresolveCountForURL(GetTestURL("/"), 1);
  // Preconnect should occur for DSE.
  EXPECT_EQ(1, preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);

  WaitForPreresolveCountForURL(GetTestURL("/"), 2);
  // Preconnect should occur again for DSE.
  EXPECT_EQ(2, preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);
}

// TODO(crbug.com/413293448): Flaky test
IN_PROC_BROWSER_TEST_P(SearchEnginePreconnectorNoDelaysBrowserTest,
                       DISABLED_PreconnectOnlyInForeground) {
  constexpr char16_t kShortName[] = u"test";
  constexpr char kSearchURL[] = "/anchors_different_area.html?q={searchTerms}";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data;
  data.SetShortName(kShortName);
  data.SetKeyword(data.short_name());
  data.SetURL(GetTestURL(kSearchURL).spec());
  data.preconnect_to_search_url = true;

  // Set the DSE to the test URL.
  TemplateURL* template_url = model->Add(std::make_unique<TemplateURL>(data));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  // Ensure that we wait long enough to trigger preconnects.
  WaitForDelay(base::Milliseconds(200));

  TemplateURLData data_fake_search;
  data_fake_search.SetShortName(kShortName);
  data_fake_search.SetKeyword(data.short_name());
  data_fake_search.SetURL(kFakeSearch);
  data_fake_search.preconnect_to_search_url = true;

  template_url = model->Add(std::make_unique<TemplateURL>(data_fake_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  // Reset the count if SearchEnginePreconnect2 is enabled, since the
  // KeyedService may already spawn the preconnector, and will re-attempt
  // preconnect automatically, resulting in flaky counts
  if (SearchEnginePreconnect2Enabled()) {
    preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()] = 0;
  }

  // Put the fake search URL to be preconnected in foreground.
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);
  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Preconnect should occur for fake search.
  EXPECT_EQ(1, preresolve_counts_[search_url]);

  // No preconnects should have been issued for the test URL.
  EXPECT_EQ(0, preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);
}

class SearchEnginePreconnectorForegroundBrowserTest
    : public SearchEnginePreconnectorBrowserTest,
      public testing::WithParamInterface<std::tuple<bool, bool, bool, bool>> {
 public:
  SearchEnginePreconnectorForegroundBrowserTest() {
    {
      std::vector<base::test::FeatureRefAndParams> enabled_features;
      std::vector<base::test::FeatureRef> disabled_features;
      if (skip_in_background()) {
        enabled_features.push_back({features::kPreconnectToSearch,
                                    {{"startup_delay_ms", "1000000"},
                                     {"skip_in_background", "true"}}});
      } else {
        enabled_features.push_back({features::kPreconnectToSearch,
                                    {{"startup_delay_ms", "1000000"},
                                     {"skip_in_background", "false"}}});
      }

      if (PreconnectFromKeyedServiceEnabled()) {
        enabled_features.push_back(
            {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
      } else {
        disabled_features.emplace_back(features::kPreconnectFromKeyedService);
      }

      if (SearchEnginePreconnect2Enabled()) {
        enabled_features.push_back(
            {net::features::kSearchEnginePreconnect2, {}});
      } else {
        disabled_features.emplace_back(net::features::kSearchEnginePreconnect2);
      }
      feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                  disabled_features);
    }
  }

  bool skip_in_background() const { return std::get<0>(GetParam()); }

  bool load_page() const { return std::get<1>(GetParam()); }

  bool PreconnectFromKeyedServiceEnabled() const override {
    return std::get<2>(GetParam());
  }

  bool SearchEnginePreconnect2Enabled() const {
    return std::get<3>(GetParam());
  }

  ~SearchEnginePreconnectorForegroundBrowserTest() override = default;

  base::SimpleTestTickClock tick_clock_;
};

INSTANTIATE_TEST_SUITE_P(All,
                         SearchEnginePreconnectorForegroundBrowserTest,
                         ::testing::Combine(::testing::Bool(),
                                            ::testing::Bool(),
                                            ::testing::Bool(),
                                            ::testing::Bool()));

// Test that search engine preconnects are done only if the browser app is
// likely in foreground.
//
// TODO(crbug.com/413293448): Disabled the test for flakiness due to test setup.
IN_PROC_BROWSER_TEST_P(SearchEnginePreconnectorForegroundBrowserTest,
                       DISABLED_PreconnectOnlyInForeground) {
  static const char16_t kShortName[] = u"test";
  static const char kSearchURL[] =
      "/anchors_different_area.html?q={searchTerms}";
  static const char kSearchURLWithQuery[] =
      "/anchors_different_area.html?q=porgs";

  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data;
  data.SetShortName(kShortName);
  data.SetKeyword(data.short_name());
  data.SetURL(GetTestURL(kSearchURL).spec());
  data.preconnect_to_search_url = true;

  // Set the DSE to the test URL.
  TemplateURL* template_url = model->Add(std::make_unique<TemplateURL>(data));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  // Ensure that we wait long enough to trigger preconnects.
  WaitForDelay(base::Milliseconds(200));

  TemplateURLData data_fake_search;
  data_fake_search.SetShortName(kShortName);
  data_fake_search.SetKeyword(data.short_name());
  const GURL fake_search_url(kFakeSearch);
  data_fake_search.SetURL(kFakeSearch);
  data_fake_search.preconnect_to_search_url = true;

  template_url = model->Add(std::make_unique<TemplateURL>(data_fake_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  tick_clock_.SetNowTicks(base::TimeTicks::Now());
  tick_clock_.Advance(base::Seconds(10000));

  GetSearchEnginePreconnector()->SetTickClockForTesting(&tick_clock_);

  // Reset the count if SearchEnginePreconnect2 is enabled, since the
  // KeyedService may already spawn the preconnector, and will re-attempt
  // preconnect automatically, resulting in flaky counts
  if (SearchEnginePreconnect2Enabled()) {
    preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()] = 0;
  }
  if (load_page()) {
    ASSERT_TRUE(ui_test_utils::NavigateToURL(browser(),
                                             GetTestURL(kSearchURLWithQuery)));
  }

  // We start recording the histogram here since the KeyedService may already
  // spawn the preconnector and the histogram may already have some entries
  // prior to this `StartPreconnect`. By placing the HistogramTester here, we
  // will can ensure we only record the histogram for this test.
  base::HistogramTester histogram_tester;
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  if (!skip_in_background() || load_page()) {
    WaitForPreresolveCountForURL(fake_search_url, 1);
  }

  // If preconnects are skipped in background and no web contents is in
  // foreground, then no preconnect should happen.
  EXPECT_EQ(skip_in_background() && !load_page() ? 0 : 1,
            preresolve_counts_[fake_search_url]);
  histogram_tester.ExpectUniqueSample(
      "NavigationPredictor.SearchEnginePreconnector."
      "IsBrowserAppLikelyInForeground",
      !!load_page(), 1);
  EXPECT_EQ(load_page() ? 1 : 0,
            preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);
}

class SearchEnginePreconnectorKeepSocketBrowserTest
    : public SearchEnginePreconnectorBrowserTest {
 public:
  SearchEnginePreconnectorKeepSocketBrowserTest() {
    feature_list_.InitWithFeaturesAndParameters(
        {{features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
         {net::features::kSearchEnginePreconnectInterval,
          {{"preconnect_interval", "60"}}}},
        {});
  }

  ~SearchEnginePreconnectorKeepSocketBrowserTest() override = default;
};

// TODO(https://crbug.com/507121988): Re-enable once the test is fixed.
#if BUILDFLAG(IS_LINUX) || BUILDFLAG(IS_WIN)
#define MAYBE_SocketWarmForSearch DISABLED_SocketWarmForSearch
#else
#define MAYBE_SocketWarmForSearch SocketWarmForSearch
#endif  // BUILDFLAG(IS_LINUX) || BUILDFLAG(IS_WIN)
IN_PROC_BROWSER_TEST_F(SearchEnginePreconnectorKeepSocketBrowserTest,
                       MAYBE_SocketWarmForSearch) {
  // Verifies that a navigation to search will use a warm socket.
  constexpr char16_t kShortName[] = u"test";
  constexpr char kSearchURL[] = "/anchors_different_area.html?q={searchTerms}";
  constexpr char kSearchURLWithQuery[] = "/anchors_different_area.html?q=porgs";

  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data;
  data.SetShortName(kShortName);
  data.SetKeyword(data.short_name());
  data.SetURL(GetTestURL(kSearchURL).spec());
  data.preconnect_to_search_url = true;

  // Set the DSE to the test URL.
  TemplateURL* template_url = model->Add(std::make_unique<TemplateURL>(data));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  // Put the fake search URL to be preconnected in foreground.
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  WaitForPreresolveCountForURL(GetTestURL(kSearchURL), 1);

  ASSERT_TRUE(
      ui_test_utils::NavigateToURL(browser(), GetTestURL(kSearchURLWithQuery)));

  auto ukm_recorder = std::make_unique<ukm::TestAutoSetUkmRecorder>();

  ASSERT_TRUE(
      ui_test_utils::NavigateToURL(browser(), GURL(url::kAboutBlankURL)));

  const auto& entries =
      ukm_recorder->GetMergedEntriesByName(ukm::builders::PageLoad::kEntryName);
  EXPECT_EQ(1u, entries.size());

  for (const auto& kv : entries) {
    EXPECT_TRUE(ukm_recorder->EntryHasMetric(
        kv.second.get(),
        ukm::builders::PageLoad::kMainFrameResource_SocketReusedName));
  }
}

class SearchEnginePreconnectorDesktopAutoStartBrowserTest
    : public SearchEnginePreconnectorBrowserTest {
 public:
  SearchEnginePreconnectorDesktopAutoStartBrowserTest() {
    feature_list_.InitWithFeaturesAndParameters(
        {
            {features::kPreconnectToSearch, {{"startup_delay_ms", "0"}}},
            {net::features::kSearchEnginePreconnectInterval,
             {{"preconnect_interval", "0"}}},
        },
        {});
  }

  ~SearchEnginePreconnectorDesktopAutoStartBrowserTest() override = default;
};

IN_PROC_BROWSER_TEST_F(SearchEnginePreconnectorDesktopAutoStartBrowserTest,
                       AutoStartDesktop) {
  int preresolve_count =
      SearchEnginePreconnector::SearchEnginePreconnect2Enabled() ? 1 : 2;
  // Verifies that the default search is preconnected.
  WaitForPreresolveCountForURL(GURL(kGoogleSearch), preresolve_count);
}

class SearchEnginePreconnectorEnabledOnlyBrowserTest
    : public SearchEnginePreconnectorBrowserTest,
      public testing::WithParamInterface<std::tuple<bool, bool>> {
 public:
  SearchEnginePreconnectorEnabledOnlyBrowserTest() {
    {
      std::vector<base::test::FeatureRefAndParams> enabled_features{
          {features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
          {net::features::kSearchEnginePreconnectInterval,
           {{"preconnect_interval", "60"}}}};

      std::vector<base::test::FeatureRef> disabled_features;
      if (PreconnectFromKeyedServiceEnabled()) {
        enabled_features.push_back(
            {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
      } else {
        disabled_features.emplace_back(features::kPreconnectFromKeyedService);
      }

      if (SearchEnginePreconnect2Enabled()) {
        enabled_features.push_back(
            {net::features::kSearchEnginePreconnect2, {}});
      } else {
        disabled_features.emplace_back(net::features::kSearchEnginePreconnect2);
      }

      feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                  disabled_features);
    }
  }

  bool PreconnectFromKeyedServiceEnabled() const override {
    return std::get<0>(GetParam());
  }
  bool SearchEnginePreconnect2Enabled() const {
    return std::get<1>(GetParam());
  }

  ~SearchEnginePreconnectorEnabledOnlyBrowserTest() override = default;
};

INSTANTIATE_TEST_SUITE_P(All,
                         SearchEnginePreconnectorEnabledOnlyBrowserTest,
                         ::testing::Combine(::testing::Bool(),
                                            ::testing::Bool()));

IN_PROC_BROWSER_TEST_P(SearchEnginePreconnectorEnabledOnlyBrowserTest,
                       AllowedSearch) {
  constexpr char16_t kShortName[] = u"test";
  constexpr char kSearchURL[] = "/anchors_different_area.html?q={searchTerms}";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data;
  data.SetShortName(kShortName);
  data.SetKeyword(data.short_name());
  data.SetURL(GetTestURL(kSearchURL).spec());
  data.preconnect_to_search_url = false;

  // Set the DSE to the test URL.
  TemplateURL* template_url = model->Add(std::make_unique<TemplateURL>(data));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  template_url = model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Preconnect should occur for Google search.
  EXPECT_EQ(1, preresolve_counts_[search_url]);

  // No preconnects should have been issued for the test URL.
  EXPECT_EQ(0, preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);
}

}  // namespace

class SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest
    : public SearchEnginePreconnectorBrowserTest,
      public testing::WithParamInterface<bool> {
 public:
  SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest() {
    std::vector<base::test::FeatureRefAndParams> enabled_features{
        {features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
        {net::features::kSearchEnginePreconnectInterval,
         {{"preconnect_interval", "0"}}},
        {net::features::kSearchEnginePreconnect2,
         {{"FallbackInLowPowerMode", "true"}}}};
    battery::OverrideIsBatterySaverEnabledForTesting(false);

    std::vector<base::test::FeatureRef> disabled_features{
        {features::kAdjustPreconnectRetryInterval}};

    if (PreconnectFromKeyedServiceEnabled()) {
      enabled_features.push_back(
          {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
    } else {
      disabled_features.emplace_back(features::kPreconnectFromKeyedService);
    }

    feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                disabled_features);
  }

  bool PreconnectFromKeyedServiceEnabled() const override { return GetParam(); }

  void OnPreresolveFinished(
      const GURL& url,
      const net::NetworkAnonymizationKey& network_anonymization_key,
      mojo::PendingRemote<network::mojom::ConnectionChangeObserverClient>&
          observer,
      bool success) override {
    // Take the observer so that we can manually send mojo message.
    if (observer.is_valid() && !remote_.is_bound()) {
      remote_.Bind(std::move(observer));
    }

    SearchEnginePreconnectorBrowserTest::OnPreresolveFinished(
        url, network_anonymization_key, observer, success);
  }
};

class SearchEnginePreconnectorWithResetConnectionFailureOnSessionUsedBrowserTest
    : public SearchEnginePreconnectorBrowserTest,
      public testing::WithParamInterface<bool> {
 public:
  SearchEnginePreconnectorWithResetConnectionFailureOnSessionUsedBrowserTest() {
    std::vector<base::test::FeatureRefAndParams> enabled_features{
        {features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
        {net::features::kSearchEnginePreconnectInterval,
         {{"preconnect_interval", "0"}}},
        {net::features::kSearchEnginePreconnect2,
         {{"FallbackInLowPowerMode", "true"}}},
        {features::kResetConnectionFailureOnSessionUsed, {}}};
    battery::OverrideIsBatterySaverEnabledForTesting(false);

    std::vector<base::test::FeatureRef> disabled_features{
        {features::kAdjustPreconnectRetryInterval}};

    if (PreconnectFromKeyedServiceEnabled()) {
      enabled_features.push_back(
          {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
    } else {
      disabled_features.emplace_back(features::kPreconnectFromKeyedService);
    }

    feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                disabled_features);
  }

  bool PreconnectFromKeyedServiceEnabled() const override { return GetParam(); }

  void OnPreresolveFinished(
      const GURL& url,
      const net::NetworkAnonymizationKey& network_anonymization_key,
      mojo::PendingRemote<network::mojom::ConnectionChangeObserverClient>&
          observer,
      bool success) override {
    if (observer.is_valid() && !remote_.is_bound()) {
      remote_.Bind(std::move(observer));
    }

    SearchEnginePreconnectorBrowserTest::OnPreresolveFinished(
        url, network_anonymization_key, observer, success);
  }
};

INSTANTIATE_TEST_SUITE_P(
    All,
    SearchEnginePreconnectorWithResetConnectionFailureOnSessionUsedBrowserTest,
    ::testing::Bool());

INSTANTIATE_TEST_SUITE_P(
    All,
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    ::testing::Bool());

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearch) {
  constexpr char16_t kShortName[] = u"test";
  constexpr char kSearchURL[] = "/anchors_different_area.html?q={searchTerms}";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data;
  data.SetShortName(kShortName);
  data.SetKeyword(data.short_name());
  data.SetURL(GetTestURL(kSearchURL).spec());
  data.preconnect_to_search_url = false;

  // Set the DSE to the test URL.
  TemplateURL* template_url = model->Add(std::make_unique<TemplateURL>(data));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  template_url = model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Preconnect should occur for Google search.
  EXPECT_LE(1, preresolve_counts_[search_url]);

  // No preconnects should have been issued for the test URL.
  EXPECT_EQ(0, preresolve_counts_[GetTestURL("/").DeprecatedGetOriginAsURL()]);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnClose) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->SetIsShortSessionForTesting(
      /*is_short_session=*/false);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a Session Close. This should trigger a reattempt.
  GetSearchEnginePreconnector()->OnSessionClosed(
      /*was_ever_used_to_create_streams=*/true);
  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);

  // Since this is not a short session, we should be resetting the value.
  EXPECT_EQ(0, GetSearchEnginePreconnector()
                   ->GetConsecutiveConnectionFailureForTesting());
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnCloseWithShortSession) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->SetIsShortSessionForTesting(
      /*is_short_session=*/true);

  int failure_before_testing =
      GetSearchEnginePreconnector()
          ->GetConsecutiveConnectionFailureForTesting();

  base::HistogramTester histogram_tester;
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a Session Close. This should trigger a reattempt.
  GetSearchEnginePreconnector()->OnSessionClosed(
      /*was_ever_used_to_create_streams=*/false);
  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);

  // Since this is a short session and not used, we should increment the value.
  EXPECT_EQ(1, GetSearchEnginePreconnector()
                       ->GetConsecutiveConnectionFailureForTesting() -
                   failure_before_testing);

  histogram_tester.ExpectUniqueSample(
      "NavigationPredictor.SearchEnginePreconnector."
      "TriggerEvent",
      static_cast<int>(
          SearchEnginePreconnector::PreconnectTriggerEvent::kSessionClosed),
      1);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithResetConnectionFailureOnSessionUsedBrowserTest,
    PreconnectSearchAfterOnCloseWithShortSessionAndUsed) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->SetIsShortSessionForTesting(
      /*is_short_session=*/true);

  // Artificially increment failure count to test if it resets.
  GetSearchEnginePreconnector()->SetConsecutiveFailureForTesting(1);
  ASSERT_EQ(1, GetSearchEnginePreconnector()
                   ->GetConsecutiveConnectionFailureForTesting());

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a Session Close with was_ever_used_to_create_streams=true.
  // This should trigger a reattempt and reset failure count.
  GetSearchEnginePreconnector()->OnSessionClosed(
      /*was_ever_used_to_create_streams=*/true);
  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);

  // Since was_ever_used_to_create_streams=true, failure count should be reset
  // to 0.
  EXPECT_EQ(0, GetSearchEnginePreconnector()
                   ->GetConsecutiveConnectionFailureForTesting());
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnCloseWithShortSessionAndUsed_FeatureDisabled) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->SetIsShortSessionForTesting(
      /*is_short_session=*/true);

  // Artificially increment failure count to test if it resets.
  GetSearchEnginePreconnector()->SetConsecutiveFailureForTesting(1);
  ASSERT_EQ(1, GetSearchEnginePreconnector()
                   ->GetConsecutiveConnectionFailureForTesting());

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a Session Close with was_ever_used_to_create_streams=true.
  // This should trigger a reattempt but NOT reset failure count because feature
  // is disabled.
  GetSearchEnginePreconnector()->OnSessionClosed(
      /*was_ever_used_to_create_streams=*/true);
  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);

  // Since feature is disabled, failure count should NOT be reset, it should
  // increment!
  EXPECT_EQ(2, GetSearchEnginePreconnector()
                   ->GetConsecutiveConnectionFailureForTesting());
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnFailure) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  int failure_before_testing =
      GetSearchEnginePreconnector()
          ->GetConsecutiveConnectionFailureForTesting();

  base::HistogramTester histogram_tester;
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a connection failure. This should trigger a reattempt.
  GetSearchEnginePreconnector()->OnConnectionFailed();
  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);

  // Since this is a short session, we should be resetting the value.
  EXPECT_EQ(1, GetSearchEnginePreconnector()
                       ->GetConsecutiveConnectionFailureForTesting() -
                   failure_before_testing);

  histogram_tester.ExpectUniqueSample(
      "NavigationPredictor.SearchEnginePreconnector."
      "TriggerEvent",
      static_cast<int>(
          SearchEnginePreconnector::PreconnectTriggerEvent::kConnectionFailed),
      1);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnConnect) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  base::HistogramTester histogram_tester;
  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a new connection. This should trigger a reattempt.
  GetSearchEnginePreconnector()->OnNetworkEvent(
      net::NetworkChangeEvent::kConnected);
  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);

  histogram_tester.ExpectUniqueSample(
      "NavigationPredictor.SearchEnginePreconnector."
      "TriggerEvent",
      static_cast<int>(
          SearchEnginePreconnector::PreconnectTriggerEvent::kNetworkEvent),
      1);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    CalculateBackoffMultiplier) {
  GetSearchEnginePreconnector()->StopPreconnecting();

  int failures = 0;
  for (; failures < std::numeric_limits<int32_t>::digits; failures++) {
    GetSearchEnginePreconnector()->SetConsecutiveFailureForTesting(failures);
    ASSERT_EQ(failures, GetSearchEnginePreconnector()
                            ->GetConsecutiveConnectionFailureForTesting());
    ASSERT_EQ(1 << failures,
              GetSearchEnginePreconnector()->CalculateBackoffMultiplier());
  }

  GetSearchEnginePreconnector()->SetConsecutiveFailureForTesting(failures);
  ASSERT_EQ(failures, GetSearchEnginePreconnector()
                          ->GetConsecutiveConnectionFailureForTesting());
  ASSERT_EQ(1 << (std::numeric_limits<int32_t>::digits - 1),
            GetSearchEnginePreconnector()->CalculateBackoffMultiplier());
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnClosedViaMojoPipe) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a new connection. This should trigger a reattempt.
  remote_->OnSessionClosed(/*was_ever_used_to_create_streams=*/true);

  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnFailureViaMojoPipe) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a new connection. This should trigger a reattempt.
  remote_->OnConnectionFailed();

  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchAfterOnNetworkEventViaMojoPipoe) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually trigger a new connection. This should trigger a reattempt.
  remote_->OnNetworkEvent(net::NetworkChangeEvent::kConnected);

  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    PreconnectSearchOnMojoPipeDisconnect) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  // Manually unbind to trigger a new preconnect
  ASSERT_TRUE(remote_.is_bound());
  remote_.Unbind().reset();

  WaitForPreresolveCountForURL(search_url, 2);

  // Preconnect should occur for Google search.
  EXPECT_EQ(2, preresolve_counts_[search_url]);
}

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest,
    CheckConnectionKeepAliveConfig) {
  auto config = GetSearchEnginePreconnector()->GetConnectionKeepAliveConfig();
  EXPECT_TRUE(config.enable_connection_keep_alive);

  battery::OverrideIsBatterySaverEnabledForTesting(true);

  config = GetSearchEnginePreconnector()->GetConnectionKeepAliveConfig();
  EXPECT_FALSE(config.enable_connection_keep_alive);
  EXPECT_EQ(config.idle_timeout_in_seconds, 60);
}

class SearchEnginePreconnectorWithBindReceiversEverytimeFeatureBrowserTest
    : public SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest {
 public:
  SearchEnginePreconnectorWithBindReceiversEverytimeFeatureBrowserTest() {
    feature_list_.Reset();
    std::vector<base::test::FeatureRefAndParams> enabled_features{
        {features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
        {net::features::kSearchEnginePreconnectInterval,
         {{"preconnect_interval", "0"}}},
        {net::features::kSearchEnginePreconnect2,
         {{"FallbackInLowPowerMode", "true"}}},
        {features::kRebindPreconnectReceivers, {}}};
    battery::OverrideIsBatterySaverEnabledForTesting(false);

    std::vector<base::test::FeatureRef> disabled_features;

    if (PreconnectFromKeyedServiceEnabled()) {
      enabled_features.push_back(
          {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
    } else {
      disabled_features.emplace_back(features::kPreconnectFromKeyedService);
    }

    feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                disabled_features);
  }

  void OnPreresolveFinished(
      const GURL& url,
      const net::NetworkAnonymizationKey& network_anonymization_key,
      mojo::PendingRemote<network::mojom::ConnectionChangeObserverClient>&
          observer,
      bool success) override {
    if (observer.is_valid()) {
      // This will disconnect the old remote if it is bound.
      remote_.reset();
      remote_.Bind(std::move(observer));
    }

    const GURL origin = url.DeprecatedGetOriginAsURL();
    if (!preresolve_counts_.contains(origin)) {
      return;
    }

    if (origin == GetTestURL("/").DeprecatedGetOriginAsURL()) {
      EXPECT_TRUE(success);
    }

    ++preresolve_counts_[origin];
    if (run_loops_[origin]) {
      run_loops_[origin]->Quit();
    }
  }
};

INSTANTIATE_TEST_SUITE_P(
    All,
    SearchEnginePreconnectorWithBindReceiversEverytimeFeatureBrowserTest,
    ::testing::Bool());

// TODO(crbug.com/506949513): Flaky on Linux.
#if BUILDFLAG(IS_LINUX)
#define MAYBE_BindNewRemoteOnEachPreconnect \
  DISABLED_BindNewRemoteOnEachPreconnect
#else
#define MAYBE_BindNewRemoteOnEachPreconnect BindNewRemoteOnEachPreconnect
#endif
IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithBindReceiversEverytimeFeatureBrowserTest,
    MAYBE_BindNewRemoteOnEachPreconnect) {
  constexpr char16_t kShortName[] = u"test";
  TemplateURLService* model =
      TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
  ASSERT_TRUE(model);
  search_test_utils::WaitForTemplateURLServiceToLoad(model);
  ASSERT_TRUE(model->loaded());

  TemplateURLData data_allowed_search;
  data_allowed_search.SetShortName(kShortName);
  data_allowed_search.SetKeyword(data_allowed_search.short_name());
  data_allowed_search.SetURL(kGoogleSearch);
  data_allowed_search.preconnect_to_search_url = true;

  auto* template_url =
      model->Add(std::make_unique<TemplateURL>(data_allowed_search));
  ASSERT_TRUE(template_url);
  model->SetUserSelectedDefaultSearchProvider(template_url);

  GetSearchEnginePreconnector()->StartPreconnecting(
      /*with_startup_delay=*/false);

  const GURL search_url = template_url->GenerateSearchURL({});
  WaitForPreresolveCountForURL(search_url, 1);

  ASSERT_TRUE(remote_.is_bound());
  mojo::Remote<network::mojom::ConnectionChangeObserverClient> remote_1 =
      std::move(remote_);

  ASSERT_FALSE(remote_.is_bound());
  ASSERT_TRUE(remote_1.is_bound());

  base::RunLoop disconnect_run_loop;
  remote_1.set_disconnect_handler(disconnect_run_loop.QuitClosure());

  GetSearchEnginePreconnector()->OnSessionClosed(
      /*was_ever_used_to_create_streams=*/true);
  WaitForPreresolveCountForURL(search_url, 2);

  disconnect_run_loop.Run();
  ASSERT_FALSE(remote_1.is_connected());

  remote_1.reset_on_disconnect();

  EXPECT_FALSE(remote_1.is_bound());
  EXPECT_TRUE(remote_.is_bound());
}

class
    SearchEnginePreconnectorWithAdjustPreconnectRetryIntervalFeatureBrowserTest
    : public SearchEnginePreconnectorWithPreconnect2FeatureBrowserTest {
 public:
  constexpr static double kBackoffMultiplier = 2.0;
  SearchEnginePreconnectorWithAdjustPreconnectRetryIntervalFeatureBrowserTest() {
    feature_list_.Reset();
    std::vector<base::test::FeatureRefAndParams> enabled_features{
        {features::kPreconnectToSearch, {{"startup_delay_ms", "1000000"}}},
        {net::features::kSearchEnginePreconnectInterval,
         {{"preconnect_interval", "0"}}},
        {net::features::kSearchEnginePreconnect2,
         {{"FallbackInLowPowerMode", "true"}}},
        {features::kAdjustPreconnectRetryInterval,
         {{"kPreconnectBackoffMultiplier",
           base::NumberToString(kBackoffMultiplier)}}}};

    std::vector<base::test::FeatureRef> disabled_features;

    if (PreconnectFromKeyedServiceEnabled()) {
      enabled_features.push_back(
          {features::kPreconnectFromKeyedService, {{"run_on_otr", "false"}}});
    } else {
      disabled_features.emplace_back(features::kPreconnectFromKeyedService);
    }

    feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                disabled_features);
  }
};

INSTANTIATE_TEST_SUITE_P(
    All,
    SearchEnginePreconnectorWithAdjustPreconnectRetryIntervalFeatureBrowserTest,
    ::testing::Bool());

IN_PROC_BROWSER_TEST_P(
    SearchEnginePreconnectorWithAdjustPreconnectRetryIntervalFeatureBrowserTest,
    CalculateBackoffMultiplier) {
  GetSearchEnginePreconnector()->StopPreconnecting();

  for (int failures = 0; base::ClampedNumeric<int32_t>(
                             std::pow(kBackoffMultiplier, failures - 1)) <
                         std::numeric_limits<int32_t>::max();
       failures++) {
    GetSearchEnginePreconnector()->SetConsecutiveFailureForTesting(failures);
    ASSERT_EQ(failures, GetSearchEnginePreconnector()
                            ->GetConsecutiveConnectionFailureForTesting());
    ASSERT_EQ(base::ClampedNumeric<int32_t>(
                  std::pow(kBackoffMultiplier, failures - 1)),
              GetSearchEnginePreconnector()->CalculateBackoffMultiplier());
  }
}

#if BUILDFLAG(ENABLE_DEVICE_BOUND_SESSIONS)
class SearchEnginePreconnectorDeviceBoundSessionBrowserTest
    : public SearchEnginePreconnectorBrowserTest {
 public:
  SearchEnginePreconnectorDeviceBoundSessionBrowserTest() {
    feature_list_.InitWithFeaturesAndParameters(
        {{features::kDeviceBoundSessionsDsePrewarmer, {}},
         {features::kPreconnectToSearch,
          {{"skip_in_background", "false"}, {"startup_delay_ms", "0"}}}},
        {});
  }

  void SetUpOnMainThread() override {
    SearchEnginePreconnectorBrowserTest::SetUpOnMainThread();
    TemplateURLService* model =
        TemplateURLServiceFactory::GetForProfile(browser()->GetProfile());
    ASSERT_TRUE(model);
    search_test_utils::WaitForTemplateURLServiceToLoad(model);
    ASSERT_TRUE(model->loaded());
  }

  void TearDownOnMainThread() override {
    browser()
        ->GetProfile()
        ->GetDefaultStoragePartition()
        ->OverrideDeviceBoundSessionManagerForTesting(nullptr);
    SearchEnginePreconnectorBrowserTest::TearDownOnMainThread();
  }

 private:
  base::test::ScopedFeatureList feature_list_;
};

IN_PROC_BROWSER_TEST_F(SearchEnginePreconnectorDeviceBoundSessionBrowserTest,
                       PrewarmsDeviceBoundSessionOnStart) {
  auto mock_manager =
      std::make_unique<network::MockDeviceBoundSessionManager>();

  base::RunLoop run_loop;
  EXPECT_CALL(*mock_manager, PrewarmSessionsForUrl)
      .WillOnce([&run_loop](const GURL& url,
                            network::mojom::DeviceBoundSessionManager::
                                PrewarmSessionsForUrlCallback callback) {
        std::move(callback).Run({}, std::nullopt);
        run_loop.Quit();
      });

  auto* preconnector =
      SearchEnginePreconnectorKeyedServiceFactory::GetForProfile(
          browser()->GetProfile());
  ASSERT_TRUE(preconnector);

  // Stop preconnector so the pre-existing pre-warmer is destroyed before we
  // replace the proxy.
  preconnector->StopPreconnecting();

  browser()
      ->GetProfile()
      ->GetDefaultStoragePartition()
      ->OverrideDeviceBoundSessionManagerForTesting(std::move(mock_manager));

  preconnector->StartPreconnecting(/*with_startup_delay=*/true);

  run_loop.Run();
  EXPECT_TRUE(preconnector->HasDeviceBoundSessionPrewarmerForTesting());
}

IN_PROC_BROWSER_TEST_F(SearchEnginePreconnectorDeviceBoundSessionBrowserTest,
                       DoesNotPrewarmWhenSearchSuggestDisabled) {
  browser()->GetProfile()->GetPrefs()->SetBoolean(prefs::kSearchSuggestEnabled,
                                                  false);

  auto mock_manager =
      std::make_unique<network::MockDeviceBoundSessionManager>();

  EXPECT_CALL(*mock_manager, PrewarmSessionsForUrl).Times(0);

  auto* preconnector =
      SearchEnginePreconnectorKeyedServiceFactory::GetForProfile(
          browser()->GetProfile());
  ASSERT_TRUE(preconnector);
  preconnector->StopPreconnecting();

  browser()
      ->GetProfile()
      ->GetDefaultStoragePartition()
      ->OverrideDeviceBoundSessionManagerForTesting(std::move(mock_manager));

  preconnector->StartPreconnecting(/*with_startup_delay=*/true);

  EXPECT_FALSE(preconnector->HasDeviceBoundSessionPrewarmerForTesting());
}

IN_PROC_BROWSER_TEST_F(SearchEnginePreconnectorDeviceBoundSessionBrowserTest,
                       PrewarmerCancelledWhenStopped) {
  auto mock_manager =
      std::make_unique<network::MockDeviceBoundSessionManager>();

  base::RunLoop run_loop;
  EXPECT_CALL(*mock_manager, PrewarmSessionsForUrl)
      .WillOnce([&run_loop](const GURL& url,
                            network::mojom::DeviceBoundSessionManager::
                                PrewarmSessionsForUrlCallback callback) {
        // Schedule next prewarm in 10 seconds.
        std::move(callback).Run({}, base::Time::Now() + base::Seconds(10));
        run_loop.Quit();
      });

  auto* preconnector =
      SearchEnginePreconnectorKeyedServiceFactory::GetForProfile(
          browser()->GetProfile());
  ASSERT_TRUE(preconnector);
  preconnector->StopPreconnecting();

  browser()
      ->GetProfile()
      ->GetDefaultStoragePartition()
      ->OverrideDeviceBoundSessionManagerForTesting(std::move(mock_manager));

  preconnector->StartPreconnecting(/*with_startup_delay=*/true);
  run_loop.Run();
  EXPECT_TRUE(preconnector->HasDeviceBoundSessionPrewarmerForTesting());

  // Stop preconnector (e.g. when app goes to background).
  preconnector->StopPreconnecting();
  EXPECT_FALSE(preconnector->HasDeviceBoundSessionPrewarmerForTesting());
}
#endif  // BUILDFLAG(ENABLE_DEVICE_BOUND_SESSIONS)