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

#include "content/browser/preloading/prerender/prerender_host_registry.h"

#include <cstdint>

#include "base/memory_coordinator/memory_coordinator_features.h"
#include "base/memory_coordinator/test_memory_consumer_registry.h"
#include "base/strings/string_number_conversions.h"
#include "base/test/bind.h"
#include "base/test/scoped_feature_list.h"
#include "components/variations/scoped_variations_ids_provider.h"
#include "content/browser/preloading/preload_pipeline_info_impl.h"
#include "content/browser/preloading/preloading.h"
#include "content/browser/preloading/preloading_confidence.h"
#include "content/browser/preloading/preloading_config.h"
#include "content/browser/preloading/prerender/prerender_features.h"
#include "content/browser/preloading/prerender/prerender_final_status.h"
#include "content/browser/preloading/prerender/prerender_host.h"
#include "content/browser/preloading/prerender/prerender_metrics.h"
#include "content/browser/preloading/speculation_rules/speculation_host_impl.h"
#include "content/browser/renderer_host/render_frame_host_impl.h"
#include "content/browser/site_instance_impl.h"
#include "content/public/browser/preload_pipeline_info.h"
#include "content/public/browser/render_frame_host.h"
#include "content/public/browser/storage_partition.h"
#include "content/public/test/preloading_test_util.h"
#include "content/public/test/prerender_test_util.h"
#include "content/public/test/test_browser_context.h"
#include "content/test/mock_commit_deferring_condition.h"
#include "content/test/navigation_simulator_impl.h"
#include "content/test/test_render_view_host.h"
#include "content/test/test_web_contents.h"
#include "net/base/load_flags.h"
#include "third_party/blink/public/common/features.h"
#include "third_party/blink/public/mojom/loader/mixed_content.mojom.h"
#include "third_party/blink/public/mojom/speculation_rules/speculation_rules.mojom.h"

namespace content {
namespace {

blink::mojom::SpeculationCandidatePtr CreatePrerenderCandidate(
    const GURL& url) {
  auto candidate = blink::mojom::SpeculationCandidate::New();
  candidate->action = blink::mojom::SpeculationAction::kPrerender;
  candidate->url = url;
  candidate->referrer = blink::mojom::Referrer::New();
  candidate->eagerness = blink::mojom::SpeculationEagerness::kImmediate;
  candidate->tags = {std::nullopt};
  return candidate;
}

void SendCandidates(const std::vector<GURL>& urls,
                    mojo::Remote<blink::mojom::SpeculationHost>& remote) {
  std::vector<blink::mojom::SpeculationCandidatePtr> candidates;
  candidates.resize(urls.size());
  std::ranges::transform(urls, candidates.begin(), &CreatePrerenderCandidate);
  remote->UpdateSpeculationCandidates(std::move(candidates), false);
  remote.FlushForTesting();
}

void SendCandidate(const GURL& url,
                   mojo::Remote<blink::mojom::SpeculationHost>& remote) {
  SendCandidates({url}, remote);
}

std::unique_ptr<NavigationSimulatorImpl> CreateActivation(
    const GURL& prerendering_url,
    WebContentsImpl& web_contents) {
  std::unique_ptr<NavigationSimulatorImpl> navigation =
      NavigationSimulatorImpl::CreateRendererInitiated(
          prerendering_url, web_contents.GetPrimaryMainFrame());
  navigation->SetReferrer(blink::mojom::Referrer::New(
      web_contents.GetPrimaryMainFrame()->GetLastCommittedURL(),
      network::mojom::ReferrerPolicy::kStrictOriginWhenCrossOrigin));
  return navigation;
}

// Finish a prerendering navigation that was already started with
// CreateAndStartHost().
void CommitPrerenderNavigation(PrerenderHost& host) {
  // Normally we could use EmbeddedTestServer to provide a response, but these
  // tests use RenderViewHostImplTestHarness so the load goes through a
  // TestNavigationURLLoader which we don't have access to in order to
  // complete. Use NavigationSimulator to finish the navigation.
  FrameTreeNode* ftn = FrameTreeNode::From(host.GetPrerenderedMainFrameHost());
  std::unique_ptr<NavigationSimulator> sim =
      NavigationSimulatorImpl::CreateFromPendingInFrame(ftn);
  sim->Commit();
  EXPECT_TRUE(host.is_ready_for_activation());
}

class PrerenderHostRegistryTest : public RenderViewHostImplTestHarness {
 public:
  PrerenderHostRegistryTest()
      : scoped_variations_ids_provider_(
            variations::test::ScopedVariationsIdsProvider(
                variations::VariationsIdsProvider::Mode::kUseSignedInState)) {
    scoped_feature_list_prerender2_fallback_.InitWithFeaturesAndParameters(
        {
            {
                features::kPrerender2FallbackPrefetchSpecRules,
                {
                    {
                        features::
                            kPrerender2FallbackPrefetchUseBlockUntilHeadTimetout
                                .name,
                        "false",
                    },
                    {
                        features::kPrerender2FallbackPrefetchSchedulerPolicy
                            .name,
                        "NotUse",
                    },
                },
            },
        },
        {});
  }
  ~PrerenderHostRegistryTest() override = default;

  void SetUp() override {
    RenderViewHostImplTestHarness::SetUp();

    web_contents_delegate_ =
        std::make_unique<test::ScopedPrerenderWebContentsDelegate>(*contents());
    contents()->NavigateAndCommit(GURL("https://example.com/"));
  }

  RenderFrameHostImpl* NavigatePrimaryPage(TestWebContents* web_contents,
                                           const GURL& dest_url) {
    std::unique_ptr<NavigationSimulatorImpl> navigation =
        NavigationSimulatorImpl::CreateRendererInitiated(
            dest_url, web_contents->GetPrimaryMainFrame());
    navigation->SetTransition(ui::PAGE_TRANSITION_LINK);
    navigation->Start();
    navigation->Commit();
    RenderFrameHostImpl* render_frame_host =
        web_contents->GetPrimaryMainFrame();
    EXPECT_EQ(render_frame_host->GetLastCommittedURL(), dest_url);
    return render_frame_host;
  }

  // Helper method to test the navigation param matching logic which allows a
  // prerender host to be used in a potential activation navigation only if its
  // params match the potential activation navigation params. Use setup_callback
  // to set the parameters. Returns true if the host was selected as a
  // potential candidate for activation, and false otherwise.
  [[nodiscard]] bool CheckIsActivatedForParams(
      base::OnceCallback<void(NavigationSimulatorImpl*)> setup_callback) {
    RenderFrameHostImpl* render_frame_host = contents()->GetPrimaryMainFrame();

    const GURL kPrerenderingUrl("https://example.com/next");
    registry().CreateAndStartHost(GeneratePrerenderAttributes(
        kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
        blink::mojom::SpeculationEagerness::kImmediate, render_frame_host));
    PrerenderHost* prerender_host =
        registry().FindHostByUrlForTesting(kPrerenderingUrl);
    CommitPrerenderNavigation(*prerender_host);

    std::unique_ptr<NavigationSimulatorImpl> navigation =
        NavigationSimulatorImpl::CreateRendererInitiated(kPrerenderingUrl,
                                                         render_frame_host);
    // Set a default referrer policy that matches the initial prerender
    // navigation.
    // TODO(falken): Fix NavigationSimulatorImpl to do this itself.
    navigation->SetReferrer(blink::mojom::Referrer::New(
        contents()->GetPrimaryMainFrame()->GetLastCommittedURL(),
        network::mojom::ReferrerPolicy::kStrictOriginWhenCrossOrigin));

    // Change a parameter to differentiate the activation request from the
    // prerendering request.
    std::move(setup_callback).Run(navigation.get());
    navigation->Start();
    NavigationRequest* navigation_request = navigation->GetNavigationHandle();
    // Use is_running_potential_prerender_activation_checks() instead of
    // IsPrerenderedPageActivation() because the NavigationSimulator does not
    // proceed past CommitDeferringConditions on potential activations,
    // so IsPrerenderedPageActivation() will fail with a CHECK because
    // prerender_frame_tree_node_id_ is not populated.
    // TODO(crbug.com/40784651): Fix NavigationSimulator to wait for
    // commit deferring conditions as it does throttles.
    return navigation_request
        ->is_running_potential_prerender_activation_checks();
  }

  // Helper method to perform a prerender activation that includes specialized
  // handling or setup on the initial prerender navigation via the
  // setup_callback parameter.
  void SetupPrerenderAndCommit(
      base::OnceCallback<void(NavigationSimulatorImpl*)> setup_callback) {
    const GURL kPrerenderingUrl("https://example.com/next");
    const PrerenderHostId prerender_host_id =
        registry().CreateAndStartHost(GeneratePrerenderAttributes(
            kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
            blink::mojom::SpeculationEagerness::kImmediate,
            contents()->GetPrimaryMainFrame()));
    ASSERT_TRUE(prerender_host_id);
    PrerenderHost* prerender_host =
        registry().FindNonReservedHostById(prerender_host_id);

    // Complete the initial prerender navigation.
    FrameTreeNode* ftn =
        FrameTreeNode::From(prerender_host->GetPrerenderedMainFrameHost());
    std::unique_ptr<NavigationSimulatorImpl> sim =
        NavigationSimulatorImpl::CreateFromPendingInFrame(ftn);
    std::move(setup_callback).Run(sim.get());
    sim->Commit();
    EXPECT_TRUE(prerender_host->is_ready_for_activation());

    // Activate the prerendered page.
    contents()->ActivatePrerenderedPage(kPrerenderingUrl);
  }

  PrerenderAttributes GeneratePrerenderAttributes(
      const GURL& url,
      PreloadingTriggerType trigger_type,
      const std::string& histogram_suffix,
      std::optional<blink::mojom::SpeculationEagerness> eagerness,
      RenderFrameHostImpl* rfh) {
    switch (trigger_type) {
      case PreloadingTriggerType::kSpeculationRule:
      case PreloadingTriggerType::kSpeculationRuleFromIsolatedWorld:
      case PreloadingTriggerType::kSpeculationRuleFromAutoSpeculationRules:
        return PrerenderAttributes(
            url, trigger_type, histogram_suffix,
            std::make_optional(SpeculationRulesParams(
                blink::mojom::SpeculationTargetHint::kNoHint,
                eagerness.value_or(
                    blink::mojom::SpeculationEagerness::kImmediate),
                SpeculationRulesTags())),
            Referrer(),
            /*no_vary_search_hint=*/std::nullopt, rfh, contents()->GetWeakPtr(),
            ui::PAGE_TRANSITION_LINK,
            /*should_warm_up_compositor=*/false,
            /*should_prepare_paint_tree=*/false,
            blink::mojom::SpeculationAction::kPrerender,
            /*url_match_predicate=*/{},
            /*prerender_navigation_handle_callback=*/{},
            PreloadPipelineInfoImpl::Create(
                /*planned_max_preloading_type=*/PreloadingType::kPrerender),
            /*allow_reuse=*/false,
            /*form_submission=*/false);
      case PreloadingTriggerType::kEmbedder:
        return PrerenderAttributes(
            url, trigger_type, histogram_suffix,
            /*speculation_rules_params=*/std::nullopt, Referrer(),
            /*no_vary_search_hint=*/std::nullopt,
            /*initiator_render_frame_host=*/nullptr, contents()->GetWeakPtr(),
            ui::PageTransitionFromInt(ui::PAGE_TRANSITION_TYPED |
                                      ui::PAGE_TRANSITION_FROM_ADDRESS_BAR),
            /*should_warm_up_compositor=*/false,
            /*should_prepare_paint_tree=*/false,
            blink::mojom::SpeculationAction::kPrerender,
            /*url_match_predicate=*/{},
            /*prerender_navigation_handle_callback=*/{},
            PreloadPipelineInfoImpl::Create(
                /*planned_max_preloading_type=*/PreloadingType::kPrerender),
            /*allow_reuse=*/false,
            /*form_submission=*/false);
    }
  }

  void ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus status,
      base::HistogramBase::Count32 count = 1) {
    histogram_tester_.ExpectUniqueSample(
        "Prerender.Experimental.PrerenderHostFinalStatus.SpeculationRule",
        status, count);
  }

  void ExpectBucketCountOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus status,
      base::HistogramBase::Count32 count = 1) {
    histogram_tester_.ExpectBucketCount(
        "Prerender.Experimental.PrerenderHostFinalStatus.SpeculationRule",
        status, count);
  }

  void ExpectUniqueSampleOfEmbedderFinalStatus(
      PrerenderFinalStatus status,
      const std::string& histogram_suffix,
      base::HistogramBase::Count32 count = 1) {
    histogram_tester_.ExpectUniqueSample(
        "Prerender.Experimental.PrerenderHostFinalStatus.Embedder_" +
            histogram_suffix,
        status, count);
  }

  void ExpectBucketCountOfEmbedderFinalStatus(
      PrerenderFinalStatus status,
      const std::string& histogram_suffix,
      base::HistogramBase::Count32 count = 1) {
    histogram_tester_.ExpectBucketCount(
        "Prerender.Experimental.PrerenderHostFinalStatus.Embedder_" +
            histogram_suffix,
        status, count);
  }

  void ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch result,
      base::HistogramBase::Count32 count = 1) {
    histogram_tester_.ExpectUniqueSample(
        "Prerender.Experimental.ActivationNavigationParamsMatch."
        "SpeculationRule",
        result, count);
  }

  void ExpectBucketCountOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch result,
      base::HistogramBase::Count32 count = 1) {
    histogram_tester_.ExpectBucketCount(
        "Prerender.Experimental.ActivationNavigationParamsMatch."
        "SpeculationRule",
        result, count);
  }

  PrerenderHostRegistry& registry() {
    return *contents()->GetPrerenderHostRegistry();
  }

  base::HistogramTester& histogram_tester() { return histogram_tester_; }

 private:
  test::ScopedPrerenderFeatureList scoped_feature_list_;
  base::test::ScopedFeatureList scoped_feature_list_prerender2_fallback_;
  base::HistogramTester histogram_tester_;

 protected:
  base::TestMemoryConsumerRegistry test_registry_;
  std::unique_ptr<test::ScopedPrerenderWebContentsDelegate>
      web_contents_delegate_;
  // Prevent `DCHECK(g_instance)` failure in
  // `VariationsIdsProvider::GetInstance()` via
  // `PrefetchContainer::MakeInitialResourceRequest()`.
  variations::test::ScopedVariationsIdsProvider scoped_variations_ids_provider_;
};

TEST_F(PrerenderHostRegistryTest, CreateAndStartHost_SpeculationRule) {
  const GURL kPrerenderingUrl("https://example.com/next");
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindNonReservedHostById(prerender_host_id);
  CommitPrerenderNavigation(*prerender_host);

  contents()->ActivatePrerenderedPage(kPrerenderingUrl);

  // "Navigation.TimeToActivatePrerender.SpeculationRule" histogram should be
  // recorded on every prerender activation.
  histogram_tester().ExpectTotalCount(
      "Navigation.TimeToActivatePrerender.SpeculationRule", 1u);
}

TEST_F(PrerenderHostRegistryTest, CreateAndStartHost_Embedder_DirectURLInput) {
  const GURL kPrerenderingUrl("https://example.com/next");
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kEmbedder, "DirectURLInput",
          std::nullopt, contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindNonReservedHostById(prerender_host_id);
  CommitPrerenderNavigation(*prerender_host);

  contents()->ActivatePrerenderedPageFromAddressBar(kPrerenderingUrl);

  // "Navigation.TimeToActivatePrerender.Embedder_DirectURLInput" histogram
  // should be recorded on every prerender activation.
  histogram_tester().ExpectTotalCount(
      "Navigation.TimeToActivatePrerender.Embedder_DirectURLInput", 1u);
}

TEST_F(PrerenderHostRegistryTest, CreateAndStartHost_PreloadingConfigHoldback) {
  content::test::PreloadingConfigOverride preloading_config_override;
  preloading_config_override.SetHoldback(
      PreloadingType::kPrerender,
      content_preloading_predictor::kSpeculationRules, true);
  const GURL kPrerenderingUrl("https://example.com/next");
  auto* preloading_data = PreloadingData::GetOrCreateForWebContents(contents());
  PreloadingURLMatchCallback same_url_matcher =
      PreloadingData::GetSameURLMatcher(kPrerenderingUrl);
  PreloadingAttempt* preloading_attempt = preloading_data->AddPreloadingAttempt(
      content_preloading_predictor::kSpeculationRules,
      PreloadingType::kPrerender, std::move(same_url_matcher),
      contents()->GetPrimaryMainFrame()->GetPageUkmSourceId());
  const PrerenderHostId prerender_host_id = registry().CreateAndStartHost(
      GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()),
      preloading_attempt);
  EXPECT_FALSE(prerender_host_id);
}

TEST_F(PrerenderHostRegistryTest,
       CreateAndStartHost_HoldbackOverride_Holdback) {
  const GURL kPrerenderingUrl("https://example.com/next");
  auto* preloading_data = PreloadingData::GetOrCreateForWebContents(contents());
  PreloadingURLMatchCallback same_url_matcher =
      PreloadingData::GetSameURLMatcher(kPrerenderingUrl);
  PreloadingAttempt* preloading_attempt = preloading_data->AddPreloadingAttempt(
      content_preloading_predictor::kSpeculationRules,
      PreloadingType::kPrerender, std::move(same_url_matcher),
      contents()->GetPrimaryMainFrame()->GetPageUkmSourceId());

  auto attributes = GeneratePrerenderAttributes(
      kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
      blink::mojom::SpeculationEagerness::kImmediate,
      contents()->GetPrimaryMainFrame());
  attributes.holdback_status_override = PreloadingHoldbackStatus::kHoldback;

  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(attributes, preloading_attempt);

  EXPECT_FALSE(prerender_host_id);
}

TEST_F(PrerenderHostRegistryTest, CreateAndStartHost_HoldbackOverride_Allowed) {
  content::test::PreloadingConfigOverride preloading_config_override;
  preloading_config_override.SetHoldback(
      PreloadingType::kPrerender,
      content_preloading_predictor::kSpeculationRules, true);
  const GURL kPrerenderingUrl("https://example.com/next");
  auto* preloading_data = PreloadingData::GetOrCreateForWebContents(contents());
  PreloadingURLMatchCallback same_url_matcher =
      PreloadingData::GetSameURLMatcher(kPrerenderingUrl);
  PreloadingAttempt* preloading_attempt = preloading_data->AddPreloadingAttempt(
      content_preloading_predictor::kSpeculationRules,
      PreloadingType::kPrerender, std::move(same_url_matcher),
      contents()->GetPrimaryMainFrame()->GetPageUkmSourceId());

  auto attributes = GeneratePrerenderAttributes(
      kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
      blink::mojom::SpeculationEagerness::kImmediate,
      contents()->GetPrimaryMainFrame());
  attributes.holdback_status_override = PreloadingHoldbackStatus::kAllowed;

  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(attributes, preloading_attempt);
  ASSERT_TRUE(prerender_host_id);

  PrerenderHost* prerender_host =
      registry().FindNonReservedHostById(prerender_host_id);
  CommitPrerenderNavigation(*prerender_host);

  contents()->ActivatePrerenderedPage(kPrerenderingUrl);

  // "Navigation.TimeToActivatePrerender.SpeculationRule" histogram should be
  // recorded on every prerender activation.
  histogram_tester().ExpectTotalCount(
      "Navigation.TimeToActivatePrerender.SpeculationRule", 1u);
}

TEST_F(PrerenderHostRegistryTest, CreateAndStartHostForSameURL) {
  const GURL kPrerenderingUrl("https://example.com/next");

  const PrerenderHostId prerender_host_id1 =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_TRUE(prerender_host_id1);
  PrerenderHost* prerender_host1 =
      registry().FindNonReservedHostById(prerender_host_id1);

  // Start the prerender host for the same URL. This second host should be
  // ignored, and the first host should still be findable.
  const PrerenderHostId prerender_host_id2 =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_FALSE(prerender_host_id2);
  EXPECT_EQ(registry().FindHostByUrlForTesting(kPrerenderingUrl),
            prerender_host1);
  CommitPrerenderNavigation(*prerender_host1);

  contents()->ActivatePrerenderedPage(kPrerenderingUrl);
}

// Tests that PrerenderHostRegistry limits the number of started prerenders
// to a specific number, and after once the prerender page was activated,
// PrerenderHostRegistry can start prerendering a new one.
TEST_F(PrerenderHostRegistryTest, NumberLimit_Activation) {
  std::vector<PrerenderHostId> prerender_host_ids;
  std::vector<GURL> prerendering_urls;
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders +
           1;
       i++) {
    const GURL prerendering_url("https://example.com/next" +
                                base::NumberToString(i));
    PrerenderHostId prerender_host_id =
        registry().CreateAndStartHost(GeneratePrerenderAttributes(
            prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
            blink::mojom::SpeculationEagerness::kImmediate,
            contents()->GetPrimaryMainFrame()));

    prerender_host_ids.push_back(prerender_host_id);
    prerendering_urls.push_back(prerendering_url);
  }

  // PrerenderHostRegistry should only start prerendering within the limit.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders;
       i++) {
    EXPECT_TRUE(prerender_host_ids[i]);
  }
  EXPECT_FALSE(
      prerender_host_ids[PrerenderHostRegistry::
                             kMaxRunningSpeculationRulesImmediatePrerenders]);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded);

  // Activate the first prerender.
  PrerenderHost* prerender_host =
      registry().FindHostByUrlForTesting(prerendering_urls[0]);
  CommitPrerenderNavigation(*prerender_host);
  contents()->ActivatePrerenderedPage(prerendering_urls[0]);

  // After the first prerender page was activated, PrerenderHostRegistry can
  // start prerendering a new one.
  PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          prerendering_urls[PrerenderHostRegistry::
                                kMaxRunningSpeculationRulesImmediatePrerenders],
          PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_TRUE(prerender_host_id);
  ExpectBucketCountOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded);
}

// Tests that PrerenderHostRegistry limits the number of started prerenders
// to a specific number, and new candidates can be processed after the initiator
// page navigates to a new same-origin page.
TEST_F(PrerenderHostRegistryTest, NumberLimit_SameOriginNavigateAway) {
  RenderFrameHostImpl* render_frame_host = contents()->GetPrimaryMainFrame();
  ASSERT_TRUE(render_frame_host);

  mojo::Remote<blink::mojom::SpeculationHost> remote1;
  SpeculationHostImpl::Bind(render_frame_host,
                            remote1.BindNewPipeAndPassReceiver());
  ASSERT_TRUE(remote1.is_connected());

  std::vector<GURL> prerendering_urls;
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders +
           1;
       i++) {
    prerendering_urls.emplace_back("https://example.com/next" +
                                   base::NumberToString(i));
  }
  SendCandidates(prerendering_urls, remote1);

  // PrerenderHostRegistry should only start prerenderings within the limit.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders;
       i++) {
    ASSERT_NE(registry().FindHostByUrlForTesting(prerendering_urls[i]),
              nullptr);
  }
  ASSERT_EQ(registry().FindHostByUrlForTesting(
                prerendering_urls
                    [PrerenderHostRegistry::
                         kMaxRunningSpeculationRulesImmediatePrerenders]),
            nullptr);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded);

  // The initiator document navigates away.
  render_frame_host =
      NavigatePrimaryPage(contents(), GURL("https://example.com/elsewhere"));

  // After the initiator page navigates away, the started prerendering should be
  // cancelled, and PrerenderHostRegistry can start prerendering a new one.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders +
           1;
       i++) {
    EXPECT_EQ(registry().FindHostByUrlForTesting(prerendering_urls[i]),
              nullptr);
  }
  mojo::Remote<blink::mojom::SpeculationHost> remote2;
  SpeculationHostImpl::Bind(render_frame_host,
                            remote2.BindNewPipeAndPassReceiver());
  SendCandidate(
      prerendering_urls[PrerenderHostRegistry::
                            kMaxRunningSpeculationRulesImmediatePrerenders],
      remote2);

  EXPECT_NE(registry().FindHostByUrlForTesting(
                prerendering_urls
                    [PrerenderHostRegistry::
                         kMaxRunningSpeculationRulesImmediatePrerenders]),
            nullptr);
  ExpectBucketCountOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded);
}

// Tests that PrerenderHostRegistry limits the number of started prerenders
// to a specific number, and new candidates can be processed after the initiator
// page navigates to a new cross-origin page.
TEST_F(PrerenderHostRegistryTest, NumberLimit_CrossOriginNavigateAway) {
  RenderFrameHostImpl* render_frame_host = contents()->GetPrimaryMainFrame();
  ASSERT_TRUE(render_frame_host);

  mojo::Remote<blink::mojom::SpeculationHost> remote1;
  SpeculationHostImpl::Bind(render_frame_host,
                            remote1.BindNewPipeAndPassReceiver());
  ASSERT_TRUE(remote1.is_connected());

  std::vector<GURL> prerendering_urls;
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders +
           1;
       i++) {
    prerendering_urls.emplace_back("https://example.com/next" +
                                   base::NumberToString(i));
  }
  SendCandidates(prerendering_urls, remote1);

  // PrerenderHostRegistry should only start prerenderings within the limit.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders;
       i++) {
    ASSERT_NE(registry().FindHostByUrlForTesting(prerendering_urls[i]),
              nullptr);
  }
  ASSERT_EQ(registry().FindHostByUrlForTesting(
                prerendering_urls
                    [PrerenderHostRegistry::
                         kMaxRunningSpeculationRulesImmediatePrerenders]),
            nullptr);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded);

  // The initiator document navigates away to a cross-origin page.
  render_frame_host =
      NavigatePrimaryPage(contents(), GURL("https://example.org/"));

  // After the initiator page navigates away, the started prerendering should be
  // cancelled, and PrerenderHostRegistry can start prerendering a new one.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders +
           1;
       i++) {
    EXPECT_EQ(registry().FindHostByUrlForTesting(prerendering_urls[i]),
              nullptr);
  }
  mojo::Remote<blink::mojom::SpeculationHost> remote2;
  SpeculationHostImpl::Bind(render_frame_host,
                            remote2.BindNewPipeAndPassReceiver());
  const GURL prerendering_url("https://example.org/next");
  SendCandidate(prerendering_url, remote2);
  EXPECT_NE(registry().FindHostByUrlForTesting(prerendering_url), nullptr);
  ExpectBucketCountOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded);
}

class PrerenderHostRegistryLimitGroupTest
    : public PrerenderHostRegistryTest,
      public testing::WithParamInterface<bool> {
 public:
  using PrerenderLimitGroup = PrerenderHostRegistry::PrerenderLimitGroup;

  void SetUp() override {
    feature_list_.InitWithFeatures({base::kStatefulMemoryPressure,
                                    blink::features::kPrerender2MemoryControls},
                                   {});
    PrerenderHostRegistryTest::SetUp();
  }

  const std::string histogram_suffix = "EmbedderSuffixForTest";

  bool IsNewTabTrigger(PrerenderLimitGroup limit_group) {
    return GetParam() && limit_group != PrerenderLimitGroup::kEmbedder;
  }

  PrerenderHostId CreateAndStartHostByLimitGroup(
      PrerenderLimitGroup limit_group) {
    static int unique_id = 0;
    const GURL prerendering_url("https://example.com/next_" +
                                base::NumberToString(unique_id));
    unique_id++;
    auto prerender_attributes = [&] {
      switch (limit_group) {
        case PrerenderLimitGroup::kSpeculationRulesImmediate:
          return GeneratePrerenderAttributes(
              prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
              blink::mojom::SpeculationEagerness::kImmediate,
              contents()->GetPrimaryMainFrame());
        case PrerenderLimitGroup::kSpeculationRulesNonImmediate:
          return GeneratePrerenderAttributes(
              prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
              blink::mojom::SpeculationEagerness::kModerate,
              contents()->GetPrimaryMainFrame());
        case PrerenderLimitGroup::kEmbedder:
          return GeneratePrerenderAttributes(
              prerendering_url, PreloadingTriggerType::kEmbedder,
              histogram_suffix, std::nullopt, nullptr);
      }
    }();

    PreloadingPredictor embedder_predictor(100, "Embedder");

    PreloadingPredictor creating_predictor = [&] {
      switch (limit_group) {
        case PrerenderLimitGroup::kSpeculationRulesImmediate:
        case PrerenderLimitGroup::kSpeculationRulesNonImmediate:
          return content_preloading_predictor::kSpeculationRules;
        case PrerenderLimitGroup::kEmbedder:
          return embedder_predictor;
      }
    }();
    PreloadingPredictor enacting_predictor = [&] {
      switch (limit_group) {
        case PrerenderLimitGroup::kSpeculationRulesImmediate:
          return content_preloading_predictor::kSpeculationRules;
        case PrerenderLimitGroup::kSpeculationRulesNonImmediate:
          // Arbitrarily chosen non-immediate predictor.
          return preloading_predictor::kUrlPointerDownOnAnchor;
        case PrerenderLimitGroup::kEmbedder:
          return embedder_predictor;
      }
    }();

    return IsNewTabTrigger(limit_group)
               ? registry().CreateAndStartHostForNewTab(
                     prerender_attributes, creating_predictor,
                     enacting_predictor, PreloadingConfidence{100})
               : registry().CreateAndStartHost(prerender_attributes);
  }

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

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

TEST_P(PrerenderHostRegistryLimitGroupTest, Immediate) {
  // Starts the immediate prerenders as many times as the specific limit.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders;
       i++) {
    PrerenderHostId prerender_host_id = CreateAndStartHostByLimitGroup(
        PrerenderLimitGroup::kSpeculationRulesImmediate);
    EXPECT_TRUE(prerender_host_id);
  }

  // If we try to start immediate prerenders after reaching the limit, that
  // should be canceled with kMaxNumOfRunningImmediatePrerendersExceeded.
  PrerenderHostId prerender_host_id_immediate_exceeded =
      CreateAndStartHostByLimitGroup(
          PrerenderLimitGroup::kSpeculationRulesImmediate);
  EXPECT_FALSE(prerender_host_id_immediate_exceeded);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded, 1);

  // On the other hand, prerenders belonging to different limit
  // group(non-immediate, embedder) can still be started.
  PrerenderHostId prerender_host_id_non_immediate =
      CreateAndStartHostByLimitGroup(
          PrerenderLimitGroup::kSpeculationRulesNonImmediate);
  PrerenderHostId prerender_host_id_embedder =
      CreateAndStartHostByLimitGroup(PrerenderLimitGroup::kEmbedder);
  EXPECT_TRUE(prerender_host_id_non_immediate);
  EXPECT_TRUE(prerender_host_id_embedder);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded, 1);
  ExpectUniqueSampleOfEmbedderFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningEmbedderPrerendersExceeded,
      histogram_suffix, 0);
}

TEST_P(PrerenderHostRegistryLimitGroupTest, NonImmediate) {
  std::vector<PrerenderHostId> started_prerender_ids;

  // Starts the non-immediate prerenders as many times as the specific limit.
  for (int i = 0;
       i <
       PrerenderHostRegistry::kMaxRunningSpeculationRulesNonImmediatePrerenders;
       i++) {
    PrerenderHostId prerender_host_id = CreateAndStartHostByLimitGroup(
        PrerenderLimitGroup::kSpeculationRulesNonImmediate);
    started_prerender_ids.push_back(prerender_host_id);
    EXPECT_TRUE(prerender_host_id);
  }

  // Even after the limit of non-immediate speculation rules is reached, it is
  // permissible to start a new prerender. Instead, the oldest prerender will be
  // canceled with kMaxNumOfRunningNonImmediatePrerendersExceeded to make room
  // for a new one.
  PrerenderHostId prerender_host_id_non_immediate_exceeded =
      CreateAndStartHostByLimitGroup(
          PrerenderLimitGroup::kSpeculationRulesNonImmediate);
  ASSERT_TRUE(prerender_host_id_non_immediate_exceeded);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningNonImmediatePrerendersExceeded, 1);

  for (auto id : started_prerender_ids) {
    FrameTreeNodeId prerender_frame_tree_node_id =
        PrerenderHost::GetFrameTreeNodeIdForId(id);
    auto* web_contents_impl = static_cast<WebContentsImpl*>(
        WebContents::FromFrameTreeNodeId(prerender_frame_tree_node_id));
    PrerenderHost* prerender_host = nullptr;
    if (web_contents_impl) {
      prerender_host = web_contents_impl->GetPrerenderHostRegistry()
                           ->FindNonReservedHostById(id);
    }
    if (id == started_prerender_ids[0]) {
      // The oldest prerender has been canceled.
      EXPECT_EQ(prerender_host, nullptr);
    } else {
      EXPECT_NE(prerender_host, nullptr);
    }
  }

  // On the other hand, prerenders belonging to different limit group(immediate,
  // embedder) can still be started and not invoke cancellation, as these limits
  // are separated.
  PrerenderHostId prerender_host_id_immediate = CreateAndStartHostByLimitGroup(
      PrerenderLimitGroup::kSpeculationRulesImmediate);
  PrerenderHostId prerender_host_id_embedder =
      CreateAndStartHostByLimitGroup(PrerenderLimitGroup::kEmbedder);
  EXPECT_TRUE(prerender_host_id_immediate);
  EXPECT_TRUE(prerender_host_id_embedder);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningNonImmediatePrerendersExceeded, 1);
  ExpectUniqueSampleOfEmbedderFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningEmbedderPrerendersExceeded,
      histogram_suffix, 0);
}

TEST_P(PrerenderHostRegistryLimitGroupTest, Embedder) {
  // Starts the embedder prerenders as many times as the specific limit.
  const int max_embedder_prerenders =
      web_contents()->GetDelegate()->AllowedPrerenderingCount(*web_contents());
  for (int i = 0; i < max_embedder_prerenders; i++) {
    PrerenderHostId prerender_host_id =
        CreateAndStartHostByLimitGroup(PrerenderLimitGroup::kEmbedder);
    EXPECT_TRUE(prerender_host_id);
  }

  // If we try to start embedder prerenders after reaching the limit, that
  // should be canceled with kMaxNumOfRunningEmbedderPrerendersExceeded.
  PrerenderHostId prerender_host_id_embedder_exceeded =
      CreateAndStartHostByLimitGroup(PrerenderLimitGroup::kEmbedder);
  EXPECT_FALSE(prerender_host_id_embedder_exceeded);
  ExpectUniqueSampleOfEmbedderFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningEmbedderPrerendersExceeded,
      histogram_suffix, 1);

  // On the other hand, prerenders belonging to different limit group(immediate,
  // non-egaer) can still be started.
  PrerenderHostId prerender_host_id_immediate = CreateAndStartHostByLimitGroup(
      PrerenderLimitGroup::kSpeculationRulesImmediate);
  PrerenderHostId prerender_host_id_non_immediate =
      CreateAndStartHostByLimitGroup(
          PrerenderLimitGroup::kSpeculationRulesNonImmediate);
  EXPECT_TRUE(prerender_host_id_immediate);
  EXPECT_TRUE(prerender_host_id_non_immediate);
  ExpectBucketCountOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningImmediatePrerendersExceeded, 0);
  ExpectBucketCountOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningNonImmediatePrerendersExceeded, 0);
  ExpectUniqueSampleOfEmbedderFinalStatus(
      PrerenderFinalStatus::kMaxNumOfRunningEmbedderPrerendersExceeded,
      histogram_suffix, 1);
}

TEST_P(PrerenderHostRegistryLimitGroupTest, StatefulNumberLimit) {
  // Set memory limit to 50%.
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(50, run_loop.QuitClosure());
    run_loop.Run();
  }

  const int default_limit =
      PrerenderHostRegistry::kMaxRunningSpeculationRulesImmediatePrerenders;
  const int expected_limit = default_limit / 2;

  std::vector<PrerenderHostId> prerender_host_ids;

  for (int i = 0; i < expected_limit + 1; i++) {
    PrerenderHostId prerender_host_id = CreateAndStartHostByLimitGroup(
        PrerenderLimitGroup::kSpeculationRulesImmediate);
    prerender_host_ids.push_back(prerender_host_id);
  }

  for (int i = 0; i < expected_limit; i++) {
    EXPECT_TRUE(prerender_host_ids[i]);
    if (IsNewTabTrigger(PrerenderLimitGroup::kSpeculationRulesImmediate)) {
      EXPECT_TRUE(
          registry().HasNewTabHandleByIdForTesting(prerender_host_ids[i]));
    } else {
      EXPECT_TRUE(registry().FindNonReservedHostById(prerender_host_ids[i]) !=
                  nullptr);
    }
  }
  EXPECT_FALSE(prerender_host_ids[expected_limit]);
  EXPECT_EQ(expected_limit,
            registry().GetHostCountByLimitGroupForTesting(
                PrerenderLimitGroup::kSpeculationRulesImmediate));
}

class PrerenderHostRegistryLegacyMemoryControlsTest
    : public PrerenderHostRegistryTest {
 public:
  void SetUp() override {
    feature_list_.InitWithFeatures({blink::features::kPrerender2MemoryControls},
                                   {base::kStatefulMemoryPressure});
    PrerenderHostRegistryTest::SetUp();
  }

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

TEST_F(PrerenderHostRegistryLegacyMemoryControlsTest,
       StatefulNumberLimit_Disabled_NonCritical) {
  const GURL prerendering_url("https://example.com/next");
  PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_TRUE(prerender_host_id);

  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(80, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  EXPECT_TRUE(registry().FindNonReservedHostById(prerender_host_id));
}

TEST_F(PrerenderHostRegistryLegacyMemoryControlsTest,
       StatefulNumberLimit_Disabled_Critical) {
  const GURL prerendering_url("https://example.com/next");
  PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_TRUE(prerender_host_id);

  // Under the legacy memory pressure system (when kStatefulMemoryPressure is
  // disabled), critical memory pressure notifications are dispatched with a
  // simulated memory limit of 0% (<= kCriticalMemoryPressureThreshold).
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(0, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  EXPECT_FALSE(registry().FindNonReservedHostById(prerender_host_id));

  // Verify that after critical memory pressure cancellation in legacy mode,
  // future prerenders are not prevented from being created.
  const GURL prerendering_url2("https://example.com/next2");
  PrerenderHostId new_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          prerendering_url2, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_TRUE(new_host_id);
}

TEST_P(PrerenderHostRegistryLimitGroupTest, StatefulNumberLimit_Critical) {
  const int initial_count = 3;
  std::vector<PrerenderHostId> prerender_host_ids;
  for (int i = 0; i < initial_count; i++) {
    PrerenderHostId prerender_host_id = CreateAndStartHostByLimitGroup(
        PrerenderLimitGroup::kSpeculationRulesImmediate);
    prerender_host_ids.push_back(prerender_host_id);
  }
  EXPECT_EQ(initial_count,
            registry().GetHostCountByLimitGroupForTesting(
                PrerenderLimitGroup::kSpeculationRulesImmediate));

  // Under stateful memory pressure (when kStatefulMemoryPressure is enabled),
  // a critical memory pressure notification (0%) actively evicts all running
  // hosts AND statefully updates the current memory limit to 0%, preventing
  // any new prerenders from starting.
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(0, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  EXPECT_EQ(0, registry().GetHostCountByLimitGroupForTesting(
                   PrerenderLimitGroup::kSpeculationRulesImmediate));

  // Verify that under stateful memory controls, future prerenders are indeed
  // prevented from starting while critical memory pressure persists.
  PrerenderHostId new_host_id = CreateAndStartHostByLimitGroup(
      PrerenderLimitGroup::kSpeculationRulesImmediate);
  EXPECT_EQ(PrerenderHostId(), new_host_id);
  EXPECT_EQ(0, registry().GetHostCountByLimitGroupForTesting(
                   PrerenderLimitGroup::kSpeculationRulesImmediate));

  // Reset memory limit back to 100% for subsequent tests.
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(100, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }
}

TEST_F(PrerenderHostRegistryTest, MemoryControlsDisabled_NoOp) {
  const GURL prerendering_url("https://example.com/next");
  PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_TRUE(prerender_host_id);

  // When kPrerender2MemoryControls is disabled, memory pressure notifications
  // must be completely ignored as a no-op, avoiding CancelAllHosts.
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(0, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  EXPECT_TRUE(registry().FindNonReservedHostById(prerender_host_id));
}

TEST_P(PrerenderHostRegistryLimitGroupTest, StatefulActiveEviction) {
  const int initial_count = 10;
  std::vector<PrerenderHostId> prerender_host_ids;
  for (int i = 0; i < initial_count; i++) {
    PrerenderHostId host_id = CreateAndStartHostByLimitGroup(
        PrerenderLimitGroup::kSpeculationRulesImmediate);
    ASSERT_TRUE(host_id);
    prerender_host_ids.push_back(host_id);
  }

  EXPECT_EQ(initial_count,
            registry().GetHostCountByLimitGroupForTesting(
                PrerenderLimitGroup::kSpeculationRulesImmediate));

  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(50, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  const int expected_remaining = initial_count / 2;
  EXPECT_EQ(expected_remaining,
            registry().GetHostCountByLimitGroupForTesting(
                PrerenderLimitGroup::kSpeculationRulesImmediate));

  // The oldest 5 hosts should remain running.
  for (int i = 0; i < expected_remaining; i++) {
    if (IsNewTabTrigger(PrerenderLimitGroup::kSpeculationRulesImmediate)) {
      EXPECT_TRUE(
          registry().HasNewTabHandleByIdForTesting(prerender_host_ids[i]));
    } else {
      EXPECT_TRUE(registry().FindNonReservedHostById(prerender_host_ids[i]));
    }
  }

  // The newest 5 hosts should be evicted.
  for (int i = expected_remaining; i < initial_count; i++) {
    if (IsNewTabTrigger(PrerenderLimitGroup::kSpeculationRulesImmediate)) {
      EXPECT_FALSE(
          registry().HasNewTabHandleByIdForTesting(prerender_host_ids[i]));
    } else {
      EXPECT_FALSE(registry().FindNonReservedHostById(prerender_host_ids[i]));
    }
  }
}

TEST_P(PrerenderHostRegistryLimitGroupTest,
       StatefulActiveEviction_MixedTriggers) {
  if (GetParam()) {
    // Pending queue logic is tested for standard in-tab prerendering.
    return;
  }

  const int initial_count = 10;
  std::vector<PrerenderHostId> prerender_host_ids;
  for (int i = 0; i < initial_count; i++) {
    const GURL prerendering_url("https://example.com/next_mixed_" +
                                base::NumberToString(i));
    auto attributes = GeneratePrerenderAttributes(
        prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
        blink::mojom::SpeculationEagerness::kImmediate,
        contents()->GetPrimaryMainFrame());

    PrerenderHostId host_id;
    if (i % 2 == 0) {
      host_id = registry().CreateAndStartHost(attributes);
    } else {
      host_id = registry().CreateAndStartHostForNewTab(
          attributes, content_preloading_predictor::kSpeculationRules,
          content_preloading_predictor::kSpeculationRules,
          PreloadingConfidence{100});
    }
    ASSERT_TRUE(host_id);
    prerender_host_ids.push_back(host_id);
  }

  EXPECT_EQ(initial_count, registry().GetHostCountByLimitGroupForTesting(
                               PrerenderHostRegistry::PrerenderLimitGroup::
                                   kSpeculationRulesImmediate));

  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(50, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  const int expected_remaining = initial_count / 2;
  EXPECT_EQ(expected_remaining, registry().GetHostCountByLimitGroupForTesting(
                                    PrerenderHostRegistry::PrerenderLimitGroup::
                                        kSpeculationRulesImmediate));

  // The oldest 5 hosts (indices 0 to 4) should remain running.
  for (int i = 0; i < expected_remaining; i++) {
    if (i % 2 == 0) {
      EXPECT_TRUE(registry().FindNonReservedHostById(prerender_host_ids[i]));
    } else {
      EXPECT_TRUE(
          registry().HasNewTabHandleByIdForTesting(prerender_host_ids[i]));
    }
  }

  // The newest 5 hosts (indices 5 to 9) should be evicted.
  for (int i = expected_remaining; i < initial_count; i++) {
    if (i % 2 == 0) {
      EXPECT_FALSE(registry().FindNonReservedHostById(prerender_host_ids[i]));
    } else {
      EXPECT_FALSE(
          registry().HasNewTabHandleByIdForTesting(prerender_host_ids[i]));
    }
  }
}

TEST_P(PrerenderHostRegistryLimitGroupTest,
       StatefulActiveEviction_NonImmediate) {
  PrerenderHostId host1 = CreateAndStartHostByLimitGroup(
      PrerenderLimitGroup::kSpeculationRulesNonImmediate);
  ASSERT_TRUE(host1);

  PrerenderHostId host2 = CreateAndStartHostByLimitGroup(
      PrerenderLimitGroup::kSpeculationRulesNonImmediate);
  ASSERT_TRUE(host2);

  EXPECT_EQ(2, registry().GetHostCountByLimitGroupForTesting(
                   PrerenderLimitGroup::kSpeculationRulesNonImmediate));

  // Set memory limit to 50% (scaled limit becomes 2 * 50% = 1).
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(50, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  // Under FIFO eviction for non-immediate prerenders, host1 (oldest) should be
  // evicted and host2 (newest) should remain running.
  EXPECT_EQ(1, registry().GetHostCountByLimitGroupForTesting(
                   PrerenderLimitGroup::kSpeculationRulesNonImmediate));
  if (IsNewTabTrigger(PrerenderLimitGroup::kSpeculationRulesNonImmediate)) {
    EXPECT_FALSE(registry().HasNewTabHandleByIdForTesting(host1));
    EXPECT_TRUE(registry().HasNewTabHandleByIdForTesting(host2));
  } else {
    EXPECT_FALSE(registry().FindNonReservedHostById(host1));
    EXPECT_TRUE(registry().FindNonReservedHostById(host2));
  }
}

struct StatefulPerGroupScalingParams {
  bool scale_immediate;
  bool scale_non_immediate;
  bool scale_embedder;
};

class PrerenderHostRegistryStatefulPerGroupTest
    : public PrerenderHostRegistryTest,
      public ::testing::WithParamInterface<StatefulPerGroupScalingParams> {
 public:
  void SetUp() override {
    const auto& params = GetParam();
    feature_list_.InitWithFeaturesAndParameters(
        {{base::kStatefulMemoryPressure,
          {{"PrerenderScaleImmediate",
            params.scale_immediate ? "true" : "false"},
           {"PrerenderScaleNonImmediate",
            params.scale_non_immediate ? "true" : "false"},
           {"PrerenderScaleEmbedder",
            params.scale_embedder ? "true" : "false"}}},
         {blink::features::kPrerender2MemoryControls, {}}},
        {});
    PrerenderHostRegistryTest::SetUp();
  }

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

TEST_P(PrerenderHostRegistryStatefulPerGroupTest, ScalingBehavior_AllGroups) {
  const auto& params = GetParam();

  // Set memory limit to 50%.
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(50, run_loop.QuitClosure());
    run_loop.Run();
  }

  // Immediate Group
  const int expected_immediate =
      params.scale_immediate
          ? PrerenderHostRegistry::
                    kMaxRunningSpeculationRulesImmediatePrerenders /
                2
          : PrerenderHostRegistry::
                kMaxRunningSpeculationRulesImmediatePrerenders;

  // Non-Immediate Group
  const int expected_non_immediate =
      params.scale_non_immediate
          ? PrerenderHostRegistry::
                    kMaxRunningSpeculationRulesNonImmediatePrerenders /
                2
          : PrerenderHostRegistry::
                kMaxRunningSpeculationRulesNonImmediatePrerenders;

  // Embedder Group
  const int max_embedder =
      web_contents()->GetDelegate()->AllowedPrerenderingCount(*web_contents());
  const int expected_embedder =
      params.scale_embedder ? max_embedder / 2 : max_embedder;

  EXPECT_EQ(expected_immediate, registry().GetScaledLimitForTesting(
                                    PrerenderHostRegistry::PrerenderLimitGroup::
                                        kSpeculationRulesImmediate));

  EXPECT_EQ(expected_non_immediate,
            registry().GetScaledLimitForTesting(
                PrerenderHostRegistry::PrerenderLimitGroup::
                    kSpeculationRulesNonImmediate));

  EXPECT_EQ(expected_embedder,
            registry().GetScaledLimitForTesting(
                PrerenderHostRegistry::PrerenderLimitGroup::kEmbedder));
}

TEST_P(PrerenderHostRegistryStatefulPerGroupTest,
       CriticalMemoryPressure_SafetyNet) {
  // Start 1 host of immediate group
  PrerenderHostId imm_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          GURL("https://example.com/imm"),
          PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(imm_id);

  // Trigger critical memory pressure (0%)
  {
    base::RunLoop run_loop;
    test_registry_.NotifyUpdateMemoryLimitAsync(0, base::DoNothing());
    test_registry_.NotifyReleaseMemoryAsync(run_loop.QuitClosure());
    run_loop.Run();
  }

  // Verify all hosts are cancelled under critical memory pressure safety net
  EXPECT_EQ(0, registry().GetHostCountByLimitGroupForTesting(
                   PrerenderHostRegistry::PrerenderLimitGroup::
                       kSpeculationRulesImmediate));
  EXPECT_EQ(0, registry().GetHostCountByLimitGroupForTesting(
                   PrerenderHostRegistry::PrerenderLimitGroup::
                       kSpeculationRulesNonImmediate));
}

INSTANTIATE_TEST_SUITE_P(
    AllGroupScalingCombinations,
    PrerenderHostRegistryStatefulPerGroupTest,
    ::testing::Values(StatefulPerGroupScalingParams{false, false, false},
                      StatefulPerGroupScalingParams{true, false, false},
                      StatefulPerGroupScalingParams{false, true, false},
                      StatefulPerGroupScalingParams{false, false, true},
                      StatefulPerGroupScalingParams{true, true, true}));

TEST_F(PrerenderHostRegistryTest,
       ReserveHostToActivateBeforeReadyForActivation) {
  const GURL original_url = contents()->GetLastCommittedURL();
  const GURL kPrerenderingUrl("https://example.com/next");

  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindHostByUrlForTesting(kPrerenderingUrl);
  FrameTreeNode* ftn =
      FrameTreeNode::From(prerender_host->GetPrerenderedMainFrameHost());
  std::unique_ptr<NavigationSimulatorImpl> sim =
      NavigationSimulatorImpl::CreateFromPendingInFrame(ftn);
  // Ensure that navigation in prerendering frame tree does not commit and
  // PrerenderHost doesn't become ready for activation.
  sim->SetAutoAdvance(false);

  EXPECT_FALSE(prerender_host->is_ready_for_activation());

  test::PrerenderHostObserver prerender_host_observer(*contents(),
                                                      kPrerenderingUrl);

  // Start activation.
  std::unique_ptr<NavigationSimulatorImpl> navigation =
      CreateActivation(kPrerenderingUrl, *contents());
  navigation->Start();

  // Wait until PrerenderCommitDeferringCondition runs.
  // TODO(nhiroki): Avoid using base::RunUntilIdle() and instead use some
  // explicit signal of the running condition.
  base::RunLoop().RunUntilIdle();

  // The activation should be deferred by PrerenderCommitDeferringCondition
  // until the main frame navigation in the prerendering frame tree finishes.
  NavigationRequest* navigation_request = navigation->GetNavigationHandle();
  EXPECT_TRUE(
      navigation_request->IsCommitDeferringConditionDeferredForTesting());
  EXPECT_FALSE(prerender_host_observer.was_activated());
  EXPECT_EQ(contents()->GetPrimaryMainFrame()->GetLastCommittedURL(),
            original_url);

  // Finish the main frame navigation.
  sim->Commit();

  // Finish the activation.
  prerender_host_observer.WaitForDestroyed();
  EXPECT_TRUE(prerender_host_observer.was_activated());
  EXPECT_EQ(registry().FindHostByUrlForTesting(kPrerenderingUrl), nullptr);
  EXPECT_EQ(contents()->GetPrimaryMainFrame()->GetLastCommittedURL(),
            kPrerenderingUrl);
}

TEST_F(PrerenderHostRegistryTest, CancelHost) {
  const GURL kPrerenderingUrl("https://example.com/next");
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  EXPECT_NE(registry().FindHostByUrlForTesting(kPrerenderingUrl), nullptr);

  registry().CancelHost(prerender_host_id, PrerenderFinalStatus::kDestroyed);
  EXPECT_EQ(registry().FindHostByUrlForTesting(kPrerenderingUrl), nullptr);
}

// Test cancelling a prerender while a CommitDeferringCondition is running.
// This activation should fall back to a regular navigation.
TEST_F(PrerenderHostRegistryTest,
       CancelHostWhileCommitDeferringConditionIsRunning) {
  const GURL original_url = contents()->GetLastCommittedURL();

  // Start prerendering.
  const GURL kPrerenderingUrl("https://example.com/next");
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindHostByUrlForTesting(kPrerenderingUrl);
  CommitPrerenderNavigation(*prerender_host);

  test::PrerenderHostObserver prerender_host_observer(*contents(),
                                                      kPrerenderingUrl);

  // Now navigate the primary page to the prerendered URL so that we activate
  // the prerender. Use a CommitDeferringCondition to pause activation
  // before it completes.
  std::unique_ptr<NavigationSimulatorImpl> navigation;

  {
    MockCommitDeferringConditionInstaller installer(
        kPrerenderingUrl, CommitDeferringCondition::Result::kDefer);
    // Start trying to activate the prerendered page.
    navigation = CreateActivation(kPrerenderingUrl, *contents());
    navigation->Start();

    // Wait for the condition to pause the activation.
    installer.WaitUntilInstalled();
    installer.condition().WaitUntilInvoked();

    // The request should be deferred by the condition.
    auto* navigation_request =
        static_cast<NavigationRequest*>(navigation->GetNavigationHandle());
    EXPECT_TRUE(
        navigation_request->IsCommitDeferringConditionDeferredForTesting());

    // The primary page should still be the original page.
    EXPECT_EQ(contents()->GetLastCommittedURL(), original_url);

    // Cancel the prerender while the CommitDeferringCondition is running.
    registry().CancelHost(prerender_host_id, PrerenderFinalStatus::kDestroyed);
    prerender_host_observer.WaitForDestroyed();
    EXPECT_FALSE(prerender_host_observer.was_activated());
    EXPECT_EQ(registry().FindHostByUrlForTesting(kPrerenderingUrl), nullptr);

    // Resume the activation. This should fall back to a regular navigation.
    installer.condition().CallResumeClosure();
  }

  navigation->Commit();
  EXPECT_EQ(contents()->GetPrimaryMainFrame()->GetLastCommittedURL(),
            kPrerenderingUrl);
}

// Test cancelling a prerender and then starting a new prerender for the same
// URL while a CommitDeferringCondition is running. This activation should not
// reserve the second prerender and should fall back to a regular navigation.
TEST_F(PrerenderHostRegistryTest,
       CancelAndStartHostWhileCommitDeferringConditionIsRunning) {
  const GURL original_url = contents()->GetLastCommittedURL();
  const GURL kPrerenderingUrl("https://example.com/next");

  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindHostByUrlForTesting(kPrerenderingUrl);
  CommitPrerenderNavigation(*prerender_host);

  test::PrerenderHostObserver prerender_host_observer(*contents(),
                                                      kPrerenderingUrl);

  // Now navigate the primary page to the prerendered URL so that we activate
  // the prerender. Use a CommitDeferringCondition to pause activation
  // before it completes.
  std::unique_ptr<NavigationSimulatorImpl> navigation;
  base::OnceClosure resume_navigation;

  {
    MockCommitDeferringConditionInstaller installer(
        kPrerenderingUrl, CommitDeferringCondition::Result::kDefer);
    // Start trying to activate the prerendered page.
    navigation = CreateActivation(kPrerenderingUrl, *contents());
    navigation->Start();

    // Wait for the condition to pause the activation.
    installer.WaitUntilInstalled();
    installer.condition().WaitUntilInvoked();
    resume_navigation = installer.condition().TakeResumeClosure();

    // The request should be deferred by the condition.
    auto* navigation_request =
        static_cast<NavigationRequest*>(navigation->GetNavigationHandle());
    EXPECT_TRUE(
        navigation_request->IsCommitDeferringConditionDeferredForTesting());

    // The primary page should still be the original page.
    EXPECT_EQ(contents()->GetLastCommittedURL(), original_url);

    // Cancel the prerender while the CommitDeferringCondition is running.
    registry().CancelHost(prerender_host_id, PrerenderFinalStatus::kDestroyed);
    prerender_host_observer.WaitForDestroyed();
    EXPECT_FALSE(prerender_host_observer.was_activated());
    EXPECT_EQ(registry().FindHostByUrlForTesting(kPrerenderingUrl), nullptr);
  }

  {
    // Start the second prerender for the same URL.
    const PrerenderHostId prerender_host_id2 =
        registry().CreateAndStartHost(GeneratePrerenderAttributes(
            kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
            blink::mojom::SpeculationEagerness::kImmediate,
            contents()->GetPrimaryMainFrame()));
    ASSERT_TRUE(prerender_host_id2);
    PrerenderHost* prerender_host2 =
        registry().FindHostByUrlForTesting(kPrerenderingUrl);
    CommitPrerenderNavigation(*prerender_host2);

    EXPECT_NE(prerender_host_id, prerender_host_id2);
  }

  // Resume the initial activation. This should not reserve the second
  // prerender and should fall back to a regular navigation.
  std::move(resume_navigation).Run();
  navigation->Commit();
  EXPECT_EQ(contents()->GetPrimaryMainFrame()->GetLastCommittedURL(),
            kPrerenderingUrl);

  // The second prerender should still exist.
  EXPECT_NE(registry().FindHostByUrlForTesting(kPrerenderingUrl), nullptr);
}

// Tests that prerendering should be canceled if the trigger is in the
// background and its type is kEmbedder.
// For the case where the trigger type is speculation rules,
// browsertests `TestSequentialPrerenderingInBackground` covers it.
TEST_F(PrerenderHostRegistryTest,
       DontStartPrerenderWhenEmbedderTriggerIsAlreadyHidden) {
  // The visibility state to be HIDDEN will cause prerendering not started when
  // trigger type is kEmbedder.
  contents()->WasHidden();

  const GURL kPrerenderingUrl = GURL("https://example.com/empty.html");
  RenderFrameHostImpl* initiator_rfh = contents()->GetPrimaryMainFrame();
  PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kEmbedder, "DirectURLInput",
          std::nullopt, initiator_rfh));
  EXPECT_FALSE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindNonReservedHostById(prerender_host_id);
  EXPECT_EQ(prerender_host, nullptr);
  histogram_tester().ExpectUniqueSample(
      "Prerender.Experimental.PrerenderHostFinalStatus.Embedder_DirectURLInput",
      PrerenderFinalStatus::kTriggerBackgrounded, 1u);
}

// -------------------------------------------------
// Activation navigation parameter matching unit tests.
// These tests change a parameter to differentiate the activation request from
// the prerendering request.

// A positive test to show that if the navigation params are equal then the
// prerender host is selected for activation.
TEST_F(PrerenderHostRegistryTest, SameInitialAndActivationParams) {
  EXPECT_TRUE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        // Do not change any params, so activation happens.
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kOk);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_InitiatorFrameToken) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        const GURL kOriginalUrl("https://example.com/");
        navigation->SetInitiatorFrame(nullptr);
        navigation->set_initiator_origin(url::Origin::Create(kOriginalUrl));
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kInitiatorFrameToken);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_Headers) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_request_headers("User-Agent: Test");
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kHttpRequestHeader);
}

// Tests that the Purpose header is ignored when comparing request headers.
TEST_F(PrerenderHostRegistryTest, PurposeHeaderIsIgnoredForParamMatching) {
  EXPECT_TRUE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_request_headers("Sec-Purpose: Test");
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kOk);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_LoadFlags) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_load_flags(net::LOAD_ONLY_FROM_CACHE);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kLoadFlags);

  // If the potential activation request requires validation or bypass of the
  // browser cache, the prerendered page should not be activated.
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_load_flags(net::LOAD_VALIDATE_CACHE);
      })));
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_load_flags(net::LOAD_BYPASS_CACHE);
      })));
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_load_flags(net::LOAD_DISABLE_CACHE);
      })));
  ExpectBucketCountOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kCacheLoadFlags, 3);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_SkipServiceWorker) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_skip_service_worker(true);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kSkipServiceWorker);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_MixedContentContextType) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_mixed_content_context_type(
            blink::mojom::MixedContentContextType::kNotMixedContent);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kMixedContentContextType);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_IsFormSubmission) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->SetIsFormSubmission(true);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kIsFormSubmission);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_SearchableFormUrl) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        const GURL kOriginalUrl("https://example.com/");
        navigation->set_searchable_form_url(kOriginalUrl);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kSearchableFormUrl);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationBeginParams_SearchableFormEncoding) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_searchable_form_encoding("Test encoding");
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kSearchableFormEncoding);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_InitiatorOrigin) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_initiator_origin(url::Origin());
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kInitiatorOrigin);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_ShouldNotCheckMainWorldCSP) {
  // Initial navigation blocked by the main world CSP cancels prerendering.
  // So, it's safe to match the page for CSP bypassing requests from isolated
  // worlds (e.g., extensions).
  EXPECT_TRUE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_should_check_main_world_csp(
            network::mojom::CSPDisposition::DO_NOT_CHECK);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kOk);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_Method) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->SetMethod("POST");
      })));
  // The method parameter change is detected as a HTTP request header change.
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kHttpRequestHeader);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_HrefTranslate) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_href_translate("test");
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kHrefTranslate);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_Transition) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->SetTransition(ui::PAGE_TRANSITION_FORM_SUBMIT);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kTransition);

  histogram_tester().ExpectUniqueSample(
      "Prerender.Experimental.ActivationTransitionMismatch.SpeculationRule",
      ui::PAGE_TRANSITION_FORM_SUBMIT, 1);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_RequestContextType) {
  EXPECT_FALSE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_request_context_type(
            blink::mojom::RequestContextType::AUDIO);
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kRequestContextType);
}

TEST_F(PrerenderHostRegistryTest,
       CompareInitialAndActivationCommonParams_ReferrerPolicy) {
  EXPECT_TRUE(CheckIsActivatedForParams(
      base::BindLambdaForTesting([&](NavigationSimulatorImpl* navigation) {
        navigation->SetReferrer(blink::mojom::Referrer::New(
            contents()->GetPrimaryMainFrame()->GetLastCommittedURL(),
            network::mojom::ReferrerPolicy::kAlways));
      })));
  ExpectUniqueSampleOfActivationNavigationParamsMatch(
      PrerenderHost::ActivationNavigationParamsMatch::kOk);
}

// End navigation parameter matching tests ---------

// Begin replication state matching tests ----------

TEST_F(PrerenderHostRegistryTest, InsecureRequestPolicyIsSetWhilePrerendering) {
  SetupPrerenderAndCommit(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        navigation->set_insecure_request_policy(
            blink::mojom::InsecureRequestPolicy::kBlockAllMixedContent);
      }));
  EXPECT_EQ(contents()
                ->GetPrimaryMainFrame()
                ->frame_tree_node()
                ->current_replication_state()
                .insecure_request_policy,
            blink::mojom::InsecureRequestPolicy::kBlockAllMixedContent);
}

TEST_F(PrerenderHostRegistryTest,
       InsecureNavigationsSetIsSetWhilePrerendering) {
  SetupPrerenderAndCommit(
      base::BindLambdaForTesting([](NavigationSimulatorImpl* navigation) {
        const std::vector<uint32_t> insecure_navigations = {1, 2};
        navigation->set_insecure_navigations_set(insecure_navigations);
      }));
  const std::vector<uint32_t> insecure_navigations = {1, 2};
  EXPECT_EQ(contents()
                ->GetPrimaryMainFrame()
                ->frame_tree_node()
                ->current_replication_state()
                .insecure_navigations_set,
            insecure_navigations);
}

// End replication state matching tests ------------

TEST_F(PrerenderHostRegistryTest, OneTaskToDeleteAllHosts) {
  std::vector<PrerenderHostId> prerender_host_ids;
  std::vector<std::unique_ptr<test::PrerenderHostObserver>>
      prerender_host_observers;

  for (int i = 0; i < 2; i++) {
    const GURL prerendering_url("https://example.com/next" +
                                base::NumberToString(i));
    PrerenderHostId prerender_host_id =
        registry().CreateAndStartHost(GeneratePrerenderAttributes(
            prerendering_url, PreloadingTriggerType::kSpeculationRule, "",
            blink::mojom::SpeculationEagerness::kImmediate,
            contents()->GetPrimaryMainFrame()));

    prerender_host_observers.emplace_back(
        std::make_unique<test::PrerenderHostObserver>(*contents(),
                                                      prerender_host_id));
    prerender_host_ids.push_back(prerender_host_id);
  }
  int pending_task_before_posting_abandon_task =
      task_environment()->GetPendingMainThreadTaskCount();
  registry().CancelHosts(
      prerender_host_ids,
      PrerenderCancellationReason(PrerenderFinalStatus::kDestroyed));
  int pending_task_after_posting_abandon_task =
      task_environment()->GetPendingMainThreadTaskCount();
  // Only one task was posted.
  EXPECT_EQ(pending_task_before_posting_abandon_task + 1,
            pending_task_after_posting_abandon_task);
  for (auto& observer : prerender_host_observers) {
    // All PrerenderHosts were deleted, so it should not timeout.
    observer->WaitForDestroyed();
  }
}

TEST_F(PrerenderHostRegistryTest, DisallowPageHavingEffectiveUrl_TriggerUrl) {
  const GURL original_url = contents()->GetLastCommittedURL();
  const GURL kModifiedSiteUrl("custom-scheme://custom");

  // Let the trigger's URL have the effective URL.
  EffectiveURLContentBrowserClient modified_client(
      original_url, kModifiedSiteUrl,
      /*requires_dedicated_process=*/false);
  ContentBrowserClient* old_client =
      SetBrowserClientForTesting(&modified_client);

  // Start prerendering. This should fail as the initiator's URL has the
  // effective URL.
  const GURL kPrerenderingUrl("https://example.com/empty.html");
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_FALSE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindNonReservedHostById(prerender_host_id);
  EXPECT_EQ(prerender_host, nullptr);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kTriggerUrlHasEffectiveUrl);

  SetBrowserClientForTesting(old_client);
}

TEST_F(PrerenderHostRegistryTest,
       DisallowPageHavingEffectiveUrl_PrerenderingUrl) {
  const GURL original_url = contents()->GetLastCommittedURL();
  const GURL kPrerenderingUrl("https://example.com/empty.html");
  const GURL kModifiedSiteUrl("custom-scheme://custom");

  // Let the prerendering URL have the effective URL.
  EffectiveURLContentBrowserClient modified_client(
      kPrerenderingUrl, kModifiedSiteUrl,
      /*requires_dedicated_process=*/false);
  ContentBrowserClient* old_client =
      SetBrowserClientForTesting(&modified_client);

  // Start prerendering. This should fail as the prerendering URL has the
  // effective URL.
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  EXPECT_FALSE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindNonReservedHostById(prerender_host_id);
  EXPECT_EQ(prerender_host, nullptr);
  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kPrerenderingUrlHasEffectiveUrl);

  SetBrowserClientForTesting(old_client);
}

TEST_F(PrerenderHostRegistryTest,
       DisallowPageHavingEffectiveUrl_ActivationUrl) {
  const GURL original_url = contents()->GetLastCommittedURL();
  const GURL kPrerenderingUrl("https://example.com/empty.html");
  const GURL kModifiedSiteUrl("custom-scheme://custom");

  // Start prerendering.
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindHostByUrlForTesting(kPrerenderingUrl);
  CommitPrerenderNavigation(*prerender_host);

  // Let the prerendering URL have the effective URL after prerendering.
  EffectiveURLContentBrowserClient modified_client(
      kPrerenderingUrl, kModifiedSiteUrl,
      /*requires_dedicated_process=*/false);
  ContentBrowserClient* old_client =
      SetBrowserClientForTesting(&modified_client);

  // Navigate the primary page to the prerendering URL that has the effective
  // URL. This should fail to activate the prerendered page.
  contents()->NavigateAndCommit(kPrerenderingUrl);

  ExpectUniqueSampleOfSpeculationRuleFinalStatus(
      PrerenderFinalStatus::kActivationUrlHasEffectiveUrl);

  SetBrowserClientForTesting(old_client);
}

TEST_F(PrerenderHostRegistryTest, PotentialPrerenderProcessReuseUMA) {
  const GURL kPrerenderingUrl("https://example.com/next");
  // Start prerendering.
  const PrerenderHostId prerender_host_id =
      registry().CreateAndStartHost(GeneratePrerenderAttributes(
          kPrerenderingUrl, PreloadingTriggerType::kSpeculationRule, "",
          blink::mojom::SpeculationEagerness::kImmediate,
          contents()->GetPrimaryMainFrame()));
  ASSERT_TRUE(prerender_host_id);
  PrerenderHost* prerender_host =
      registry().FindHostByUrlForTesting(kPrerenderingUrl);
  CommitPrerenderNavigation(*prerender_host);

  {
    base::HistogramTester histogram_tester;
    auto navigation = CreateActivation(kPrerenderingUrl, *contents());
    navigation->Start();
    histogram_tester.ExpectUniqueSample(
        "Prerender.Experimental.PrerenderProcessReuseAvailability",
        PrerenderHostRegistry::PrerenderProcessReuseAvailability::
            kHasMatchableHosts,
        1);
  }
  {
    base::HistogramTester histogram_tester;
    const GURL kSameSiteUrl("https://example.com/other");
    auto navigation = CreateActivation(kSameSiteUrl, *contents());
    navigation->Start();
    histogram_tester.ExpectUniqueSample(
        "Prerender.Experimental.PrerenderProcessReuseAvailability",
        PrerenderHostRegistry::PrerenderProcessReuseAvailability::
            kHasSameOriginHosts,
        1);
  }
  {
    base::HistogramTester histogram_tester;
    const GURL kSameSiteUrl("https://www.example.com:8000/other");
    auto navigation = CreateActivation(kSameSiteUrl, *contents());
    navigation->Start();
    histogram_tester.ExpectUniqueSample(
        "Prerender.Experimental.PrerenderProcessReuseAvailability",
        PrerenderHostRegistry::PrerenderProcessReuseAvailability::
            kHasSameSiteHosts,
        1);
  }
  {
    base::HistogramTester histogram_tester;
    const GURL kNotSameSiteUrl("https://another.com/other");
    auto navigation = CreateActivation(kNotSameSiteUrl, *contents());
    navigation->Start();
    histogram_tester.ExpectUniqueSample(
        "Prerender.Experimental.PrerenderProcessReuseAvailability",
        PrerenderHostRegistry::PrerenderProcessReuseAvailability::
            kNoSameOriginOrSiteHosts,
        1);
  }
}

}  // namespace
}  // namespace content