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

#ifndef CONTENT_BROWSER_PRELOADING_PREFETCH_PREFETCH_TEST_UTIL_INTERNAL_H_
#define CONTENT_BROWSER_PRELOADING_PREFETCH_PREFETCH_TEST_UTIL_INTERNAL_H_

#include <string>

#include "base/test/metrics/histogram_tester.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/test_future.h"
#include "content/browser/preloading/prefetch/prefetch_service.h"
#include "content/public/test/preloading_test_util.h"
#include "content/public/test/test_renderer_host.h"
#include "content/test/test_content_browser_client.h"
#include "mojo/public/cpp/system/data_pipe.h"
#include "mojo/public/cpp/system/data_pipe_drainer.h"
#include "services/metrics/public/cpp/ukm_source_id.h"
#include "services/network/public/mojom/early_hints.mojom.h"
#include "services/network/public/mojom/url_loader.mojom.h"
#include "services/network/test/test_url_loader_factory.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"

namespace base {
class RunLoop;
}  // namespace base

namespace content {

using OnPrefetchCompleteTestFuture =
    base::test::TestFuture<network::URLLoaderCompletionStatus>;
using OnPrefetchReceiveRedirectTestFuture =
    base::test::TestFuture<net::RedirectInfo,
                           network::mojom::URLResponseHeadPtr>;

struct NotReachedTagForTests {};

// Used to specify the test behavior on a specific callback.
// - If `NotReachedTagForTests()`: the callback shouldn't be called.
// - Otherwise: `T` (either `RunLoop*` or `TestFuture*`) is unblocked when the
//   callback is called, if non-nullptr.
template <typename T>
using NotReachedTagForTestsOr = std::variant<T, NotReachedTagForTests>;

// The centralized helper to create `PrefetchStreamingURLLoader` and its
// corresponding `PrefetchResponseReader`, without `PrefetchContainer`.
std::tuple<scoped_refptr<PrefetchResponseReader>,
           base::WeakPtr<PrefetchStreamingURLLoader>>
CreateStreamingURLLoaderWithoutPrefetchContainerForTests(
    scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
    const network::ResourceRequest& prefetch_request,
    NotReachedTagForTestsOr<base::RunLoop*> on_response_received,
    NotReachedTagForTestsOr<OnPrefetchCompleteTestFuture*> on_complete,
    NotReachedTagForTestsOr<OnPrefetchReceiveRedirectTestFuture*>
        on_receive_redirect,
    NotReachedTagForTestsOr<base::RunLoop*> on_head_received,
    std::optional<PrefetchErrorOnResponseReceived> error_on_response_received =
        std::nullopt,
    base::TimeDelta timeout_duration = {});

// The centralized helper to create `PrefetchStreamingURLLoader` associated with
// - `PrefetchContainer` (possibly nullptr) and
// - `PrefetchResponseReader` (always non-nullptr).
base::WeakPtr<PrefetchStreamingURLLoader> CreateStreamingURLLoaderForTests(
    base::WeakPtr<PrefetchContainer> prefetch_container,
    base::WeakPtr<PrefetchResponseReader> response_reader,
    scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
    const network::ResourceRequest& prefetch_request,
    NotReachedTagForTestsOr<base::RunLoop*> on_response_received,
    NotReachedTagForTestsOr<OnPrefetchReceiveRedirectTestFuture*>
        on_receive_redirect,
    std::optional<PrefetchErrorOnResponseReceived> error_on_response_received =
        std::nullopt,
    base::TimeDelta timeout_duration = {});

void MakeServableStreamingURLLoaderForTest(
    PrefetchContainer* prefetch_container,
    network::mojom::URLResponseHeadPtr head,
    const std::string body,
    network::URLLoaderCompletionStatus status =
        network::URLLoaderCompletionStatus(net::OK));

network::TestURLLoaderFactory::PendingRequest
MakeManuallyServableStreamingURLLoaderForTest(
    PrefetchContainer* prefetch_container);

void MakeServableStreamingURLLoaderWithRedirectForTest(
    PrefetchContainer* prefetch_container,
    const GURL& original_url,
    const GURL& redirect_url);

void MakeServableStreamingURLLoadersWithNetworkTransitionRedirectForTest(
    PrefetchContainer* prefetch_container,
    const GURL& original_url,
    const GURL& redirect_url);

class PrefetchTestURLLoaderClient : public network::mojom::URLLoaderClient,
                                    public mojo::DataPipeDrainer::Client {
 public:
  PrefetchTestURLLoaderClient();
  ~PrefetchTestURLLoaderClient() override;

  PrefetchTestURLLoaderClient(const PrefetchTestURLLoaderClient&) = delete;
  PrefetchTestURLLoaderClient& operator=(const PrefetchTestURLLoaderClient&) =
      delete;

  mojo::PendingReceiver<network::mojom::URLLoader>
  BindURLloaderAndGetReceiver();
  mojo::PendingRemote<network::mojom::URLLoaderClient>
  BindURLLoaderClientAndGetRemote();
  void DisconnectMojoPipes();

  // By default, auto draining is enabled, i.e. body data pipe is started
  // draining when received. If auto draining is disabled by
  // `SetAutoDraining(false)`, `StartDraining()` should be explicitly called
  // only once.
  void SetAutoDraining(bool auto_draining) { auto_draining_ = auto_draining; }
  void StartDraining();

  std::string body_content() { return body_content_; }
  uint32_t total_bytes_read() { return total_bytes_read_; }
  bool body_finished() { return body_finished_; }
  int32_t total_transfer_size_diff() { return total_transfer_size_diff_; }
  std::optional<network::URLLoaderCompletionStatus> completion_status() {
    return completion_status_;
  }
  const std::vector<
      std::pair<net::RedirectInfo, network::mojom::URLResponseHeadPtr>>&
  received_redirects() {
    return received_redirects_;
  }

 private:
  // network::mojom::URLLoaderClient
  void OnReceiveEarlyHints(network::mojom::EarlyHintsPtr early_hints) override;
  void OnReceiveResponse(
      network::mojom::URLResponseHeadPtr head,
      mojo::ScopedDataPipeConsumerHandle body,
      std::optional<mojo_base::BigBuffer> cached_metadata) override;
  void OnReceiveRedirect(const net::RedirectInfo& redirect_info,
                         network::mojom::URLResponseHeadPtr head) override;
  void OnUploadProgress(int64_t current_position,
                        int64_t total_size,
                        OnUploadProgressCallback callback) override;
  void OnTransferSizeUpdated(int32_t transfer_size_diff) override;
  void OnComplete(const network::URLLoaderCompletionStatus& status) override;

  // mojo::DataPipeDrainer::Client
  void OnDataAvailable(base::span<const uint8_t> data) override;

  void OnDataComplete() override;

#if BUILDFLAG(ARKWEB_UNITTESTS)
  void OnTransferDataWithSharedMemory(
        ::base::ReadOnlySharedMemoryRegion region,
        uint64_t buffer_size) override;
#endif

  mojo::Remote<network::mojom::URLLoader> remote_;
  mojo::Receiver<network::mojom::URLLoaderClient> receiver_{this};

  std::unique_ptr<mojo::DataPipeDrainer> pipe_drainer_;
  bool auto_draining_{true};
  mojo::ScopedDataPipeConsumerHandle body_;

  std::string body_content_;
  uint32_t total_bytes_read_{0};
  bool body_finished_{false};
  int32_t total_transfer_size_diff_{0};

  std::optional<network::URLLoaderCompletionStatus> completion_status_;

  std::vector<std::pair<net::RedirectInfo, network::mojom::URLResponseHeadPtr>>
      received_redirects_;
};

class ScopedMockContentBrowserClient : public TestContentBrowserClient {
 public:
  ScopedMockContentBrowserClient();
  ~ScopedMockContentBrowserClient() override;

  MOCK_METHOD(
      void,
      WillCreateURLLoaderFactory,
      (BrowserContext * browser_context,
       RenderFrameHost* frame,
       int render_process_id,
       URLLoaderFactoryType type,
       const url::Origin& request_initiator,
       const net::IsolationInfo& isolation_info,
       std::optional<int64_t> navigation_id,
       ukm::SourceIdObj ukm_source_id,
       network::URLLoaderFactoryBuilder& factory_builder,
       mojo::PendingRemote<network::mojom::TrustedURLLoaderHeaderClient>*
           header_client,
       bool* bypass_redirect_checks,
       bool* disable_secure_dns,
       network::mojom::URLLoaderFactoryOverridePtr* factory_override,
       scoped_refptr<base::SequencedTaskRunner>
           navigation_response_task_runner),
      (override));

 private:
  raw_ptr<ContentBrowserClient> old_browser_client_;
};

class TestPrefetchService final : public PrefetchService {
 public:
  explicit TestPrefetchService(BrowserContext* browser_context);
  ~TestPrefetchService() override;

  void PrefetchUrl(
      base::WeakPtr<PrefetchContainer> prefetch_container) override;
  void OnPrefetchCompletedOrFailed(
      PrefetchContainer& prefetch_container,
      const network::URLLoaderCompletionStatus& completion_status,
      const std::optional<int>& response_code) override;
  void EvictPrefetch(size_t index);

  std::vector<base::WeakPtr<PrefetchContainer>> prefetches_;
};

// Helper for testing prefetching-side (i.e. not serving-side) metrics including
// "PrefetchProxy.Prefetch.*" UMAs, `PrefetchReferringPageMetrics` and
// Preloading_Attempt UKMs.
class PrefetchingMetricsTestBase : public RenderViewHostTestHarness {
 public:
  PrefetchingMetricsTestBase();
  ~PrefetchingMetricsTestBase() override;

  const int kTotalTimeDuration = 4321;
  const int kConnectTimeDuration = 123;

  ukm::TestAutoSetUkmRecorder* test_ukm_recorder() {
    return test_ukm_recorder_.get();
  }

  const test::PreloadingAttemptUkmEntryBuilder* attempt_entry_builder() {
    return attempt_entry_builder_.get();
  }

  void SetUp() override;
  void TearDown() override;

  // Prefetch didn't receive any net errors nor non-redirect responses.
  // Use more specific methods below to check UKMs, if applicable.
  void ExpectPrefetchNoNetErrorOrResponseReceived(
      const base::HistogramTester& histogram_tester,
      bool is_eligible,
      bool browser_initiated_prefetch = false);

  // Prefetch was not started because it was not eligible.
  void ExpectPrefetchNotEligible(const base::HistogramTester& histogram_tester,
                                 PreloadingEligibility expected_eligibility,
                                 bool is_accurate = false,
                                 bool browser_initiated_prefetch = false);

  // Prefetch was started but failed before the final response nor any network
  // error is received.
  void ExpectPrefetchFailedBeforeResponseReceived(
      const base::HistogramTester& histogram_tester,
      PrefetchStatus expected_prefetch_status,
      bool is_accurate = false);

  // Prefetch was started but failed due to a network error, before the final
  // response is received.
  void ExpectPrefetchFailedNetError(
      const base::HistogramTester& histogram_tester,
      int expected_net_error_code,
      blink::mojom::SpeculationEagerness eagerness =
          blink::mojom::SpeculationEagerness::kImmediate,
      bool is_accurate_triggering = false,
      bool browser_initiated_prefetch = false);

  // Prefetch was started but failed on or after the final response is
  // received.
  void ExpectPrefetchFailedAfterResponseReceived(
      const base::HistogramTester& histogram_tester,
      net::HttpStatusCode expected_response_code,
      int expected_body_length,
      PrefetchStatus expected_prefetch_status);

  void ExpectPrefetchSuccess(const base::HistogramTester& histogram_tester,
                             int expected_body_length,
                             blink::mojom::SpeculationEagerness eagerness =
                                 blink::mojom::SpeculationEagerness::kImmediate,
                             bool is_accurate = false);

  // `navigate_url` is used as `MockNavigationHandle`'s URL to simulate a
  // navigation possibly using the prefetch. It is passed outside
  // `ExpectCorrectUkmLogsArgs` to keep `ExpectCorrectUkmLogsArgs` non-complex.
  ukm::SourceId ForceLogsUploadAndGetUkmId(
      GURL navigate_url = GURL("http://Not.Accurate.Trigger.Url/"));
  struct ExpectCorrectUkmLogsArgs {
    PreloadingEligibility eligibility = PreloadingEligibility::kEligible;
    PreloadingHoldbackStatus holdback = PreloadingHoldbackStatus::kAllowed;
    PreloadingTriggeringOutcome outcome = PreloadingTriggeringOutcome::kReady;
    PreloadingFailureReason failure = PreloadingFailureReason::kUnspecified;
    bool is_accurate = false;
    bool expect_ready_time = false;
    blink::mojom::SpeculationEagerness eagerness =
        blink::mojom::SpeculationEagerness::kImmediate;
  };
  void ExpectCorrectUkmLogs(
      ExpectCorrectUkmLogsArgs args,
      GURL navigate_url = GURL("http://Not.Accurate.Trigger.Url/"));

 private:
  std::unique_ptr<ukm::TestAutoSetUkmRecorder> test_ukm_recorder_;
  std::unique_ptr<test::PreloadingAttemptUkmEntryBuilder>
      attempt_entry_builder_;
};

// Helpers for parametrized tests for rearchitecturing/refactoring of the core
// classes of `preloading/prefetch` like `PrefetchService`, `PrefetchContainer`,
// etc.
// Although some features/parameters might be unrelated to some of the target
// classes/unit tests, we anyway apply the same sets of parameters to all the
// target unit tests, for uniformity, comprehensiveness and convenience.
//
// Usage:
// - Make *Test classes inherit
//   - `public WithPrefetchRearchParam` and
//   - `public ::testing::WithParamInterface<PrefetchRearchParam>`
//     (or its tuple etc. if other parameters are needed)
// - Call `InitRearchFeatures()` upon setup.
// - `INSTANTIATE_TEST_SUITE_P()` with
//   `testing::ValuesIn(PrefetchRearchParam::Params())`.
//
// Do not remove these classes and keep using them, even if there is no param to
// make it easy to add another param in the future.

struct PrefetchRearchParam final {
 public:
  static std::vector<PrefetchRearchParam> Params();

  bool prefetch_scheduler;
  bool prefetch_scheduler_progress_sync_best_effort;
  bool graceful_notification;
};

class WithPrefetchRearchParam {
 public:
  explicit WithPrefetchRearchParam(PrefetchRearchParam param);
  virtual ~WithPrefetchRearchParam();

  void InitRearchFeatures();

  const PrefetchRearchParam& rearch_param() { return param_; }

 private:
  PrefetchRearchParam param_;
  base::test::ScopedFeatureList feature_list_prefetch_scheduler_;
  base::test::ScopedFeatureList feature_list_graceful_notification_;
};

// A wrapper for `PrefetchService::SetInjectedEligibilityCheckForTesting`.
// - Provide `TestFuture`-based interface.
// - Cleanup `SetInjectedEligibilityCheckForTesting()` on dtor.
class PrefetchServiceInjectedEligibilityCheckFuture final {
 public:
  explicit PrefetchServiceInjectedEligibilityCheckFuture(
      PrefetchService& prefetch_service);
  ~PrefetchServiceInjectedEligibilityCheckFuture();

  using TestFutureType = base::test::TestFuture<
      PrefetchService::InjectedEligibilityCheckResultCallbackForTesting>;

  TestFutureType* operator->() { return &result_callback_future_; }

 private:
  raw_ref<PrefetchService> prefetch_service_;
  TestFutureType result_callback_future_;
};

}  // namespace content

#endif  // CONTENT_BROWSER_PRELOADING_PREFETCH_PREFETCH_TEST_UTIL_INTERNAL_H_