components/omnibox/composebox/test_composebox_query_controller.h

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

#ifndef COMPONENTS_OMNIBOX_COMPOSEBOX_TEST_COMPOSEBOX_QUERY_CONTROLLER_H_
#define COMPONENTS_OMNIBOX_COMPOSEBOX_TEST_COMPOSEBOX_QUERY_CONTROLLER_H_

#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "base/functional/callback.h"
#include "components/endpoint_fetcher/endpoint_fetcher.h"
#include "components/variations/variations_client.h"
#include "composebox_query_controller.h"
#include "third_party/lens_server_proto/lens_overlay_server.pb.h"

namespace lens {
class LensOverlayClientContext;
}  // namespace lens

class FakeEndpointFetcher : public endpoint_fetcher::EndpointFetcher {
 public:
  explicit FakeEndpointFetcher(endpoint_fetcher::EndpointResponse response);
  void PerformRequest(
      endpoint_fetcher::EndpointFetcherCallback endpoint_fetcher_callback,
      const char* key) override;

  bool disable_responding_ = false;

 private:
  endpoint_fetcher::EndpointResponse response_;
};

// Fake VariationsClient for testing.
class FakeVariationsClient : public variations::VariationsClient {
 public:
  ~FakeVariationsClient() override = default;

  bool IsOffTheRecord() const override;

  variations::mojom::VariationsHeadersPtr GetVariationsHeaders() const override;
};

// Helper for testing features that use the ComposeboxQueryController.
// The only logic in this class should be for setting up fake network responses
// and tracking sent request data to maximize testing coverage.
class TestComposeboxQueryController : public ComposeboxQueryController {
 public:
  explicit TestComposeboxQueryController(
      signin::IdentityManager* identity_manager,
      scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
      version_info::Channel channel,
      std::string locale,
      TemplateURLService* template_url_service,
      variations::VariationsClient* variations_client,
      bool send_lns_surface);
  ~TestComposeboxQueryController() override;

  // Mutators.
  void set_fake_cluster_info_response(
      lens::LensOverlayServerClusterInfoResponse response) {
    fake_cluster_info_response_ = response;
  }

  void set_next_cluster_info_request_should_return_error(
      bool set_next_cluster_info_request_should_return_error) {
    next_cluster_info_request_should_return_error_ =
        set_next_cluster_info_request_should_return_error;
  }

  void set_next_file_upload_request_should_return_error(
      bool set_next_file_upload_request_should_return_error) {
    next_file_upload_request_should_return_error_ =
        set_next_file_upload_request_should_return_error;
  }

  void set_enable_cluster_info_ttl(bool enable_cluster_info_ttl) {
    enable_cluster_info_ttl_ = enable_cluster_info_ttl;
  }

  void set_on_query_controller_state_changed_callback(
      QueryControllerStateChangedCallback callback) {
    on_query_controller_state_changed_callback_ = std::move(callback);
  }

  // Accessors.
  const int& num_cluster_info_fetch_requests_sent() const {
    return num_cluster_info_fetch_requests_sent_;
  }

  const int& num_file_upload_requests_sent() const {
    return num_file_upload_requests_sent_;
  }

  QueryControllerState query_controller_state() const {
    return query_controller_state_;
  }

  const GURL& last_sent_fetch_url() const { return last_sent_fetch_url_; }

  // Gets the last sent file upload request.
  std::optional<lens::LensOverlayServerRequest> last_sent_file_upload_request()
      const {
    return last_sent_file_upload_request_;
  }

  // Gets the last sent cors exempt headers.
  std::vector<std::string> last_sent_cors_exempt_headers() const {
    return last_sent_cors_exempt_headers_;
  }

  // Gets the client context used for the requests.
  lens::LensOverlayClientContext client_context() const {
    return ComposeboxQueryController::CreateClientContext();
  }

 protected:
  std::unique_ptr<endpoint_fetcher::EndpointFetcher> CreateEndpointFetcher(
      std::string request_string,
      const GURL& fetch_url,
      endpoint_fetcher::HttpMethod http_method,
      base::TimeDelta timeout,
      const std::vector<std::string>& request_headers,
      const std::vector<std::string>& cors_exempt_headers,
      UploadProgressCallback upload_progress_callback) override;

  void ResetRequestClusterInfoState(int session_id) override;

  // The fake response to return for cluster info requests.
  lens::LensOverlayServerClusterInfoResponse fake_cluster_info_response_;

  // The number of cluster info fetch requests sent by the query controller.
  int num_cluster_info_fetch_requests_sent_ = 0;

  // The number of file upload requests sent by the query controller.
  int num_file_upload_requests_sent_ = 0;

  // If true, the next cluster info request will return an error.
  bool next_cluster_info_request_should_return_error_ = false;

  // If true, the next file upload request will return an error.
  bool next_file_upload_request_should_return_error_ = false;

  // If true, the cluster info will expire when the TTL expires as normal.
  // Set to false by default to prevent flakiness in tests that expect the
  // cluster info to be available.
  bool enable_cluster_info_ttl_ = false;

  // The last url for which a fetch request was sent by the query controller.
  GURL last_sent_fetch_url_;

  // The last sent file upload request.
  std::optional<lens::LensOverlayServerRequest> last_sent_file_upload_request_;

  // The last sent cors exempt headers.
  std::vector<std::string> last_sent_cors_exempt_headers_;
};

#endif  // COMPONENTS_OMNIBOX_COMPOSEBOX_TEST_COMPOSEBOX_QUERY_CONTROLLER_H_