// 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. #include "chrome/browser/ui/webui/cr_components/searchbox/contextual_searchbox_handler.h" #include #include #include #include #include #include "base/base64.h" #include "base/compiler_specific.h" #include "base/memory/raw_ptr.h" #include "base/strings/strcat.h" #include "base/test/bind.h" #include "base/test/gmock_move_support.h" #include "base/test/metrics/user_action_tester.h" #include "base/test/mock_callback.h" #include "base/test/scoped_feature_list.h" #include "base/test/test_future.h" #include "base/time/time.h" #include "base/types/expected.h" #include "base/unguessable_token.h" #include "base/version_info/channel.h" #include "chrome/browser/autocomplete/aim_eligibility_service_factory.h" #include "chrome/browser/autocomplete/chrome_autocomplete_scheme_classifier.h" #include "chrome/browser/contextual_search/contextual_search_service_factory.h" #include "chrome/browser/contextual_search/contextual_search_web_contents_helper.h" #include "chrome/browser/contextual_tasks/active_task_context_provider_impl.h" #include "chrome/browser/contextual_tasks/contextual_tasks.mojom.h" #include "chrome/browser/contextual_tasks/contextual_tasks_context_service.h" #include "chrome/browser/contextual_tasks/contextual_tasks_context_service_factory.h" #include "chrome/browser/contextual_tasks/contextual_tasks_service_factory.h" #include "chrome/browser/contextual_tasks/contextual_tasks_ui_interface.h" #include "chrome/browser/contextual_tasks/smart_tab_sharing_metrics.h" #include "chrome/browser/feature_engagement/tracker_factory.h" #include "chrome/browser/media/webrtc/fake_desktop_media_picker_factory.h" #include "chrome/browser/search_engines/template_url_service_factory.h" #include "chrome/browser/signin/identity_manager_factory.h" #include "chrome/browser/tab_list/mock_tab_list_interface.h" #include "chrome/browser/tab_list/tab_list_interface.h" #include "chrome/browser/ui/browser_window/public/browser_window_features.h" #include "chrome/browser/ui/browser_window/public/browser_window_interface.h" #include "chrome/browser/ui/browser_window/test/mock_browser_window_interface.h" #include "chrome/browser/ui/contextual_search/tab_contextualization_controller.h" #include "chrome/browser/ui/tab_ui_helper.h" #include "chrome/browser/ui/tabs/alert/tab_alert_controller.h" #include "chrome/browser/ui/tabs/public/tab_features.h" #include "chrome/browser/ui/tabs/tab_model.h" #include "chrome/browser/ui/views/drive_picker_host/drive_picker_host_controller.h" #include "chrome/browser/ui/views/drive_picker_host/drive_picker_sanitizer.h" #include "chrome/browser/ui/webui/cr_components/composebox/composebox_handler.h" #include "chrome/browser/ui/webui/drive_picker_host/drive_picker_host_request.h" #include "chrome/browser/ui/webui/new_tab_page/composebox/variations/composebox_fieldtrial.h" #include "chrome/browser/ui/webui/searchbox/contextual_searchbox_test_utils.h" #include "chrome/browser/ui/webui/searchbox/searchbox_test_utils.h" #include "chrome/browser/ui/webui/webui_embedding_context.h" #include "chrome/common/webui_url_constants.h" #include "components/contextual_search/contextual_search_metrics_recorder.h" #include "components/contextual_search/contextual_search_service.h" #include "components/contextual_search/footprints/public/drive_disclaimer_controller.h" #include "components/contextual_search/footprints/public/fpop_service.h" #include "components/contextual_search/footprints/public/proto/footprints_oneplatform.pb.h" #include "components/contextual_search/internal/test_composebox_query_controller.h" #include "components/contextual_search/pref_names.h" #include "components/contextual_tasks/public/features.h" #include "components/contextual_tasks/public/mock_contextual_tasks_service.h" #include "components/contextual_tasks/public/prefs.h" #include "components/contextual_tasks/public/query_contextualizer.h" #include "components/feature_engagement/test/mock_tracker.h" #include "components/lens/lens_features.h" #include "components/lens/lens_overlay_invocation_source.h" #include "components/lens/proto/server/lens_overlay_response.pb.h" #include "components/omnibox/browser/autocomplete_input.h" #include "components/omnibox/browser/autocomplete_match.h" #include "components/omnibox/browser/autocomplete_result.h" #include "components/omnibox/browser/fake_autocomplete_controller.h" #include "components/omnibox/browser/mock_aim_eligibility_service.h" #include "components/omnibox/browser/searchbox.mojom.h" #include "components/omnibox/common/composebox_features.h" #include "components/omnibox/common/omnibox_features.h" #include "components/omnibox/composebox/composebox_query.mojom.h" #include "components/omnibox/composebox/contextual_search_mojom_traits.h" #include "components/prefs/pref_service.h" #include "components/search/ntp_features.h" #include "components/tabs/public/mock_tab_interface.h" #include "content/public/browser/desktop_capture.h" #include "content/public/browser/navigation_entry.h" #include "content/public/browser/web_contents.h" #include "content/public/browser/web_contents_delegate.h" #include "content/public/test/desktop_capture_test_utils.h" #include "content/public/test/navigation_simulator.h" #include "content/public/test/test_navigation_observer.h" #include "content/public/test/web_contents_tester.h" #include "media/base/media_switches.h" #include "mojo/public/cpp/base/unguessable_token_mojom_traits.h" #include "mojo/public/cpp/test_support/fake_message_dispatch_context.h" #include "mojo/public/cpp/test_support/test_utils.h" #include "testing/gmock/include/gmock/gmock.h" #include "testing/gtest/include/gtest/gtest.h" #include "third_party/webrtc/modules/desktop_capture/desktop_capturer.h" #include "third_party/webrtc/modules/desktop_capture/desktop_frame.h" #include "ui/base/base_window.h" #include "ui/base/test/mock_base_window.h" #include "ui/base/unowned_user_data/unowned_user_data_host.h" #include "ui/base/webui/web_ui_util.h" #include "ui/gfx/codec/png_codec.h" #if BUILDFLAG(IS_CHROMEOS) #include "chromeos/ash/components/network/network_handler_test_helper.h" #endif // BUILDFLAG(IS_CHROMEOS) #if BUILDFLAG(IS_MAC) #include "base/mac/mac_util.h" #endif // BUILDFLAG(IS_MAC) using contextual_search::SessionState; namespace { constexpr char kClientUploadDurationQueryParameter[] = "cud"; constexpr char kQuerySubmissionTimeQueryParameter[] = "qsubts"; constexpr char kQueryText[] = "query"; constexpr char kUdmQueryParameter[] = "udm"; constexpr char kUdmQueryParameterValue[] = "50"; GURL StripTimestampsFromAimUrl(const GURL& url) { std::string qsubts_param; EXPECT_TRUE(net::GetValueForKeyInQuery( url, kQuerySubmissionTimeQueryParameter, &qsubts_param)); std::string cud_param; EXPECT_TRUE(net::GetValueForKeyInQuery( url, kClientUploadDurationQueryParameter, &cud_param)); GURL result_url = url; result_url = net::AppendOrReplaceQueryParameter( result_url, kQuerySubmissionTimeQueryParameter, std::nullopt); result_url = net::AppendOrReplaceQueryParameter( result_url, kClientUploadDurationQueryParameter, std::nullopt); return result_url; } class FakeContextualTasksUIInterface : public contextual_tasks::ContextualTasksUIInterface { public: FakeContextualTasksUIInterface() { page_receiver_ = page_remote_.BindNewPipeAndPassReceiver(); } ~FakeContextualTasksUIInterface() override = default; // TaskInfoDelegate: const std::optional& GetTaskId() override { static const std::optional id; return id; } void SetTaskId(std::optional id) override {} const std::optional& GetThreadId() override { static const std::optional id; return id; } void SetThreadId(std::optional id) override {} const std::optional& GetThreadTitle() override { static const std::optional title; return title; } void SetThreadTitle(std::optional title) override {} void SetIsAiPage(bool is_ai_page) override {} void UpdateStateFromUrl(const GURL& url) override {} bool IsShownInTab() override { return false; } BrowserWindowInterface* GetBrowser() override { return nullptr; } content::WebContents* GetWebUIWebContents() override { return nullptr; } void OnZeroStateChange(bool is_zero_state) override {} void SetInNlm(bool in_nlm) override {} void PushTaskDetailsToPage(std::optional id, const GURL& url, bool replace_navigation_entry) override {} void PrepareForTaskChange() override {} void OnTaskChanged() override {} // ContextualTasksUIInterface: mojo::Remote& GetPageRemote() override { return page_remote_; } contextual_tasks::ContextualTasksAutoSuggestionManager* GetAutoSuggestionManager() override { return nullptr; } Profile* GetProfile() override { return nullptr; } void TransferNavigationToEmbeddedPage( content::OpenURLParams params) override {} void CloseSidePanel() override {} void OnSidePanelStateChanged() override {} void OnActiveTabContextStatusChanged() override {} void OnLensOverlayStateChanged( bool is_showing, std::optional invocation_source) override {} bool IsLensOverlayShowing() const override { return false; } void StartPlatformVoiceRecognition() override {} void OnVoiceTranscribed(const std::string& query) override {} void OnPageContextEligibilityChecked(bool is_page_context_eligible) override { } bool IsActiveTabContextSuggestionShowing() const override { return false; } bool CanExpandToFullTab() const override { return false; } void MoveTaskUiToNewTab() override {} void UpdateExpandButtonEnabled(bool enabled) override {} GURL GetWebUiUrl() override { return GURL(); } void PostAimMessage(const lens::ClientToAimMessage& message) override {} contextual_search::ContextualSearchSessionHandle* GetOrCreateContextualSessionHandle() override { return nullptr; } std::unique_ptr TakeInputStateModel() override { return nullptr; } std::vector GetRestoredTabIds() override { return {}; } void OnRestoredTabsFetched( std::vector tabs) override {} void SetComposeboxHandler( contextual_tasks::ContextualTasksComposeboxHandlerInterface* handler) override {} const GURL& GetInnerFrameUrl() const override { return GURL::EmptyGURL(); } content::WebContents* GetInnerWebContents() const override { return nullptr; } bool IsContextualTasksEligibleOnInit() const override { return false; } bool IsInitComplete() override { return false; } void OnInitComplete() override {} void AddObserver(Observer* observer) override {} void RemoveObserver(Observer* observer) override {} bool is_history_thread_loading() const override { return false; } void set_is_history_thread_loading(bool loading) override {} private: mojo::Remote page_remote_; mojo::PendingReceiver page_receiver_; }; class FakeContextualSearchboxHandler : public ContextualSearchboxHandler { public: FakeContextualSearchboxHandler( mojo::PendingReceiver pending_page_handler, mojo::PendingRemote pending_page, Profile* profile, content::WebContents* web_contents, std::unique_ptr client, GetSessionHandleCallback get_session_callback) : ContextualSearchboxHandler(std::move(pending_page_handler), std::move(pending_page), profile, web_contents, std::move(client), std::move(get_session_callback)) {} ~FakeContextualSearchboxHandler() override { #if !BUILDFLAG(IS_ANDROID) OnDrivePickerDisconnected(); #endif } using SearchboxHandler::client; using SearchboxHandler::SetAutocompleteControllerForTesting; // searchbox::mojom::PageHandler void ExecuteAction(uint8_t line, uint8_t action_index, const GURL& url, base::TimeTicks match_selection_timestamp, uint8_t mouse_button, bool alt_key, bool ctrl_key, bool meta_key, bool shift_key) override {} void OnThumbnailRemoved() override {} contextual_search::ContextualSearchMetricsRecorder* GetMetricsRecorder() { return ContextualSearchboxHandler::GetMetricsRecorder(); } void NotifySessionStateChanged( contextual_search::SessionState session_state) { GetContextualSessionHandle() ->GetMetricsRecorder() ->NotifySessionStateChanged(session_state); } using ContextualSearchboxHandler::InitializeInputStateModel; contextual_search::InputStateModel* input_state_model() { return input_state_model_.get(); } void set_smart_tab_sharing_active_override(bool active) { smart_tab_sharing_active_override_ = active; } void clear_smart_tab_sharing_active_for_thread() { smart_tab_sharing_active_for_thread_.reset(); } bool IsSmartTabSharingActive() const override { if (smart_tab_sharing_active_override_.has_value()) { return *smart_tab_sharing_active_override_; } return ContextualSearchboxHandler::IsSmartTabSharingActive(); } void SetDrivePickerController( std::unique_ptr controller) { drive_picker_controller_ = std::move(controller); } void SetDriveDisclaimerController( std::unique_ptr controller) { drive_disclaimer_controller_ = std::move(controller); } drive_picker::DriveDisclaimerController* drive_disclaimer_controller() { return drive_disclaimer_controller_.get(); } bool IsDrivePickerReceiverBound() const { return drive_picker_result_handler_receiver_.is_bound(); } void OnDrivePickerDisconnected() { ContextualSearchboxHandler::OnDrivePickerDisconnected(); } omnibox::InputState GetValidInputStateForTesting() { return GetValidInputState(); } contextual_tasks::ContextualTasksUIInterface* GetContextualTasksUiInterface() override { return &fake_ui_interface_; } bool IsContextualSearchTabSharingEligible() const override { return tab_sharing_eligible_; } void set_tab_sharing_eligible(bool eligible) { tab_sharing_eligible_ = eligible; } private: std::optional smart_tab_sharing_active_override_; bool tab_sharing_eligible_ = true; FakeContextualTasksUIInterface fake_ui_interface_; }; class MockDrivePickerHostController : public DrivePickerHostController { public: MockDrivePickerHostController( Profile* profile, BrowserWindowInterface* browser_window_interface) : DrivePickerHostController(profile, browser_window_interface) {} ~MockDrivePickerHostController() override = default; MOCK_METHOD(void, ShowDrivePickerHost, (std::unique_ptr), (override)); }; class MockFpopService : public contextual_search::FpopService { public: MOCK_METHOD( void, GetFacs, (const footprints::oneplatform::GetFacsRequest& request, base::OnceCallback callback), (override)); MOCK_METHOD( void, UpdateActivityControlsSettings, (const footprints::oneplatform::UpdateActivityControlsSettingsRequest& request, base::OnceCallback callback), (override)); MOCK_METHOD( void, ShouldShowMobileConsentFlow, (const footprints::oneplatform::ShouldShowMobileConsentFlowRequest& request, base::OnceCallback callback), (override)); }; class MockContextualTasksContextService : public contextual_tasks::ContextualTasksContextService { public: explicit MockContextualTasksContextService(Profile* profile) : ContextualTasksContextService(profile) {} MOCK_METHOD(void, GetRelevantTabsForConversationThread, (const contextual_tasks::TabSelectionOptions&, const contextual_tasks::ConversationThread&, const std::vector&, base::OnceCallback< void(std::vector>)>), (override)); MOCK_METHOD(void, OnTypedQuery, (base::WeakPtr), (override)); }; } // namespace // TODO(crbug.com/458086158): Make dedicated unit tests for the // ContextualSessionHandle for testing query controller and metrics recorder // interactions. class ContextualSearchboxHandlerTest : public ContextualSearchboxHandlerTestHarness { public: ~ContextualSearchboxHandlerTest() override = default; TestingProfile::TestingFactories GetTestingFactories() const override { auto factories = ContextualSearchboxHandlerTestHarness::GetTestingFactories(); factories.push_back(TestingProfile::TestingFactory{ AimEligibilityServiceFactory::GetInstance(), base::BindRepeating([](content::BrowserContext* context) -> std::unique_ptr { Profile* profile = Profile::FromBrowserContext(context); return std::make_unique( *profile->GetPrefs(), /*template_url_service=*/nullptr, /*url_loader_factory=*/nullptr, /*identity_manager=*/nullptr); })}); return factories; } void SetUp() override { ContextualSearchboxHandlerTestHarness::SetUp(); auto* aim_service = static_cast( AimEligibilityServiceFactory::GetForProfile(profile())); if (aim_service) { aim_service->config().add_input_type_configs()->set_input_type( omnibox::InputType::INPUT_TYPE_LENS_IMAGE); aim_service->config().add_input_type_configs()->set_input_type( omnibox::InputType::INPUT_TYPE_LENS_FILE); aim_service->config().add_input_type_configs()->set_input_type( omnibox::InputType::INPUT_TYPE_BROWSER_TAB); auto* canvas_config = aim_service->config().add_tool_configs(); canvas_config->set_tool(omnibox::ToolMode::TOOL_MODE_CANVAS); auto* canvas_rule = canvas_config->mutable_rule(); canvas_rule->set_tool(omnibox::ToolMode::TOOL_MODE_CANVAS); canvas_rule->add_allowed_input_types( omnibox::InputType::INPUT_TYPE_LENS_IMAGE); } // TODO(crbug.com/503732217): Fix tests to support lazy fetching of cluster // info and enable this feature by default in tests. scoped_feature_list_.InitWithFeatures( {omnibox::kComposeboxDriveContextMenuOption, omnibox::kForceDriveDisclaimerAccepted}, {contextual_tasks::kContextualTasksLazyFetchClusterInfo}); auto query_controller_config_params = std::make_unique< contextual_search::ContextualSearchContextController::ConfigParams>(); query_controller_config_params->send_lns_surface = false; query_controller_config_params->enable_viewport_images = true; auto query_controller_ptr = std::make_unique( /*identity_manager=*/nullptr, url_loader_factory(), version_info::Channel::UNKNOWN, "en-US", template_url_service(), fake_variations_client(), std::move(query_controller_config_params)); query_controller_ = query_controller_ptr.get(); auto metrics_recorder_ptr = std::make_unique(); ON_CALL(*metrics_recorder_ptr, RecordZeroSuggestClick) .WillByDefault(testing::Invoke( metrics_recorder_ptr.get(), &MockContextualSearchMetricsRecorder::RecordZeroSuggestClickBase)); service_ = ContextualSearchServiceFactory::GetForProfile(profile()); contextual_session_handle_ = service_->CreateSessionForTesting( std::move(query_controller_ptr), std::move(metrics_recorder_ptr)); // Check the search content sharing settings to notify the session handle // that the policy has been checked. contextual_session_handle_->CheckSearchContentSharingSettings( profile()->GetPrefs()); webui::SetBrowserWindowInterface(web_contents(), &mock_browser_window_interface_); ON_CALL(mock_browser_window_interface_, GetUnownedUserDataHost()) .WillByDefault(testing::ReturnRef(unowned_user_data_host_)); ON_CALL(mock_browser_window_interface_, GetWindow()) .WillByDefault(testing::Return(&mock_base_window_)); ON_CALL(mock_browser_window_interface_, GetFeatures()) .WillByDefault(testing::ReturnRef(browser_window_features_)); ON_CALL(mock_base_window_, GetNativeWindow()) .WillByDefault( testing::Return(web_contents()->GetTopLevelNativeWindow())); web_contents()->SetDelegate(&delegate_); handler_ = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting( [&]() { return contextual_session_handle_.get(); })); auto mock_drive_picker_controller = std::make_unique( profile(), &mock_browser_window_interface_); handler().SetDrivePickerController(std::move(mock_drive_picker_controller)); // Drain the Mojo pipe and clear setup-related calls to searchbox page. mock_searchbox_page_.FlushForTesting(); base::RunLoop().RunUntilIdle(); testing::Mock::VerifyAndClearExpectations(&mock_searchbox_page_); ON_CALL(query_controller(), GetFileInfo) .WillByDefault(testing::Invoke(query_controller_.get(), &MockQueryController::FakeGetFileInfo)); ON_CALL(query_controller(), CreateSearchUrl) .WillByDefault( [](auto&& request_info, base::OnceCallback callback) { GURL url("https://www.google.com/search?q=" + request_info->query_text); for (auto const& [key, val] : request_info->additional_params) { url = net::AppendOrReplaceQueryParameter(url, key, val); } url = net::AppendOrReplaceQueryParameter( url, kQuerySubmissionTimeQueryParameter, "0"); url = net::AppendOrReplaceQueryParameter( url, kClientUploadDurationQueryParameter, "0"); std::move(callback).Run(url); }); } void SubmitQueryAndWaitForNavigation() { content::TestNavigationObserver navigation_observer(web_contents()); handler().SubmitQuery(kQueryText, 1, false, false, false, false, /*is_voice_search=*/false); auto navigation = content::NavigationSimulator::CreateFromPending( web_contents()->GetController()); ASSERT_TRUE(navigation); auto callback = delegate_.TakeCallback(); if (callback) { std::move(callback).Run(*(navigation->GetNavigationHandle())); } navigation->Commit(); navigation_observer.Wait(); } FakeContextualSearchboxHandler& handler() { return *handler_; } MockQueryController& query_controller() { return *query_controller_; } void SetupScreenshotUploadConfig() { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); scoped_config().config.mutable_composebox()->set_max_num_files(5); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_max_size_bytes(1024 * 1024); scoped_config() .config.mutable_composebox() ->mutable_image_upload() ->set_mime_types_allowed("image/png"); } void SetUpMockFpopService(bool accepted) { SetUpMockFpopServiceWithStatus( accepted ? drive_picker::DriveDisclaimerController:: ConsentEligibilityStatus::kAlreadyConsented : drive_picker::DriveDisclaimerController:: ConsentEligibilityStatus::kCanConsent); } void SetUpMockFpopServiceWithStatus( drive_picker::DriveDisclaimerController::ConsentEligibilityStatus status) { auto mock_fpop_service = std::make_unique(); EXPECT_CALL(*mock_fpop_service, ShouldShowMobileConsentFlow(testing::_, testing::_)) .WillOnce([status](const auto& request, auto callback) { footprints::oneplatform::ShouldShowMobileConsentFlowResponse response; response.mutable_should_show_flow_result() ->mutable_eligibility() ->set_status(static_cast(status)); std::move(callback).Run(true, response); }); handler().SetDriveDisclaimerController( std::make_unique( std::move(mock_fpop_service))); } MockContextualSearchMetricsRecorder* GetMetricsRecorderPtr() { if (handler_) { /* Cast since what we pass into the handler was a mock version to begin * with. */ return static_cast( handler_->GetMetricsRecorder()); } return nullptr; } void TearDown() override { query_controller_ = nullptr; mock_drive_picker_controller_ = nullptr; metrics_recorder_ = nullptr; service_ = nullptr; handler_.reset(); contextual_session_handle_.reset(); delegate_.ClearCallback(); ContextualSearchboxHandlerTestHarness::TearDown(); } protected: testing::NiceMock mock_searchbox_page_; std::unique_ptr handler_; std::unique_ptr contextual_session_handle_; testing::NiceMock mock_browser_window_interface_; BrowserWindowFeatures browser_window_features_; testing::NiceMock mock_base_window_; ui::UnownedUserDataHost unowned_user_data_host_; #if BUILDFLAG(IS_CHROMEOS) ash::NetworkHandlerTestHelper network_handler_test_helper_; #endif // BUILDFLAG(IS_CHROMEOS) raw_ptr query_controller_; raw_ptr service_; raw_ptr metrics_recorder_; private: base::test::ScopedFeatureList scoped_feature_list_; TestWebContentsDelegate delegate_; raw_ptr mock_drive_picker_controller_; }; TEST_F(ContextualSearchboxHandlerTest, DriveDisclaimerControllerDoesNotCrashForIncognito) { // Use an incognito profile. Profile* incognito_profile = profile()->GetPrimaryOTRProfile(/*create_if_needed=*/true); ASSERT_TRUE(incognito_profile->IsOffTheRecord()); // IdentityManager is null for incognito profiles in this test setup. ASSERT_EQ(nullptr, IdentityManagerFactory::GetForProfile(incognito_profile)); // Feature is enabled by default in SetUp(). // Create a handler for the incognito profile. testing::NiceMock incognito_mock_searchbox_page; auto incognito_handler = std::make_unique( mojo::PendingReceiver(), incognito_mock_searchbox_page.BindAndGetRemote(), incognito_profile, web_contents(), std::make_unique(), base::BindLambdaForTesting( [&]() { return contextual_session_handle_.get(); })); // Without the fix, the constructor above would have crashed. // With the fix, the disclaimer controller should be null. EXPECT_EQ(nullptr, incognito_handler->drive_disclaimer_controller()); } TEST_F(ContextualSearchboxHandlerTest, SessionStarted) { SessionState state_arg = SessionState::kNone; auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(query_controller(), InitializeIfNeeded); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .WillOnce(testing::SaveArg<0>(&state_arg)); handler().NotifySessionStarted(); EXPECT_EQ(state_arg, SessionState::kSessionStarted); } TEST_F(ContextualSearchboxHandlerTest, ActivateMetricsFunnel) { auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, ActivateMetricsFunnel("AiMode")).Times(1); handler().ActivateMetricsFunnel("AiMode"); } TEST_F(ContextualSearchboxHandlerTest, AddFile_Pdf) { searchbox::mojom::SelectedFileInfoPtr file_info = searchbox::mojom::SelectedFileInfo::New(); file_info->file_name = "test.pdf"; file_info->selection_time = base::Time::Now(); file_info->mime_type = "application/pdf"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); base::MockCallback callback; base::UnguessableToken controller_file_info_token; std::optional callback_token; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce(testing::SaveArg<0>(&controller_file_info_token)); EXPECT_CALL(callback, Run).WillOnce([&](const auto& result) { ASSERT_TRUE(result.has_value()); callback_token = result.value(); }); handler().AddFileContext(std::move(file_info), std::move(file_data), callback.Get()); EXPECT_EQ(callback_token, controller_file_info_token); } TEST_F(ContextualSearchboxHandlerTest, AddFile_Image) { searchbox::mojom::SelectedFileInfoPtr file_info = searchbox::mojom::SelectedFileInfo::New(); file_info->file_name = "test.png"; file_info->selection_time = base::Time::Now(); file_info->mime_type = "application/image"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); std::optional image_options; base::UnguessableToken controller_file_info_token; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& file_token, auto, std::optional options_arg) { controller_file_info_token = file_token; image_options = std::move(options_arg); }); base::MockCallback callback; std::optional callback_token; EXPECT_CALL(callback, Run).WillOnce([&](const auto& result) { ASSERT_TRUE(result.has_value()); callback_token = result.value(); }); handler().AddFileContext(std::move(file_info), std::move(file_data), callback.Get()); EXPECT_EQ(callback_token, controller_file_info_token); EXPECT_TRUE(image_options.has_value()); auto image_upload = scoped_config().Get().config.composebox().image_upload(); EXPECT_EQ(image_options->max_height, image_upload.downscale_max_image_height()); EXPECT_EQ(image_options->max_size, image_upload.downscale_max_image_size()); EXPECT_EQ(image_options->max_width, image_upload.downscale_max_image_width()); EXPECT_EQ(image_options->compression_quality, image_upload.image_compression_quality()); } TEST_F(ContextualSearchboxHandlerTest, SetActiveToolModeHistogram) { base::HistogramTester histogram_tester; handler().SetActiveToolMode(omnibox::TOOL_MODE_CANVAS, /*is_set_by_aim=*/false); handler().SetActiveToolMode(omnibox::TOOL_MODE_DEEP_SEARCH, /*is_set_by_aim=*/false); handler().SetActiveToolMode(omnibox::TOOL_MODE_CANVAS, /*is_set_by_aim=*/false); histogram_tester.ExpectTotalCount( "ContextualSearch.Tools.ChangedByAIM.NewTabPage", 0); } TEST_F(ContextualSearchboxHandlerTest, SetActiveModelModeHistogram) { base::HistogramTester histogram_tester; handler().SetActiveModelMode(omnibox::MODEL_MODE_GEMINI_PRO, /*is_set_by_aim=*/false); handler().SetActiveModelMode(omnibox::MODEL_MODE_GEMINI_REGULAR, /*is_set_by_aim=*/false); handler().SetActiveModelMode(omnibox::MODEL_MODE_GEMINI_PRO, /*is_set_by_aim=*/false); histogram_tester.ExpectTotalCount( "ContextualSearch.Models.ChangedByAIM.NewTabPage", 0); } TEST_F(ContextualSearchboxHandlerTest, SetActiveToolModeSetByServerHistogram) { base::HistogramTester histogram_tester; handler().SetActiveToolMode(omnibox::TOOL_MODE_CANVAS, /*is_set_by_aim=*/true); handler().SetActiveToolMode(omnibox::TOOL_MODE_DEEP_SEARCH, /*is_set_by_aim=*/true); handler().SetActiveToolMode(omnibox::TOOL_MODE_CANVAS, /*is_set_by_aim=*/true); histogram_tester.ExpectBucketCount( "ContextualSearch.Tools.ChangedByAIM.NewTabPage", omnibox::TOOL_MODE_CANVAS, 2); histogram_tester.ExpectBucketCount( "ContextualSearch.Tools.ChangedByAIM.NewTabPage", omnibox::TOOL_MODE_DEEP_SEARCH, 1); histogram_tester.ExpectTotalCount( "ContextualSearch.Tools.ChangedByAIM.NewTabPage", 3); } TEST_F(ContextualSearchboxHandlerTest, SetActiveModelModeSetByServerHistogram) { base::HistogramTester histogram_tester; handler().SetActiveModelMode(omnibox::MODEL_MODE_GEMINI_PRO, /*is_set_by_aim=*/true); handler().SetActiveModelMode(omnibox::MODEL_MODE_GEMINI_REGULAR, /*is_set_by_aim=*/true); handler().SetActiveModelMode(omnibox::MODEL_MODE_GEMINI_PRO, /*is_set_by_aim=*/true); histogram_tester.ExpectBucketCount( "ContextualSearch.Models.ChangedByAIM.NewTabPage", omnibox::MODEL_MODE_GEMINI_PRO, 2); histogram_tester.ExpectBucketCount( "ContextualSearch.Models.ChangedByAIM.NewTabPage", omnibox::MODEL_MODE_GEMINI_REGULAR, 1); histogram_tester.ExpectTotalCount( "ContextualSearch.Models.ChangedByAIM.NewTabPage", 3); } TEST_F(ContextualSearchboxHandlerTest, ClearFiles) { searchbox::mojom::SelectedFileInfoPtr file_info = searchbox::mojom::SelectedFileInfo::New(); file_info->file_name = "test.png"; file_info->selection_time = base::Time::Now(); file_info->mime_type = "application/image"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); std::optional image_options; base::UnguessableToken controller_file_info_token; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& file_token, auto, std::optional options_arg) { controller_file_info_token = file_token; image_options = std::move(options_arg); }); base::MockCallback callback; std::optional callback_token; EXPECT_CALL(callback, Run).WillOnce([&](const auto& result) { ASSERT_TRUE(result.has_value()); callback_token = result.value(); }); handler().AddFileContext(std::move(file_info), std::move(file_data), callback.Get()); EXPECT_EQ(handler().GetUploadedContextTokens().size(), 1u); EXPECT_CALL(query_controller(), ClearFiles).Times(0); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); handler().ClearFiles(/*should_block_auto_suggested_tabs=*/false); EXPECT_EQ(handler().GetUploadedContextTokens().size(), 0u); } TEST_F(ContextualSearchboxHandlerTest, AddFile_PolicyDisabled) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kDisabled)); searchbox::mojom::SelectedFileInfoPtr file_info = searchbox::mojom::SelectedFileInfo::New(); file_info->file_name = "test.pdf"; file_info->selection_time = base::Time::Now(); file_info->mime_type = "application/pdf"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); base::MockCallback callback; base::expected callback_result = base::ok(base::UnguessableToken()); // Initialize with dummy EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(0); EXPECT_CALL(callback, Run).WillOnce(testing::SaveArg<0>(&callback_result)); handler().AddFileContext(std::move(file_info), std::move(file_data), callback.Get()); EXPECT_FALSE(callback_result.has_value()); EXPECT_EQ(callback_result.error(), contextual_search::ContextUploadErrorType::kBrowserProcessingError); } TEST_F(ContextualSearchboxHandlerTest, AddFile_PolicyToggled) { // Start with disabled. profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kDisabled)); searchbox::mojom::SelectedFileInfoPtr file_info_1 = searchbox::mojom::SelectedFileInfo::New(); file_info_1->file_name = "test1.pdf"; file_info_1->mime_type = "application/pdf"; std::vector test_data_1 = {1}; auto test_data_span_1 = base::span(test_data_1); mojo_base::BigBuffer file_data_1(test_data_span_1); base::MockCallback callback_1; base::expected callback_result_1 = base::ok(base::UnguessableToken()); // Initialize with dummy EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(0); EXPECT_CALL(callback_1, Run) .WillOnce(testing::SaveArg<0>(&callback_result_1)); handler().AddFileContext(std::move(file_info_1), std::move(file_data_1), callback_1.Get()); EXPECT_FALSE(callback_result_1.has_value()); EXPECT_EQ(callback_result_1.error(), contextual_search::ContextUploadErrorType::kBrowserProcessingError); // Enable policy. profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); searchbox::mojom::SelectedFileInfoPtr file_info_2 = searchbox::mojom::SelectedFileInfo::New(); file_info_2->file_name = "test2.pdf"; file_info_2->mime_type = "application/pdf"; std::vector test_data_2 = {2}; auto test_data_span_2 = base::span(test_data_2); mojo_base::BigBuffer file_data_2(test_data_span_2); base::MockCallback callback_2; std::optional callback_token_2; EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(1); EXPECT_CALL(callback_2, Run).WillOnce([&](const auto& result) { ASSERT_TRUE(result.has_value()); callback_token_2 = result.value(); }); handler().AddFileContext(std::move(file_info_2), std::move(file_data_2), callback_2.Get()); EXPECT_TRUE(callback_token_2.has_value()); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_PolicyDisabled) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kDisabled)); std::string file_name = "test.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); std::optional image_options; base::MockCallback callback; base::expected callback_result = base::ok(base::UnguessableToken()); // Initialize with dummy EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(0); EXPECT_CALL(callback, Run).WillOnce(testing::SaveArg<0>(&callback_result)); handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), image_options, callback.Get()); EXPECT_FALSE(callback_result.has_value()); EXPECT_EQ(callback_result.error(), contextual_search::ContextUploadErrorType::kBrowserProcessingError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_PolicyEnabled) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); scoped_config().config.mutable_composebox()->set_max_num_files(5); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_max_size_bytes(100); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_mime_types_allowed("application/pdf"); std::string file_name = "test.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); std::optional image_options; base::MockCallback callback; base::UnguessableToken controller_file_info_token; std::optional callback_token; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce(testing::SaveArg<0>(&controller_file_info_token)); EXPECT_CALL(callback, Run).WillOnce([&](const auto& result) { ASSERT_TRUE(result.has_value()); callback_token = result.value(); }); handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), image_options, callback.Get()); EXPECT_EQ(callback_token, controller_file_info_token); EXPECT_TRUE(callback_token.has_value()); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_DeepSearchNotAllowed) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); // Set the active tool to Deep Search. handler().input_state_model()->setActiveTool(omnibox::TOOL_MODE_DEEP_SEARCH); std::string file_name = "test.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingFileUploadNotAllowedError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_DisabledInputType) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); // Set the active tool to something other than Unspecified or Deep Search. omnibox::InputState state = handler().input_state_model()->GetInputState(); state.active_tool = omnibox::TOOL_MODE_IMAGE_GEN; state.disabled_input_types.push_back(omnibox::INPUT_TYPE_LENS_FILE); handler().input_state_model()->set_state_for_testing(state); std::string file_name = "test.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingUnsupportedFileTypeError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_FileEmpty) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); std::string file_name = "empty.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingFileEmptyError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_FileTooLarge) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); // Set the maximum file size to 2 bytes. scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_max_size_bytes(2); std::string file_name = "large.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3}; // 3 bytes > 2 bytes limit mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingFileTooLargeError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_UnsupportedType) { base::test::ScopedFeatureList local_feature_list; local_feature_list.InitAndDisableFeature( lens::features::kLensSendRawFileMediaTypes); profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); // Only allow PDF. scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_mime_types_allowed("application/pdf"); scoped_config() .config.mutable_composebox() ->mutable_image_upload() ->set_mime_types_allowed(""); std::string file_name = "test.txt"; std::string mime_type = "text/plain"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingUnsupportedFileTypeError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_MaxImagesExceeded) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); omnibox::InputState state = handler().input_state_model()->GetInputState(); state.max_inputs_by_type[omnibox::INPUT_TYPE_LENS_IMAGE] = 0; handler().input_state_model()->set_state_for_testing(state); scoped_config() .config.mutable_composebox() ->mutable_image_upload() ->set_mime_types_allowed("image/jpeg"); std::string file_name = "test.jpg"; std::string mime_type = "image/jpeg"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingMaxImagesExceededError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_MaxPdfsExceeded) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); omnibox::InputState state = handler().input_state_model()->GetInputState(); state.max_inputs_by_type[omnibox::INPUT_TYPE_LENS_FILE] = 0; handler().input_state_model()->set_state_for_testing(state); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_mime_types_allowed("application/pdf"); std::string file_name = "test.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingMaxPdfsExceededError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_MaxFilesExceeded) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); omnibox::InputState state = handler().input_state_model()->GetInputState(); state.max_inputs_by_type[omnibox::INPUT_TYPE_BROWSER_TAB] = 0; handler().input_state_model()->set_state_for_testing(state); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_mime_types_allowed("tab"); std::string file_name = "test.tab"; std::string mime_type = "tab"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingMaxFilesExceededError); } TEST_F(ContextualSearchboxHandlerTest, AddFileFromBrowser_MaxTotalFilesExceeded) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); omnibox::InputState state = handler().input_state_model()->GetInputState(); state.max_total_inputs = 0; state.max_inputs_by_type[omnibox::INPUT_TYPE_LENS_FILE] = 1; // bypass specific type check handler().input_state_model()->set_state_for_testing(state); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_mime_types_allowed("application/pdf"); std::string file_name = "test.pdf"; std::string mime_type = "application/pdf"; std::vector test_data = {1, 2, 3}; mojo_base::BigBuffer file_data(test_data); base::test::TestFuture> future; handler().AddFileContextFromBrowser(file_name, mime_type, std::move(file_data), std::nullopt, future.GetCallback()); auto result = future.Take(); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), contextual_search::ContextUploadErrorType:: kBrowserProcessingMaxFilesExceededError); } TEST_F(ContextualSearchboxHandlerTest, SubmitQuery) { std::vector session_states; auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .Times(3) .WillRepeatedly([&](SessionState session_state) { session_states.push_back(session_state); }); EXPECT_CALL(*metrics_recorder_ptr, NotifyQuerySubmitted(testing::_, testing::_, testing::_, testing::_, testing::_)) .Times(1) .WillOnce(testing::Invoke( metrics_recorder_ptr, &MockContextualSearchMetricsRecorder::NotifyQuerySubmittedBase)); // Start the session. EXPECT_CALL(query_controller(), InitializeIfNeeded) .Times(1) .WillOnce(testing::Invoke(&query_controller(), &MockQueryController::InitializeIfNeededBase)); // There should be no delayed uploads if there are no cached tab snapshot. EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(0); handler().NotifySessionStarted(); SubmitQueryAndWaitForNavigation(); std::unique_ptr search_url_request_info = std::make_unique< ComposeboxQueryController::CreateSearchUrlRequestInfo>(); search_url_request_info->query_text = kQueryText; search_url_request_info->query_start_time = base::Time::Now(); search_url_request_info->additional_params = { {kUdmQueryParameter, kUdmQueryParameterValue}}; base::test::TestFuture future; query_controller().CreateSearchUrl(std::move(search_url_request_info), future.GetCallback()); GURL expected_url = future.Take(); GURL actual_url = web_contents()->GetController().GetLastCommittedEntry()->GetURL(); // Ensure navigation occurred. EXPECT_EQ(StripTimestampsFromAimUrl(expected_url), StripTimestampsFromAimUrl(actual_url)); EXPECT_THAT(session_states, testing::ElementsAre(SessionState::kSessionStarted, SessionState::kQuerySubmitted, SessionState::kNavigationOccurred)); } TEST_F(ContextualSearchboxHandlerTest, SubmitQuery_VoiceSearch) { bool was_voice_search = false; EXPECT_CALL(query_controller(), CreateSearchUrl) .WillOnce( [&](auto&& request_info, base::OnceCallback callback) { was_voice_search = request_info->is_voice_search; }); handler().SubmitQuery(kQueryText, 1, false, false, false, false, /*is_voice_search=*/true); EXPECT_TRUE(was_voice_search); } TEST_F(ContextualSearchboxHandlerTest, SubmitQuery_DelayUpload) { // Arrange std::vector session_states; auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .Times(3) .WillRepeatedly([&](SessionState session_state) { session_states.push_back(session_state); }); EXPECT_CALL(*metrics_recorder_ptr, NotifyQuerySubmitted(testing::_, testing::_, testing::_, testing::_, testing::_)) .Times(1) .WillOnce(testing::Invoke( metrics_recorder_ptr, &MockContextualSearchMetricsRecorder::NotifyQuerySubmittedBase)); // Start the session. EXPECT_CALL(query_controller(), InitializeIfNeeded) .Times(1) .WillOnce(testing::Invoke(&query_controller(), &MockQueryController::InitializeIfNeededBase)); // Set a cached tab context snapshot. auto token = base::UnguessableToken::Create(); handler() .GetContextualSessionHandle() ->GetUploadedContextTokensForTesting() .push_back(token); handler().tab_context_snapshot_ = std::make_pair(token, std::make_unique()); ASSERT_TRUE(handler().tab_context_snapshot_.has_value()); // The file should be uploaded if there is a cached tab snapshot. EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(1); handler().NotifySessionStarted(); // Act SubmitQueryAndWaitForNavigation(); // Assert // Once the cached tab snapshot is uploaded, it should be cleared. ASSERT_FALSE(handler().tab_context_snapshot_.has_value()); std::unique_ptr search_url_request_info = std::make_unique< ComposeboxQueryController::CreateSearchUrlRequestInfo>(); search_url_request_info->query_text = kQueryText; search_url_request_info->query_start_time = base::Time::Now(); search_url_request_info->file_tokens.push_back(token); search_url_request_info->additional_params = { {kUdmQueryParameter, kUdmQueryParameterValue}}; base::test::TestFuture future; query_controller().CreateSearchUrl(std::move(search_url_request_info), future.GetCallback()); GURL expected_url = future.Take(); GURL actual_url = web_contents()->GetController().GetLastCommittedEntry()->GetURL(); // Ensure navigation occurred. EXPECT_EQ(StripTimestampsFromAimUrl(expected_url), StripTimestampsFromAimUrl(actual_url)); EXPECT_THAT(session_states, testing::ElementsAre(SessionState::kSessionStarted, SessionState::kQuerySubmitted, SessionState::kNavigationOccurred)); } TEST_F(ContextualSearchboxHandlerTest, SubmitQuery_TabAttachmentCount) { // Arrange auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .Times(3); EXPECT_CALL(*metrics_recorder_ptr, NotifyQuerySubmitted(testing::_, testing::_, testing::_, testing::_, testing::_)) .WillOnce(testing::Invoke( metrics_recorder_ptr, &MockContextualSearchMetricsRecorder::NotifyQuerySubmittedBase)); // Start the session. EXPECT_CALL(query_controller(), InitializeIfNeeded) .WillOnce(testing::Invoke(&query_controller(), &MockQueryController::InitializeIfNeededBase)); auto token_tab = base::UnguessableToken::Create(); auto& uploaded_tokens = handler() .GetContextualSessionHandle() ->GetUploadedContextTokensForTesting(); uploaded_tokens.push_back(token_tab); query_controller().AddTabFileInfoForTesting( token_tab, GURL("https://www.google.com"), lens::MimeType::kAnnotatedPageContent); handler().NotifySessionStarted(); // Act SubmitQueryAndWaitForNavigation(); // Assert histogram_tester().ExpectUniqueSample( "ContextualSearch.Query.AttachmentCount.Tab.NewTabPage", 1, 1); } class SmartTabSharingTest : public ContextualSearchboxHandlerTestHarness { public: ~SmartTabSharingTest() override = default; TestingProfile::TestingFactories GetTestingFactories() const override { auto factories = ContextualSearchboxHandlerTestHarness::GetTestingFactories(); factories.emplace_back( feature_engagement::TrackerFactory::GetInstance(), base::BindRepeating([](content::BrowserContext* context) -> std::unique_ptr { auto tracker = std::make_unique< testing::NiceMock>(); ON_CALL(*tracker, IsInFeatureTestMode) .WillByDefault(testing::Return(true)); return tracker; })); return factories; } feature_engagement::test::MockTracker* mock_tracker() { return static_cast( feature_engagement::TrackerFactory::GetForBrowserContext(profile())); } void SetUp() override { ContextualSearchboxHandlerTestHarness::SetUp(); webui::SetBrowserWindowInterface(web_contents(), &mock_browser_window_interface_); ON_CALL(mock_browser_window_interface_, GetUnownedUserDataHost()) .WillByDefault(testing::ReturnRef(unowned_user_data_host_)); ON_CALL(mock_browser_window_interface_, GetWeakPtr()) .WillByDefault(testing::Return(weak_factory_.GetWeakPtr())); feature_list_.InitWithFeaturesAndParameters( {{contextual_tasks::kContextualTasks, {}}, {contextual_tasks::kContextualTasksForceEntryPointEligibility, {}}, {contextual_tasks::kContextualTasksContext, {{"ContextualTasksContextSmartTabSharing", "true"}}}, {contextual_tasks:: kContextualTasksContextSmartTabSharingDefaultOnAvailability, {}}}, {}); auto query_controller_config_params = std::make_unique< contextual_search::ContextualSearchContextController::ConfigParams>(); query_controller_config_params->send_lns_surface = false; query_controller_config_params->enable_viewport_images = true; auto query_controller_ptr = std::make_unique( /*identity_manager=*/nullptr, url_loader_factory(), version_info::Channel::UNKNOWN, "en-US", template_url_service(), fake_variations_client(), std::move(query_controller_config_params)); query_controller_ = query_controller_ptr.get(); auto metrics_recorder_ptr = std::make_unique(); ON_CALL(*metrics_recorder_ptr, RecordZeroSuggestClick) .WillByDefault(testing::Invoke( metrics_recorder_ptr.get(), &MockContextualSearchMetricsRecorder::RecordZeroSuggestClickBase)); service_ = ContextualSearchServiceFactory::GetForProfile(profile()); contextual_session_handle_ = service_->CreateSessionForTesting( std::move(query_controller_ptr), std::move(metrics_recorder_ptr)); contextual_session_handle_->CheckSearchContentSharingSettings( profile()->GetPrefs()); web_contents()->SetDelegate(&delegate_); // Set testing factory BEFORE creating handler. contextual_tasks::ContextualTasksContextServiceFactory::GetInstance() ->SetTestingFactory( profile(), base::BindRepeating( [](SmartTabSharingTest* test, content::BrowserContext* context) -> std::unique_ptr { auto service = std::make_unique( static_cast(context)); test->mock_service_ = service.get(); return service; }, this)); contextual_tasks::ContextualTasksServiceFactory::GetInstance() ->SetTestingFactory( profile(), base::BindRepeating([](content::BrowserContext* context) -> std::unique_ptr { return std::make_unique>(); })); handler_ = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting( [&]() { return contextual_session_handle_.get(); })); // Drain the Mojo pipe and clear setup-related calls to searchbox page. mock_searchbox_page_.FlushForTesting(); base::RunLoop().RunUntilIdle(); testing::Mock::VerifyAndClearExpectations(&mock_searchbox_page_); ON_CALL(query_controller(), CreateSearchUrl) .WillByDefault( [](auto&& request_info, base::OnceCallback callback) { GURL url(base::StrCat({"https://www.google.com/search?q=", request_info->query_text})); url = net::AppendOrReplaceQueryParameter(url, "qsubts", "0"); url = net::AppendOrReplaceQueryParameter(url, "cud", "0"); std::move(callback).Run(url); }); } void SubmitQueryAndWaitForNavigation() { content::TestNavigationObserver navigation_observer(web_contents()); handler().SubmitQuery(kQueryText, 1, false, false, false, false, /*is_voice_search=*/false); auto navigation = content::NavigationSimulator::CreateFromPending( web_contents()->GetController()); ASSERT_TRUE(navigation); auto callback = delegate_.TakeCallback(); if (callback) { std::move(callback).Run(*(navigation->GetNavigationHandle())); } navigation->Commit(); navigation_observer.Wait(); } FakeContextualSearchboxHandler& handler() { return *handler_; } MockQueryController& query_controller() { return *query_controller_; } void TearDown() override { auto* mock_service = static_cast( contextual_tasks::ContextualTasksContextServiceFactory::GetForProfile( profile())); if (mock_service) { testing::Mock::VerifyAndClearExpectations(mock_service); } mock_service_ = nullptr; query_controller_ = nullptr; handler_.reset(); service_ = nullptr; delegate_.ClearCallback(); ContextualSearchboxHandlerTestHarness::TearDown(); } protected: testing::NiceMock mock_searchbox_page_; std::unique_ptr handler_; base::test::ScopedFeatureList feature_list_; raw_ptr mock_service_; std::unique_ptr contextual_session_handle_; TestWebContentsDelegate delegate_; raw_ptr query_controller_; raw_ptr service_; testing::NiceMock mock_browser_window_interface_; ui::UnownedUserDataHost unowned_user_data_host_; base::WeakPtrFactory weak_factory_{ &mock_browser_window_interface_}; }; TEST_F(SmartTabSharingTest, IsSmartTabSharingActive_ReadsPref) { profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, true); EXPECT_TRUE(handler().IsSmartTabSharingActive()); profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, false); EXPECT_FALSE(handler().IsSmartTabSharingActive()); } TEST_F(SmartTabSharingTest, IsSmartTabSharingActive_AvailabilityDisabled) { base::test::ScopedFeatureList feature_list; feature_list.InitAndDisableFeature( contextual_tasks:: kContextualTasksContextSmartTabSharingDefaultOnAvailability); profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, true); EXPECT_FALSE(handler().IsSmartTabSharingActive()); } TEST_F(SmartTabSharingTest, QueryAutocomplete_CallsOnTypedQueryWhenActive) { handler().set_smart_tab_sharing_active_override(true); ASSERT_TRUE(mock_service_); EXPECT_CALL(*mock_service_, OnTypedQuery(testing::_)).Times(1); handler().QueryAutocomplete( 0, /*tab_id=*/std::nullopt, u"test", /*prevent_inline_autocomplete=*/false, 0, omnibox::SuggestInventory::SUGGEST_INVENTORY_DEFAULT, /*is_on_focus=*/false, /*keyword=*/"", searchbox::mojom::InputMethod::kKeyboard); } TEST_F(SmartTabSharingTest, QueryAutocomplete_DoesNotCallOnTypedQueryWhenInactive) { handler().set_smart_tab_sharing_active_override(false); ASSERT_TRUE(mock_service_); EXPECT_CALL(*mock_service_, OnTypedQuery(testing::_)).Times(0); handler().QueryAutocomplete( 0, /*tab_id=*/std::nullopt, u"test", /*prevent_inline_autocomplete=*/false, 0, omnibox::SuggestInventory::SUGGEST_INVENTORY_DEFAULT, /*is_on_focus=*/false, /*keyword=*/"", searchbox::mojom::InputMethod::kKeyboard); } TEST_F(SmartTabSharingTest, SubmitQuery_SmartTabSharingOverrideDisabled) { handler().set_smart_tab_sharing_active_override(false); ASSERT_TRUE(mock_service_); EXPECT_CALL(*mock_service_, GetRelevantTabsForConversationThread(testing::_, testing::_, testing::_, testing::_)) .Times(0); SubmitQueryAndWaitForNavigation(); histogram_tester().ExpectUniqueSample( "ContextualSearch.SmartTabSharing.ThreadWithTabsSubmitted", false, 1); } TEST_F(SmartTabSharingTest, SubmitQuery_SmartTabSharingOverrideEnabledAndActive) { ASSERT_TRUE(mock_service_); handler().set_smart_tab_sharing_active_override(true); EXPECT_CALL(*mock_service_, GetRelevantTabsForConversationThread(testing::_, testing::_, testing::_, testing::_)) .Times(1) .WillOnce([](const auto& options, const auto& conversation_thread, const auto& explicit_urls, auto callback) { ASSERT_FALSE(options.min_model_score.has_value()); std::move(callback).Run({}); }); SubmitQueryAndWaitForNavigation(); histogram_tester().ExpectUniqueSample( "ContextualSearch.SmartTabSharing.ThreadWithTabsSubmitted", true, 1); } TEST_F(SmartTabSharingTest, SetSmartTabSharingActive_FeatureDisabled) { base::test::ScopedFeatureList feature_list; feature_list.InitAndDisableFeature(contextual_tasks::kContextualTasksContext); EXPECT_FALSE(handler().IsSmartTabSharingActive()); handler().SetSmartTabSharingActive(true); EXPECT_FALSE(handler().IsSmartTabSharingActive()); } TEST_F(SmartTabSharingTest, SetSmartTabSharingActive_AvailabilityDisabled) { base::test::ScopedFeatureList feature_list; feature_list.InitAndDisableFeature( contextual_tasks:: kContextualTasksContextSmartTabSharingDefaultOnAvailability); EXPECT_CALL(*mock_tracker(), NotifyEvent("smart_tab_sharing_activated")) .Times(1); EXPECT_CALL(*mock_tracker(), ShouldTriggerHelpUI(testing::_)).Times(0); profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, false); handler().SetSmartTabSharingActive(true); } TEST_F(SmartTabSharingTest, SetSmartTabSharingActive_AvailabilityEnabled_RequestPromo) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature( contextual_tasks:: kContextualTasksContextSmartTabSharingDefaultOnAvailability); EXPECT_CALL(*mock_tracker(), NotifyEvent("smart_tab_sharing_activated")) .Times(1); EXPECT_CALL(*mock_tracker(), ShouldTriggerHelpUI(testing::Ref( feature_engagement::kIPHSmartTabSharingDefaultOnFeature))) .Times(1) .WillOnce(testing::Return(false)); profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, false); handler().SetSmartTabSharingActive(true); } TEST_F(SmartTabSharingTest, SetSmartTabSharingActive_AvailabilityEnabled_StsAlreadyDefaultOn) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature( contextual_tasks:: kContextualTasksContextSmartTabSharingDefaultOnAvailability); EXPECT_CALL(*mock_tracker(), NotifyEvent("smart_tab_sharing_activated")) .Times(1); EXPECT_CALL(*mock_tracker(), ShouldTriggerHelpUI(testing::_)).Times(0); profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, true); handler().SetSmartTabSharingActive(true); } TEST_F(SmartTabSharingTest, GetSmartTabSharingActive) { handler().SetSmartTabSharingActive(true); base::test::TestFuture future; handler().GetSmartTabSharingActive(future.GetCallback()); EXPECT_TRUE(future.Get()); } TEST_F(SmartTabSharingTest, InitializationFromPref) { profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, true); // Recreate handler to test initialization. mock_searchbox_page_.receiver_.reset(); auto handler = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting([&]() { return contextual_session_handle_.get(); })); EXPECT_TRUE(handler->IsSmartTabSharingActive()); } TEST_F(SmartTabSharingTest, FallbackToPrefChanges) { EXPECT_FALSE(handler().IsSmartTabSharingActive()); profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, true); EXPECT_TRUE(handler().IsSmartTabSharingActive()); handler().SetSmartTabSharingActive(false); EXPECT_FALSE(handler().IsSmartTabSharingActive()); // Now that it is overridden, changing the pref should NOT take effect! profile()->GetPrefs()->SetBoolean( contextual_tasks::kContextualTasksShareOpenTabsEveryThread, true); EXPECT_FALSE(handler().IsSmartTabSharingActive()); } TEST_F(SmartTabSharingTest, SubmitQuery_PersistsSmartTabSharingActive) { handler().SetSmartTabSharingActive(true); EXPECT_TRUE(handler().IsSmartTabSharingActive()); ASSERT_TRUE(mock_service_); EXPECT_CALL(*mock_service_, GetRelevantTabsForConversationThread(testing::_, testing::_, testing::_, testing::_)) .Times(1) .WillOnce([](const auto& options, const auto& conversation_thread, const auto& explicit_urls, auto callback) { std::move(callback).Run({}); }); SubmitQueryAndWaitForNavigation(); auto* helper = ContextualSearchWebContentsHelper::FromWebContents(web_contents()); ASSERT_TRUE(helper); auto* new_session_handle = helper->session_handle(); ASSERT_TRUE(new_session_handle); EXPECT_TRUE(new_session_handle->smart_tab_sharing_active().has_value()); EXPECT_TRUE(*new_session_handle->smart_tab_sharing_active()); } TEST_F(SmartTabSharingTest, SubmitQuery_PersistsSmartTabSharingInactive) { handler().SetSmartTabSharingActive(false); EXPECT_FALSE(handler().IsSmartTabSharingActive()); ASSERT_TRUE(mock_service_); EXPECT_CALL(*mock_service_, GetRelevantTabsForConversationThread(testing::_, testing::_, testing::_, testing::_)) .Times(0); SubmitQueryAndWaitForNavigation(); auto* helper = ContextualSearchWebContentsHelper::FromWebContents(web_contents()); ASSERT_TRUE(helper); auto* new_session_handle = helper->session_handle(); ASSERT_TRUE(new_session_handle); EXPECT_TRUE(new_session_handle->smart_tab_sharing_active().has_value()); EXPECT_FALSE(*new_session_handle->smart_tab_sharing_active()); } TEST_F(SmartTabSharingTest, LogMenuOptionClickedMetrics) { ASSERT_TRUE(contextual_tasks::ContextualTasksContextService:: GetIsSmartTabSharingEnabled(profile())); handler().SetSmartTabSharingActive(true); histogram_tester().ExpectUniqueSample( "ContextualSearch.SmartTabSharing.MenuOptionClicked", contextual_tasks::SmartTabSharingToggleState::kToggledOn, 1); handler().SetSmartTabSharingActive(false); histogram_tester().ExpectBucketCount( "ContextualSearch.SmartTabSharing.MenuOptionClicked", contextual_tasks::SmartTabSharingToggleState::kToggledOff, 1); } TEST_F(SmartTabSharingTest, LogOptOutMidThread) { ASSERT_TRUE(contextual_tasks::ContextualTasksContextService:: GetIsSmartTabSharingEnabled(profile())); // Enable STS. handler().SetSmartTabSharingActive(true); histogram_tester().ExpectUniqueSample( "ContextualSearch.SmartTabSharing.MenuOptionClicked", contextual_tasks::SmartTabSharingToggleState::kToggledOn, 1); // Add a turn to session handle. contextual_tasks::ThreadTurn turn; turn.query = "test"; contextual_session_handle_->AddThreadTurn(turn); // Disable STS. handler().SetSmartTabSharingActive(false); // Verify kToggledOff is logged. histogram_tester().ExpectBucketCount( "ContextualSearch.SmartTabSharing.MenuOptionClicked", contextual_tasks::SmartTabSharingToggleState::kToggledOff, 1); // Verify OptOutMidThread is logged. histogram_tester().ExpectUniqueSample( "ContextualSearch.SmartTabSharing.OptOutMidThread", true, 1); } TEST_F(ContextualSearchboxHandlerTest, OnInputStateChanged) { omnibox::InputState received_state_1; omnibox::InputState received_state_2; EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged) .Times(2) .WillOnce( [&](const omnibox::InputState& state) { received_state_1 = state; }) .WillOnce( [&](const omnibox::InputState& state) { received_state_2 = state; }); EXPECT_CALL(*GetMetricsRecorderPtr(), RecordToolMode(omnibox::ToolMode::TOOL_MODE_CANVAS)) .WillOnce(testing::Invoke( GetMetricsRecorderPtr(), &MockContextualSearchMetricsRecorder::RecordToolModeBase)); handler_->SetActiveToolMode(omnibox::ToolMode::TOOL_MODE_CANVAS, /*is_set_by_aim=*/false); handler_->RecordToolSelectionAction(omnibox::ToolMode::TOOL_MODE_CANVAS); mock_searchbox_page_.FlushForTesting(); EXPECT_EQ(received_state_1.active_tool, omnibox::ToolMode::TOOL_MODE_CANVAS); histogram_tester().ExpectUniqueSample("ContextualSearch.Tools.NewTabPage", omnibox::ToolMode::TOOL_MODE_CANVAS, 1); EXPECT_CALL(*GetMetricsRecorderPtr(), RecordModelMode(omnibox::ModelMode::MODEL_MODE_GEMINI_REGULAR)) .WillOnce(testing::Invoke( GetMetricsRecorderPtr(), &MockContextualSearchMetricsRecorder::RecordModelModeBase)); handler().SetActiveModelMode(omnibox::ModelMode::MODEL_MODE_GEMINI_REGULAR, /*is_set_by_aim=*/false); handler().RecordModelSelectionAction( omnibox::ModelMode::MODEL_MODE_GEMINI_REGULAR); mock_searchbox_page_.FlushForTesting(); EXPECT_EQ(received_state_2.active_tool, omnibox::ToolMode::TOOL_MODE_CANVAS); EXPECT_EQ(received_state_2.active_model, omnibox::ModelMode::MODEL_MODE_GEMINI_REGULAR); histogram_tester().ExpectUniqueSample( "ContextualSearch.Models.NewTabPage", omnibox::ModelMode::MODEL_MODE_GEMINI_REGULAR, 1); } TEST_F(ContextualSearchboxHandlerTest, OnDriveUploadClicked_DoubleClick) { EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); auto mock_drive_picker_controller = std::make_unique( profile(), &mock_browser_window_interface_); auto* mock_ptr = mock_drive_picker_controller.get(); handler().SetDrivePickerController(std::move(mock_drive_picker_controller)); // Verify that the controller is only shown once even if clicked multiple // times. EXPECT_CALL(*mock_ptr, ShowDrivePickerHost(testing::_)).Times(1); base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); // Subsequent clicks while the picker is active should be ignored. base::test::TestFuture second_future; handler().OnDriveUploadClicked(second_future.GetCallback()); } TEST_F(ContextualSearchboxHandlerTest, OnDrivePickerDisconnected) { base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); handler().OnDrivePickerDisconnected(); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); EXPECT_TRUE(future.Wait()); } TEST_F(ContextualSearchboxHandlerTest, OnDrivePickerResult_OnSelection) { base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); std::vector files; handler().OnSelection(std::move(files)); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); EXPECT_TRUE(future.Wait()); } TEST_F(ContextualSearchboxHandlerTest, OnDrivePickerResult_OnSelection_InvalidFiles) { base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); // Set max_total_inputs to 5. omnibox::InputState state; state.max_total_inputs = 5; handler().input_state_model()->set_state_for_testing(state); std::vector files; // 1. Valid file auto file1 = drive_picker_host::mojom::DriveFile::New(); file1->id = "valid_id"; file1->name = "valid_name"; file1->mime_type = "text/plain"; file1->type = "document"; file1->size_bytes = 100; files.push_back(std::move(file1)); // 2. Invalid file (invalid ID) auto file2 = drive_picker_host::mojom::DriveFile::New(); file2->id = "invalid id with spaces"; file2->name = "invalid_id_name"; file2->mime_type = "text/plain"; file2->type = "document"; file2->size_bytes = 100; files.push_back(std::move(file2)); // 3. Invalid file (untrusted thumbnail URL) auto file3 = drive_picker_host::mojom::DriveFile::New(); file3->id = "valid_id_3"; file3->name = "untrusted_thumb_name"; file3->mime_type = "image/jpeg"; file3->type = "photo"; file3->size_bytes = 100; file3->thumbnail_url = GURL("https://malicious.com/thumb.jpg"); files.push_back(std::move(file3)); // Expect only file 1 to be processed because processing aborts on invalid // file 2. EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(1); { mojo::FakeMessageDispatchContext context; handler().OnSelection(std::move(files)); } auto response = future.Take(); ASSERT_TRUE(response); // An empty response should be returned if any file is invalid. EXPECT_TRUE(response->files.empty()); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); } TEST_F(ContextualSearchboxHandlerTest, OnDrivePickerResult_OnCancel) { base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); handler().OnCancel(); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); EXPECT_TRUE(future.Wait()); } TEST_F(ContextualSearchboxHandlerTest, OnDrivePickerResult_OnError) { base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); handler().OnError(drive_picker_host::mojom::DrivePickerError::kUnknown); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); EXPECT_TRUE(future.Wait()); } TEST_F(ContextualSearchboxHandlerTest, DriveDisclaimer_NotAccepted) { base::test::ScopedFeatureList feature_list; feature_list.InitWithFeatures({omnibox::kComposeboxDriveContextMenuOption}, {omnibox::kForceDriveDisclaimerAccepted}); SetUpMockFpopService(/*accepted=*/false); base::test::TestFuture future; handler().GetDriveDisclaimerStatus(future.GetCallback()); EXPECT_EQ(searchbox::mojom::DriveDisclaimerStatus::kNotAccepted, future.Get()); EXPECT_EQ( profile()->GetPrefs()->GetInteger(contextual_search::kDriveConsentState), static_cast(contextual_search::DriveConsentState::kNotConsent)); } TEST_F(ContextualSearchboxHandlerTest, DriveDisclaimer_Accepted) { base::test::ScopedFeatureList feature_list; feature_list.InitWithFeatures({omnibox::kComposeboxDriveContextMenuOption}, {omnibox::kForceDriveDisclaimerAccepted}); SetUpMockFpopService(/*accepted=*/true); base::test::TestFuture future; handler().GetDriveDisclaimerStatus(future.GetCallback()); EXPECT_EQ(searchbox::mojom::DriveDisclaimerStatus::kAccepted, future.Get()); EXPECT_EQ( profile()->GetPrefs()->GetInteger(contextual_search::kDriveConsentState), static_cast(contextual_search::DriveConsentState::kConsent)); } TEST_F(ContextualSearchboxHandlerTest, DriveDisclaimer_Restricted) { base::test::ScopedFeatureList feature_list; feature_list.InitWithFeatures({omnibox::kComposeboxDriveContextMenuOption}, {omnibox::kForceDriveDisclaimerAccepted}); SetUpMockFpopServiceWithStatus(drive_picker::DriveDisclaimerController:: ConsentEligibilityStatus::kCannotConsent); base::test::TestFuture future; handler().GetDriveDisclaimerStatus(future.GetCallback()); EXPECT_EQ(searchbox::mojom::DriveDisclaimerStatus::kRestricted, future.Get()); EXPECT_EQ( profile()->GetPrefs()->GetInteger(contextual_search::kDriveConsentState), static_cast(contextual_search::DriveConsentState::kRestricted)); } TEST_F(ContextualSearchboxHandlerTest, DriveDisclaimer_NoEagerRpcOnInitialize) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kComposeboxDriveContextMenuOption); auto mock_fpop_service = std::make_unique(); EXPECT_CALL(*mock_fpop_service, ShouldShowMobileConsentFlow(testing::_, testing::_)) .Times(0); handler().SetDriveDisclaimerController( std::make_unique( std::move(mock_fpop_service))); // Initializing the input state model (as done on page load) should NOT // trigger the RPC. handler().InitializeInputStateModel(); } TEST_F(ContextualSearchboxHandlerTest, DriveDisclaimer_FlagDisabled) { base::test::ScopedFeatureList feature_list; feature_list.InitAndDisableFeature( omnibox::kComposeboxDriveContextMenuOption); base::test::TestFuture future; handler().GetDriveDisclaimerStatus(future.GetCallback()); EXPECT_EQ(searchbox::mojom::DriveDisclaimerStatus::kRestricted, future.Get()); } TEST_F(ContextualSearchboxHandlerTest, OnDriveUploadClicked) { EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); omnibox::InputState state; state.max_total_inputs = 10; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); auto mock_drive_picker_controller = std::make_unique( profile(), &mock_browser_window_interface_); auto* mock_ptr = mock_drive_picker_controller.get(); handler().SetDrivePickerController(std::move(mock_drive_picker_controller)); // The call should bind the receiver and show the picker. mojo::PendingRemote pending_remote; EXPECT_CALL(*mock_ptr, ShowDrivePickerHost(testing::_)) .WillOnce( [&](std::unique_ptr request) { pending_remote = request->TakeResultHandler(); }); base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); // Set max_total_inputs to 1. omnibox::InputState state_1; state_1.max_total_inputs = 1; handler().input_state_model()->set_state_for_testing(state_1); // Simulate user selecting a file. EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::_, testing::_)) .WillOnce([&](const base::UnguessableToken& token, std::unique_ptr input_data, std::optional image_options) { EXPECT_EQ(input_data->drive_id, "mock_id"); EXPECT_EQ(input_data->resource_key, "mock_resource_key"); EXPECT_EQ(input_data->mime_type_string, "text/plain"); }); std::vector files; auto file = drive_picker_host::mojom::DriveFile::New(); file->id = "mock_id"; file->name = "mock_name"; file->mime_type = "text/plain"; file->type = "document"; file->size_bytes = 100; file->resource_key = "mock_resource_key"; files.push_back(std::move(file)); handler().OnSelection(std::move(files)); // The callback should be run with the selected file. auto response = future.Take(); ASSERT_TRUE(response); ASSERT_EQ(response->files.size(), 1u); EXPECT_EQ(response->files[0]->file_name, "mock_name"); EXPECT_FALSE(response->error.has_value()); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); EXPECT_FALSE(handler().has_drive_picker_deactivation_blocker_for_testing()); } TEST_F(ContextualSearchboxHandlerTest, OnDrivePickerResult_OnSelection_MaxTotalInputsZero) { // Set max_total_inputs to 0. omnibox::InputState state; state.max_total_inputs = 0; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); EXPECT_TRUE(handler().IsDrivePickerReceiverBound()); std::vector files; auto file = drive_picker_host::mojom::DriveFile::New(); file->id = "mock_id"; file->name = "mock_name"; file->mime_type = "text/plain"; file->type = "document"; file->size_bytes = 100; file->resource_key = "mock_resource_key"; files.push_back(std::move(file)); // If max_total_inputs is 0, OnSelection should return an empty list of files. EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(0); handler().OnSelection(std::move(files)); auto response = future.Take(); ASSERT_TRUE(response); EXPECT_EQ(0u, response->files.size()); EXPECT_FALSE(response->error.has_value()); EXPECT_FALSE(handler().IsDrivePickerReceiverBound()); } // TODO(crbug.com/508693783): Update these tests once the Drive file upload flow // is implemented. TEST_F(ContextualSearchboxHandlerTest, OnDriveUploadClicked_SizeLimitExceeded) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kComposeboxDriveContextMenuOption); omnibox::InputState state; state.max_total_inputs = 10; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); auto mock_drive_picker_controller = std::make_unique( profile(), &mock_browser_window_interface_); auto* mock_ptr = mock_drive_picker_controller.get(); handler().SetDrivePickerController(std::move(mock_drive_picker_controller)); EXPECT_CALL(*mock_ptr, ShowDrivePickerHost(testing::_)); base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); std::vector files; // 1. Valid file auto file1 = drive_picker_host::mojom::DriveFile::New(); file1->id = "id1"; file1->name = "name1"; file1->mime_type = "text/plain"; file1->type = "document"; file1->size_bytes = 100; files.push_back(std::move(file1)); // 2. Large file (exceeds size limit) auto file2 = drive_picker_host::mojom::DriveFile::New(); file2->id = "id2"; file2->name = "name2"; file2->mime_type = "text/plain"; file2->type = "document"; file2->size_bytes = 101 * 1000 * 1000; // 101MB files.push_back(std::move(file2)); EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(1); handler().OnSelection(std::move(files)); auto response = future.Take(); ASSERT_TRUE(response); EXPECT_EQ(response->files.size(), 1u); EXPECT_EQ(response->error, searchbox::mojom::DriveUploadError::kSizeLimitExceeded); } TEST_F(ContextualSearchboxHandlerTest, OnDriveUploadClicked_MaxFilesExceeded) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kComposeboxDriveContextMenuOption); omnibox::InputState state; state.max_total_inputs = 1; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); auto mock_drive_picker_controller = std::make_unique( profile(), &mock_browser_window_interface_); auto* mock_ptr = mock_drive_picker_controller.get(); handler().SetDrivePickerController(std::move(mock_drive_picker_controller)); EXPECT_CALL(*mock_ptr, ShowDrivePickerHost(testing::_)); base::test::TestFuture future; handler().OnDriveUploadClicked(future.GetCallback()); std::vector files; auto file1 = drive_picker_host::mojom::DriveFile::New(); file1->id = "id1"; file1->name = "name1"; file1->mime_type = "text/plain"; file1->type = "document"; file1->size_bytes = 100; files.push_back(std::move(file1)); auto file2 = drive_picker_host::mojom::DriveFile::New(); file2->id = "id2"; file2->name = "name2"; file2->mime_type = "text/plain"; file2->type = "document"; file2->size_bytes = 100; files.push_back(std::move(file2)); EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(1); handler().OnSelection(std::move(files)); auto response = future.Take(); ASSERT_TRUE(response); EXPECT_EQ(response->files.size(), 1u); EXPECT_EQ(response->error, searchbox::mojom::DriveUploadError::kMaxFilesExceeded); } TEST_F(ContextualSearchboxHandlerTest, OpenAutocompleteMatch_ZeroSuggestClick) { base::UserActionTester user_action_tester; // Set up a zero-suggest input. AutocompleteInput input(std::u16string(), metrics::OmniboxEventProto::NTP_OMNIBOX_COMPOSEBOX, ChromeAutocompleteSchemeClassifier(profile())); input.set_focus_type(metrics::OmniboxFocusType::INTERACTION_FOCUS); // Set the page classification on the client's location bar model. static_cast(handler().client()) ->location_bar_model() ->set_page_classification( metrics::OmniboxEventProto::NTP_OMNIBOX_COMPOSEBOX); // 1. Test normal zero-suggest click. { // Use FakeAutocompleteController to easily set input and results. auto fake_controller = std::make_unique(nullptr); fake_controller->input_ = input; AutocompleteMatch match; match.provider = &fake_controller->GetFakeProvider(); match.destination_url = GURL("https://www.google.com"); fake_controller->published_result_.AppendMatches({match}); handler().SetAutocompleteControllerForTesting(std::move(fake_controller)); EXPECT_CALL(*GetMetricsRecorderPtr(), RecordZeroSuggestClick(false)) .WillOnce(testing::Invoke( GetMetricsRecorderPtr(), &MockContextualSearchMetricsRecorder::RecordZeroSuggestClickBase)); auto modifiers = searchbox::mojom::ActionModifiers::New(); handler().OpenAutocompleteMatch(0, GURL("https://www.google.com"), /*are_matches_showing=*/true, /*mouse_button=*/0, std::move(modifiers), /*via_keyboard=*/false); histogram_tester().ExpectBucketCount( "ContextualSearch.ZeroSuggestClickV2.IsContextual.NewTabPage", false, 1); EXPECT_EQ( user_action_tester.GetActionCount( "ContextualSearch.ZeroSuggestClickV2.NonContextual.NewTabPage"), 1); } // 2. Test contextual zero-suggest click. { // Need a new controller since we moved it. auto fake_controller = std::make_unique(nullptr); fake_controller->input_ = input; AutocompleteMatch match; match.provider = &fake_controller->GetFakeProvider(); match.destination_url = GURL("https://www.contextual.com"); match.subtypes.insert(omnibox::SuggestSubtype::SUBTYPE_CONTEXTUAL_SEARCH); fake_controller->published_result_.AppendMatches({match}); handler().SetAutocompleteControllerForTesting(std::move(fake_controller)); EXPECT_CALL(*GetMetricsRecorderPtr(), RecordZeroSuggestClick(true)) .WillOnce(testing::Invoke( GetMetricsRecorderPtr(), &MockContextualSearchMetricsRecorder::RecordZeroSuggestClickBase)); auto modifiers = searchbox::mojom::ActionModifiers::New(); handler().OpenAutocompleteMatch(0, GURL("https://www.contextual.com"), /*are_matches_showing=*/true, /*mouse_button=*/0, std::move(modifiers), /*via_keyboard=*/false); histogram_tester().ExpectBucketCount( "ContextualSearch.ZeroSuggestClickV2.IsContextual.NewTabPage", true, 1); EXPECT_EQ(user_action_tester.GetActionCount( "ContextualSearch.ZeroSuggestClickV2.Contextual.NewTabPage"), 1); } } TEST_F(ContextualSearchboxHandlerTest, OpenAutocompleteMatch_TypedSuggestNavigation) { base::UserActionTester user_action_tester; // Set up a typed input (non-zero suggest). AutocompleteInput input(u"test", metrics::OmniboxEventProto::NTP_OMNIBOX_COMPOSEBOX, ChromeAutocompleteSchemeClassifier(profile())); // Set the page classification on the client's location bar model. static_cast(handler().client()) ->location_bar_model() ->set_page_classification( metrics::OmniboxEventProto::NTP_OMNIBOX_COMPOSEBOX); // 1. Test verbatim typed click (line == 0). { auto fake_controller = std::make_unique(nullptr); fake_controller->input_ = input; AutocompleteMatch match; match.provider = &fake_controller->GetFakeProvider(); match.destination_url = GURL("https://www.google.com"); match.type = AutocompleteMatchType::SEARCH_WHAT_YOU_TYPED; fake_controller->published_result_.AppendMatches({match}); handler().SetAutocompleteControllerForTesting(std::move(fake_controller)); EXPECT_CALL(*GetMetricsRecorderPtr(), RecordTypedSuggestNavigation(true)) .WillOnce(testing::Invoke(GetMetricsRecorderPtr(), &MockContextualSearchMetricsRecorder:: RecordTypedSuggestNavigationBase)); auto modifiers = searchbox::mojom::ActionModifiers::New(); handler().OpenAutocompleteMatch(0, GURL("https://www.google.com"), /*are_matches_showing=*/true, /*mouse_button=*/0, std::move(modifiers), /*via_keyboard=*/false); histogram_tester().ExpectBucketCount( "ContextualSearch.TypedSuggestNavigation.IsVerbatim.NewTabPage", true, 1); EXPECT_EQ( user_action_tester.GetActionCount( "ContextualSearch.TypedSuggestNavigation.Verbatim.NewTabPage"), 1); } // 2. Test non-verbatim typed click (line != 0). { auto fake_controller = std::make_unique(nullptr); fake_controller->input_ = input; AutocompleteMatch match0; match0.provider = &fake_controller->GetFakeProvider(); match0.destination_url = GURL("https://www.google.com"); match0.type = AutocompleteMatchType::SEARCH_WHAT_YOU_TYPED; AutocompleteMatch match1; match1.provider = &fake_controller->GetFakeProvider(); match1.destination_url = GURL("https://www.google.com/search?q=suggestion"); match1.type = AutocompleteMatchType::SEARCH_SUGGEST; fake_controller->published_result_.AppendMatches({match0, match1}); handler().SetAutocompleteControllerForTesting(std::move(fake_controller)); EXPECT_CALL(*GetMetricsRecorderPtr(), RecordTypedSuggestNavigation(false)) .WillOnce(testing::Invoke(GetMetricsRecorderPtr(), &MockContextualSearchMetricsRecorder:: RecordTypedSuggestNavigationBase)); auto modifiers = searchbox::mojom::ActionModifiers::New(); handler().OpenAutocompleteMatch( 1, GURL("https://www.google.com/search?q=suggestion"), /*are_matches_showing=*/true, /*mouse_button=*/0, std::move(modifiers), /*via_keyboard=*/false); histogram_tester().ExpectBucketCount( "ContextualSearch.TypedSuggestNavigation.IsVerbatim.NewTabPage", false, 1); EXPECT_EQ( user_action_tester.GetActionCount( "ContextualSearch.TypedSuggestNavigation.SearchSuggest.NewTabPage"), 1); } } TEST_F(ContextualSearchboxHandlerTest, SubmitQueryWithAdditionalParams) { // Ensure udm param is always set as an additional param. SubmitQueryAndWaitForNavigation(); GURL query_url = web_contents()->GetController().GetLastCommittedEntry()->GetURL(); std::string udm_param; EXPECT_TRUE(net::GetValueForKeyInQuery(query_url, "udm", &udm_param)); EXPECT_EQ("50", udm_param); } TEST_F(ContextualSearchboxHandlerTest, SubmitQuery_NoContextualTasksService) { // Force the ContextualTasksService to be null. contextual_tasks::ContextualTasksServiceFactory::GetInstance() ->SetTestingFactory( profile(), base::BindRepeating([](content::BrowserContext* context) -> std::unique_ptr { return nullptr; })); // Recreate handler to test initialization with a null service. mock_searchbox_page_.receiver_.reset(); auto handler_without_service = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting( [&]() { return contextual_session_handle_.get(); })); content::TestNavigationObserver navigation_observer(web_contents()); handler_without_service->SubmitQuery(kQueryText, 1, false, false, false, false, /*is_voice_search=*/false); auto navigation = content::NavigationSimulator::CreateFromPending( web_contents()->GetController()); ASSERT_TRUE(navigation); navigation->Commit(); navigation_observer.Wait(); GURL query_url = web_contents()->GetController().GetLastCommittedEntry()->GetURL(); EXPECT_TRUE(query_url.spec().find(kQueryText) != std::string::npos); } TEST_F(ContextualSearchboxHandlerTest, QueryAutocomplete_SetsLensInputs) { // Set suggest inputs on the client. lens::proto::LensOverlaySuggestInputs suggest_inputs; suggest_inputs.set_encoded_image_signals("xyz"); EXPECT_CALL(*static_cast(handler().client()), GetLensOverlaySuggestInputs()) .WillRepeatedly(testing::Return(suggest_inputs)); // Set input state where `image_gen_upload_active` is false and active_tool is // not CANVAS. omnibox::InputState state; state.image_gen_upload_active = false; state.active_tool = omnibox::ToolMode::TOOL_MODE_UNSPECIFIED; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); // Set mock AutocompleteController. auto autocomplete_controller = std::make_unique>( std::make_unique(), 0); AutocompleteInput input; EXPECT_CALL(*autocomplete_controller, Start(_)) .WillOnce(testing::SaveArg<0>(&input)); handler().SetAutocompleteControllerForTesting( std::move(autocomplete_controller)); handler().QueryAutocomplete( 0, /*tab_id=*/std::nullopt, u"test", /*prevent_inline_autocomplete=*/false, 0, omnibox::SuggestInventory::SUGGEST_INVENTORY_DEFAULT, /*is_on_focus=*/false, /*keyword=*/"", searchbox::mojom::InputMethod::kKeyboard); EXPECT_TRUE(input.lens_overlay_suggest_inputs().has_value()); EXPECT_EQ(input.lens_overlay_suggest_inputs()->encoded_image_signals(), "xyz"); } TEST_F(ContextualSearchboxHandlerTest, QueryAutocomplete_NullBrowserWindowInterfaceDoesNotCrash) { // Ensure BrowserWindowInterface is null for web_contents. webui::SetBrowserWindowInterface(web_contents(), nullptr); auto autocomplete_controller = std::make_unique>( std::make_unique(), 0); EXPECT_CALL(*autocomplete_controller, Start(_)); handler().SetAutocompleteControllerForTesting( std::move(autocomplete_controller)); // Should execute cleanly without crashing even when BrowserWindowInterface // is null. handler().QueryAutocomplete( 0, /*tab_id=*/std::nullopt, u"test", /*prevent_inline_autocomplete=*/false, 0, omnibox::SuggestInventory::SUGGEST_INVENTORY_DEFAULT, /*is_on_focus=*/false, /*keyword=*/"", searchbox::mojom::InputMethod::kKeyboard); } TEST_F(ContextualSearchboxHandlerTest, QueryAutocomplete_SetsLensInputs_InToolModes) { lens::proto::LensOverlaySuggestInputs suggest_inputs; suggest_inputs.set_encoded_image_signals("xyz"); EXPECT_CALL(*static_cast(handler().client()), GetLensOverlaySuggestInputs()) .WillRepeatedly(testing::Return(suggest_inputs)); // 1. Case: `image_gen_upload_active = true`. { omnibox::InputState state; state.image_gen_upload_active = true; state.active_tool = omnibox::ToolMode::TOOL_MODE_UNSPECIFIED; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); auto autocomplete_controller = std::make_unique>( std::make_unique(), 0); AutocompleteInput input; EXPECT_CALL(*autocomplete_controller, Start(_)) .WillOnce(testing::SaveArg<0>(&input)); handler().SetAutocompleteControllerForTesting( std::move(autocomplete_controller)); handler().QueryAutocomplete( 0, /*tab_id=*/std::nullopt, u"test", /*prevent_inline_autocomplete=*/false, 0, omnibox::SuggestInventory::SUGGEST_INVENTORY_DEFAULT, /*is_on_focus=*/false, /*keyword=*/"", searchbox::mojom::InputMethod::kKeyboard); EXPECT_TRUE(input.lens_overlay_suggest_inputs().has_value()); EXPECT_EQ(input.lens_overlay_suggest_inputs()->encoded_image_signals(), "xyz"); } // 2. Case: `active_tool = TOOL_MODE_CANVAS`. { omnibox::InputState state; state.image_gen_upload_active = false; state.active_tool = omnibox::ToolMode::TOOL_MODE_CANVAS; handler().input_state_model()->set_state_for_testing(state); handler().OnInputStateChangedForTesting(state); auto autocomplete_controller = std::make_unique>( std::make_unique(), 0); AutocompleteInput input; EXPECT_CALL(*autocomplete_controller, Start(_)) .WillOnce(testing::SaveArg<0>(&input)); handler().SetAutocompleteControllerForTesting( std::move(autocomplete_controller)); handler().QueryAutocomplete( 0, /*tab_id=*/std::nullopt, u"test", /*prevent_inline_autocomplete=*/false, 0, omnibox::SuggestInventory::SUGGEST_INVENTORY_DEFAULT, /*is_on_focus=*/false, /*keyword=*/"", searchbox::mojom::InputMethod::kKeyboard); EXPECT_TRUE(input.lens_overlay_suggest_inputs().has_value()); EXPECT_EQ(input.lens_overlay_suggest_inputs()->encoded_image_signals(), "xyz"); } } class ContextualSearchboxHandlerTestTabsTest : public ContextualSearchboxHandlerTest { public: ContextualSearchboxHandlerTestTabsTest() = default; ~ContextualSearchboxHandlerTestTabsTest() override = default; void SetUp() override { ContextualSearchboxHandlerTest::SetUp(); last_active_time_ticks_ = base::TimeTicks::Now(); contextual_tasks::ContextualTasksServiceFactory::GetInstance() ->SetTestingFactory( profile(), base::BindRepeating([](content::BrowserContext* context) -> std::unique_ptr { return std::make_unique>(); })); tab_list_ = std::make_unique>(); tab_list_registration_ = std::make_unique>( user_data_host_, *tab_list_); ON_CALL(browser_window_interface_, GetUnownedUserDataHost) .WillByDefault(::testing::ReturnRef(user_data_host_)); webui::SetBrowserWindowInterface(web_contents(), &browser_window_interface_); } void TearDown() override { // Clear TabContextualizationController to avoid dangling pointers. for (tabs::TabInterface* tab : all_tabs_) { if (tab && tab->GetTabFeatures()) { tab->GetTabFeatures()->SetTabContextualizationControllerForTesting( nullptr); } } handler_.reset(); tab_list_registration_.reset(); tab_list_.reset(); all_tabs_.clear(); owned_tabs_.clear(); ContextualSearchboxHandlerTest::TearDown(); } base::TimeTicks IncrementTimeTicksAndGet() { last_active_time_ticks_ += base::Seconds(1); return last_active_time_ticks_; } MockTabListInterface* tab_list() { return tab_list_.get(); } MockBrowserWindowInterface* browser_window_interface() { return &browser_window_interface_; } std::unique_ptr CreateWebContents() { return content::WebContentsTester::CreateTestWebContents(profile(), nullptr); } tabs::TabInterface* AddTab(GURL url) { auto contents_unique_ptr = CreateWebContents(); content::WebContentsTester::For(contents_unique_ptr.get()) ->NavigateAndCommit(url); content::WebContents* content_ptr = contents_unique_ptr.get(); content::WebContentsTester::For(content_ptr) ->SetLastActiveTimeTicks(IncrementTimeTicksAndGet()); auto tab_model = std::make_unique( std::move(contents_unique_ptr), nullptr); tabs::TabInterface* tab_interface = tab_model.get(); tabs::TabFeatures* const tab_features = tab_interface->GetTabFeatures(); std::unique_ptr tab_ui_helper = tabs::TabFeatures::GetUserDataFactoryForTesting() .CreateInstance(*tab_interface, *tab_interface); tab_features->SetTabUIHelperForTesting(std::move(tab_ui_helper)); std::unique_ptr tab_contextualization_controller = tabs::TabFeatures::GetUserDataFactoryForTesting() .CreateInstance( *tab_interface, tab_interface); tab_features->SetTabContextualizationControllerForTesting( std::move(tab_contextualization_controller)); ON_CALL(*static_cast( lens::TabContextualizationController::From(tab_interface)), GetInitialPageContextEligibility()) .WillByDefault(testing::Return(true)); std::unique_ptr tab_alert_controller = tabs::TabFeatures::GetUserDataFactoryForTesting() .CreateInstance(*tab_interface, *tab_interface); tab_features->SetTabAlertControllerForTesting( std::move(tab_alert_controller)); owned_tabs_.push_back(std::move(tab_model)); all_tabs_.push_back(tab_interface); ON_CALL(*tab_list_, GetAllTabs()).WillByDefault(testing::Return(all_tabs_)); ON_CALL(*tab_list_, GetTabCount()) .WillByDefault(testing::Return(static_cast(all_tabs_.size()))); ON_CALL(*tab_list_, GetIndexOfTab(tab_interface->GetHandle())) .WillByDefault(testing::Return(static_cast(all_tabs_.size()) - 1)); ON_CALL(*tab_list_, GetActiveTab()) .WillByDefault(testing::Return(tab_interface)); return tab_interface; } protected: base::TimeTicks last_active_time_ticks_; std::vector> owned_tabs_; std::vector all_tabs_; std::unique_ptr tab_list_; std::unique_ptr> tab_list_registration_; ui::UnownedUserDataHost user_data_host_; MockBrowserWindowInterface browser_window_interface_; base::HistogramTester histogram_tester_; const tabs::TabModel::PreventFeatureInitializationForTesting prevent_; }; TEST_F(ContextualSearchboxHandlerTestTabsTest, AddTabContext) { auto sample_url = GURL("https://www.google.com"); tabs::TabInterface* tab = AddTab(sample_url); const int sample_tab_id = tab->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*tab_contextualization_controller, GetPageContext(testing::_)) .Times(1) .WillRepeatedly( [](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run( std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(1); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); base::MockCallback callback; EXPECT_CALL(callback, Run).Times(1); auto sample_contextual_input_data = std::make_unique(); sample_contextual_input_data->page_url = sample_url; handler().AddTabContext(sample_tab_id, /*delay_upload=*/false, searchbox::mojom::TabAttachmentSource::kContextMenu, callback.Get()); // Flush the mojo pipe to ensure the callback is run. mock_searchbox_page_.FlushForTesting(); histogram_tester_.ExpectUniqueSample( "ContextualSearch.AttachmentButtonUsed.NewTabPage", contextual_search::ContextualSearchAttachmentButtonType::kCurrentTab, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, ClearFiles_KeepTabs) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kContextManagementInComposebox); // Add a file context: searchbox::mojom::SelectedFileInfoPtr file_info = searchbox::mojom::SelectedFileInfo::New(); file_info->file_name = "test.png"; file_info->selection_time = base::Time::Now(); file_info->mime_type = "application/image"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); base::UnguessableToken file_token; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, auto, auto) { file_token = token; query_controller().AddFileInfoForTesting(token, lens::MimeType::kImage); }); base::MockCallback file_callback; EXPECT_CALL(file_callback, Run).Times(1); handler().AddFileContext(std::move(file_info), std::move(file_data), file_callback.Get()); // Add a tab context: auto sample_url = GURL("https://www.google.com"); tabs::TabInterface* tab = AddTab(sample_url); const int sample_tab_id = tab->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*tab_contextualization_controller, GetPageContext(testing::_)) .Times(1) .WillRepeatedly( [](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run( std::make_unique()); }); base::UnguessableToken tab_token; EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .WillOnce([&](const base::UnguessableToken& token, auto, auto) { tab_token = token; query_controller().AddTabFileInfoForTesting( token, sample_url, lens::MimeType::kAnnotatedPageContent, SessionID::FromSerializedValue(sample_tab_id)); }); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(2); base::MockCallback tab_callback; EXPECT_CALL(tab_callback, Run).Times(1); handler().AddTabContext(sample_tab_id, /*delay_upload=*/false, searchbox::mojom::TabAttachmentSource::kContextMenu, tab_callback.Get()); // Verify both tokens are uploaded: EXPECT_EQ(handler().GetUploadedContextTokens().size(), 2u); handler().ClearFiles(/*should_block_auto_suggested_tabs=*/false, /*query_submitted=*/true); // Verify all uploaded tokens are cleared: EXPECT_EQ(handler().GetUploadedContextTokens().size(), 0u); // Verify tab token remains in submitted tabs: const auto& submitted_tabs = contextual_session_handle_->persisted_tabs(); EXPECT_EQ(submitted_tabs.size(), 1u); auto it = submitted_tabs.find(SessionID::FromSerializedValue(sample_tab_id)); ASSERT_NE(it, submitted_tabs.end()); EXPECT_EQ(it->second.first, tab_token); } TEST_F(ContextualSearchboxHandlerTestTabsTest, ClearFiles_DoNotKeepTabsIfFlagDisabled) { base::test::ScopedFeatureList feature_list; feature_list.InitAndDisableFeature(omnibox::kContextManagementInComposebox); // Add a file context: searchbox::mojom::SelectedFileInfoPtr file_info = searchbox::mojom::SelectedFileInfo::New(); file_info->file_name = "test.png"; file_info->selection_time = base::Time::Now(); file_info->mime_type = "application/image"; std::vector test_data = {1, 2, 3, 4}; auto test_data_span = base::span(test_data); mojo_base::BigBuffer file_data(test_data_span); base::UnguessableToken file_token; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, auto, auto) { file_token = token; query_controller().AddFileInfoForTesting(token, lens::MimeType::kImage); }); base::MockCallback file_callback; EXPECT_CALL(file_callback, Run).Times(1); handler().AddFileContext(std::move(file_info), std::move(file_data), file_callback.Get()); // Add a tab context: auto sample_url = GURL("https://www.google.com"); tabs::TabInterface* tab = AddTab(sample_url); const int sample_tab_id = tab->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*tab_contextualization_controller, GetPageContext(testing::_)) .Times(1) .WillRepeatedly( [](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run( std::make_unique()); }); base::UnguessableToken tab_token; EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .WillOnce([&](const base::UnguessableToken& token, auto, auto) { tab_token = token; query_controller().AddTabFileInfoForTesting( token, sample_url, lens::MimeType::kAnnotatedPageContent, SessionID::FromSerializedValue(sample_tab_id)); }); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(2); base::MockCallback tab_callback; EXPECT_CALL(tab_callback, Run).Times(1); handler().AddTabContext(sample_tab_id, /*delay_upload=*/false, searchbox::mojom::TabAttachmentSource::kContextMenu, tab_callback.Get()); // Verify both tokens are uploaded: EXPECT_EQ(handler().GetUploadedContextTokens().size(), 2u); handler().ClearFiles(/*should_block_auto_suggested_tabs=*/false, /*query_submitted=*/true); // Verify no tokens remain since the feature is disabled: auto remaining_tokens = handler().GetUploadedContextTokens(); EXPECT_EQ(remaining_tokens.size(), 0u); } TEST_F(ContextualSearchboxHandlerTestTabsTest, AddTabContextNotFound) { base::MockCallback callback; base::expected callback_result = base::ok(base::UnguessableToken()); // Initialize with dummy EXPECT_CALL(callback, Run).WillOnce(testing::SaveArg<0>(&callback_result)); handler().AddTabContext(0, false, searchbox::mojom::TabAttachmentSource::kContextMenu, callback.Get()); // Flush the mojo pipe to ensure the callback is run. mock_searchbox_page_.FlushForTesting(); EXPECT_FALSE(callback_result.has_value()); EXPECT_EQ(callback_result.error(), contextual_search::ContextUploadErrorType::kBrowserProcessingError); } TEST_F(ContextualSearchboxHandlerTestTabsTest, AddTabContext_PolicyDisabled) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kDisabled)); auto sample_url = GURL("https://www.google.com"); tabs::TabInterface* tab = AddTab(sample_url); const int sample_tab_id = tab->GetHandle().raw_value(); EXPECT_CALL(query_controller(), StartFileUploadFlow).Times(0); base::MockCallback callback; base::expected callback_result = base::ok(base::UnguessableToken()); // Initialize with dummy EXPECT_CALL(callback, Run).WillOnce(testing::SaveArg<0>(&callback_result)); handler().AddTabContext(sample_tab_id, /*delay_upload=*/false, searchbox::mojom::TabAttachmentSource::kContextMenu, callback.Get()); // Flush the mojo pipe to ensure the callback is run. mock_searchbox_page_.FlushForTesting(); EXPECT_FALSE(callback_result.has_value()); EXPECT_EQ(callback_result.error(), contextual_search::ContextUploadErrorType::kBrowserProcessingError); } TEST_F(ContextualSearchboxHandlerTestTabsTest, AddTabContext_DelayUpload) { // Arrange auto sample_url = GURL("https://www.google.com"); tabs::TabInterface* tab = AddTab(sample_url); const int sample_tab_id = tab->GetHandle().raw_value(); contextual_search::ContextUploadStatus status; MockTabContextualizationController* tab_contextualization_controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*tab_contextualization_controller, GetPageContext(testing::_)) .Times(1) .WillRepeatedly( [](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run( std::make_unique()); }); // `ComposeboxQueryController::StartFileUploadFlow` should not be called when // tab context is added, upload should be delayed, and the tab context // should be cached. ASSERT_FALSE(handler().tab_context_snapshot_.has_value()); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(0); EXPECT_CALL(mock_searchbox_page_, OnContextualInputStatusChanged) .Times(1) .WillOnce( [&status]( const base::UnguessableToken& context_token, contextual_search::ContextUploadStatus context_upload_status, std::optional context_upload_error_type) { status = context_upload_status; }); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); base::MockCallback callback; EXPECT_CALL(callback, Run).Times(1); auto sample_contextual_input_data = std::make_unique(); sample_contextual_input_data->page_url = sample_url; // Act handler().AddTabContext(sample_tab_id, /*delay_upload=*/true, searchbox::mojom::TabAttachmentSource::kContextMenu, callback.Get()); // Flush the mojo pipe to ensure the callback is run and captures the status. mock_searchbox_page_.FlushForTesting(); // Assert ASSERT_TRUE(handler().tab_context_snapshot_.has_value()); ASSERT_TRUE(handler().context_input_data().has_value()); ASSERT_EQ(contextual_search::ContextUploadStatus::kProcessing, status); histogram_tester_.ExpectUniqueSample( "ContextualSearch.AttachmentButtonUsed.NewTabPage", contextual_search::ContextualSearchAttachmentButtonType::kSuggestedTab, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, AddTabContext_RecentTab) { auto sample_url_1 = GURL("https://www.google.com"); auto sample_url_2 = GURL("https://www.youtube.com"); tabs::TabInterface* tab1 = AddTab(sample_url_1); AddTab(sample_url_2); const int sample_tab_id_1 = tab1->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller = static_cast( lens::TabContextualizationController::From(tab1)); EXPECT_CALL(*tab_contextualization_controller, GetPageContext(testing::_)) .WillOnce( [](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run( std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged); base::MockCallback callback; EXPECT_CALL(callback, Run); handler().AddTabContext(sample_tab_id_1, /*delay_upload=*/false, searchbox::mojom::TabAttachmentSource::kContextMenu, callback.Get()); // Flush the mojo pipe to ensure the callback is run. mock_searchbox_page_.FlushForTesting(); histogram_tester_.ExpectUniqueSample( "ContextualSearch.AttachmentButtonUsed.NewTabPage", contextual_search::ContextualSearchAttachmentButtonType::kRecentTab, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, DeleteContext_DelayUpload) { // Arrange auto sample_url = GURL("https://www.google.com"); tabs::TabInterface* tab = AddTab(sample_url); const int sample_tab_id = tab->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*tab_contextualization_controller, GetPageContext(testing::_)) .Times(1) .WillRepeatedly( [](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run( std::make_unique()); }); // `ComposeboxQueryController::StartFileUploadFlow` should not be called when // tab context is added and upload should be delayed. EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(0); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(2); base::test::TestFuture> future; auto sample_contextual_input_data = std::make_unique(); sample_contextual_input_data->page_url = sample_url; handler().AddTabContext(sample_tab_id, /*delay_upload=*/true, searchbox::mojom::TabAttachmentSource::kContextMenu, future.GetCallback()); // Flush the mojo pipe to ensure the callback is run. mock_searchbox_page_.FlushForTesting(); ASSERT_TRUE(future.Get().has_value()); auto file_token = future.Get().value(); ASSERT_TRUE(handler().tab_context_snapshot_.has_value()); ASSERT_EQ(file_token, handler().tab_context_snapshot_.value().first); // Act // Delete context. handler().DeleteContext(file_token, /*from_automatic_chip=*/false); mock_searchbox_page_.FlushForTesting(); // Assert ASSERT_FALSE(handler().tab_context_snapshot_.has_value()); ASSERT_FALSE(handler().context_input_data().has_value()); } // Tests that context input data is set when there is delayed // context, and cleared when in multi-context mode. TEST_F(ContextualSearchboxHandlerTestTabsTest, ContextInputDataChangesOnMultiContext) { // Add delayed context and ensure context input data is set. auto url1 = GURL("https://www.google.com"); tabs::TabInterface* tab1 = AddTab(url1); const int tab_id1 = tab1->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller1 = static_cast( lens::TabContextualizationController::From(tab1)); EXPECT_CALL(*tab_contextualization_controller1, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(3); base::MockCallback callback; handler().AddTabContext(tab_id1, /*delay_upload=*/true, searchbox::mojom::TabAttachmentSource::kContextMenu, callback.Get()); // Flush the mojo pipe to ensure the callback is run. mock_searchbox_page_.FlushForTesting(); ASSERT_TRUE(handler().context_input_data().has_value()); // Add non-delayed context and ensure context input data has no value. auto url2 = GURL("https://www.wikipedia.com"); tabs::TabInterface* tab2 = AddTab(url2); const int tab_id2 = tab2->GetHandle().raw_value(); MockTabContextualizationController* tab_contextualization_controller2 = static_cast( lens::TabContextualizationController::From(tab2)); EXPECT_CALL(*tab_contextualization_controller2, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(1); base::test::TestFuture> future; handler().AddTabContext(tab_id2, /*delay_upload=*/false, searchbox::mojom::TabAttachmentSource::kContextMenu, future.GetCallback()); mock_searchbox_page_.FlushForTesting(); ASSERT_FALSE(handler().context_input_data().has_value()); // Delete the non-delayed context and ensure context input data now has a // value again because the only context left is of delayed type. auto token = future.Get().value(); handler().DeleteContext(token, /*from_automatic_chip=*/false); mock_searchbox_page_.FlushForTesting(); mock_searchbox_page_.FlushForTesting(); ASSERT_TRUE(handler().context_input_data().has_value()); } TEST_F(ContextualSearchboxHandlerTestTabsTest, TabContextAddedMetric) { // Add a tab. tabs::TabInterface* tab = AddTab(GURL("https://example.com")); const int tab_id = tab->GetHandle().raw_value(); // Mock the call to AddTabContext. MockTabContextualizationController* controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*controller, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(1); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .WillRepeatedly(testing::Invoke( metrics_recorder_ptr, &MockContextualSearchMetricsRecorder::NotifySessionStateChangedBase)); base::test::TestFuture> future; handler().NotifySessionStarted(); handler().AddTabContext(tab_id, false, searchbox::mojom::TabAttachmentSource::kContextMenu, future.GetCallback()); ASSERT_TRUE(future.Wait()); EXPECT_TRUE(future.Get().has_value()); // Check that the histogram was recorded. handler().NotifySessionStateChanged(SessionState::kSessionAbandoned); histogram_tester().ExpectUniqueSample( "ContextualSearch.TabContextAdded.V2.NewTabPage", 1, 1); histogram_tester().ExpectUniqueSample( "ContextualSearch.TabWithDuplicateTitleClicked.V2.NewTabPage", 0, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, OnContextUploadStatusChanged_TerminalFailureRemovesUnderline) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kContextManagementInComposebox); auto* contextual_tasks_service = static_cast( contextual_tasks::ContextualTasksServiceFactory::GetForProfile( profile())); ASSERT_TRUE(contextual_tasks_service); auto active_task_context_provider = std::make_unique( browser_window_interface(), contextual_tasks_service); base::UnguessableToken context_token = base::UnguessableToken::Create(); int32_t tab_id = 456; struct TestObserver : public contextual_tasks::ActiveTaskContextProvider::Observer { void OnContextTabsChanged( const std::set& context_tabs) override { context_tabs_ = context_tabs; } std::set context_tabs_; } observer; active_task_context_provider->AddObserver(&observer); handler().selected_tabs[context_token] = tab_id; active_task_context_provider->AddLocalTabUnderline(tabs::TabHandle(tab_id)); // Verify that the underline is initially added. EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id)) == observer.context_tabs_.end()); handler().OnContextUploadStatusChanged( context_token, lens::MimeType::kUnknown, contextual_search::ContextUploadStatus::kValidationFailed, contextual_search::ContextUploadErrorType::kImageProcessingError); EXPECT_TRUE(handler().selected_tabs.find(context_token) == handler().selected_tabs.end()); EXPECT_TRUE(observer.context_tabs_.find(tabs::TabHandle(tab_id)) == observer.context_tabs_.end()); active_task_context_provider->RemoveObserver(&observer); } TEST_F(ContextualSearchboxHandlerTestTabsTest, OnContextUploadStatusChanged_SuccessDoesNotRemoveUnderline) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kContextManagementInComposebox); auto* contextual_tasks_service = static_cast( contextual_tasks::ContextualTasksServiceFactory::GetForProfile( profile())); ASSERT_TRUE(contextual_tasks_service); auto active_task_context_provider = std::make_unique( browser_window_interface(), contextual_tasks_service); base::UnguessableToken context_token = base::UnguessableToken::Create(); int32_t tab_id = 456; struct TestObserver : public contextual_tasks::ActiveTaskContextProvider::Observer { void OnContextTabsChanged( const std::set& context_tabs) override { context_tabs_ = context_tabs; } std::set context_tabs_; } observer; active_task_context_provider->AddObserver(&observer); handler().selected_tabs[context_token] = tab_id; active_task_context_provider->AddLocalTabUnderline(tabs::TabHandle(tab_id)); // Verify that the underline is initially added. EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id)) == observer.context_tabs_.end()); handler().OnContextUploadStatusChanged( context_token, lens::MimeType::kUnknown, contextual_search::ContextUploadStatus::kUploadSuccessful, std::nullopt); EXPECT_FALSE(handler().selected_tabs.find(context_token) == handler().selected_tabs.end()); EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id)) == observer.context_tabs_.end()); active_task_context_provider->RemoveObserver(&observer); } TEST_F(ContextualSearchboxHandlerTestTabsTest, OnContextUploadStatusChanged_MissingTokenDoesNotAffectOtherUnderlines) { base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeature(omnibox::kContextManagementInComposebox); auto* contextual_tasks_service = static_cast( contextual_tasks::ContextualTasksServiceFactory::GetForProfile( profile())); ASSERT_TRUE(contextual_tasks_service); auto active_task_context_provider = std::make_unique( browser_window_interface(), contextual_tasks_service); base::UnguessableToken context_token1 = base::UnguessableToken::Create(); base::UnguessableToken context_token2 = base::UnguessableToken::Create(); int32_t tab_id1 = 456; int32_t tab_id2 = 789; struct TestObserver : public contextual_tasks::ActiveTaskContextProvider::Observer { void OnContextTabsChanged( const std::set& context_tabs) override { context_tabs_ = context_tabs; } std::set context_tabs_; } observer; active_task_context_provider->AddObserver(&observer); handler().selected_tabs[context_token1] = tab_id1; active_task_context_provider->AddLocalTabUnderline(tabs::TabHandle(tab_id1)); active_task_context_provider->AddLocalTabUnderline(tabs::TabHandle(tab_id2)); // Verify both underlines are initially added. EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id1)) == observer.context_tabs_.end()); EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id2)) == observer.context_tabs_.end()); handler().OnContextUploadStatusChanged( context_token2, lens::MimeType::kUnknown, contextual_search::ContextUploadStatus::kValidationFailed, contextual_search::ContextUploadErrorType::kImageProcessingError); EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id1)) == observer.context_tabs_.end()); EXPECT_FALSE(observer.context_tabs_.find(tabs::TabHandle(tab_id2)) == observer.context_tabs_.end()); active_task_context_provider->RemoveObserver(&observer); } TEST_F(ContextualSearchboxHandlerTestTabsTest, TabStripModelObserverIsAddedWithValidSession) { EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); handler().OnTabAdded(*tab_list(), nullptr, 0); mock_searchbox_page_.FlushForTesting(); } TEST_F(ContextualSearchboxHandlerTestTabsTest, TabStripModelObserverSelectivelyNotifies) { { // Tab insert updates notify. EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); handler().OnTabAdded(*tab_list(), nullptr, 0); mock_searchbox_page_.FlushForTesting(); } { // Tab active updates notify. EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); handler().OnActiveTabChanged(*tab_list(), nullptr); mock_searchbox_page_.FlushForTesting(); } } TEST_F(ContextualSearchboxHandlerTestTabsTest, ActiveTabNavigationObserverNotifies) { tabs::TabInterface* tab = AddTab(GURL("https://example.com")); mock_searchbox_page_.receiver_.reset(); handler_ = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting([&]() { return contextual_session_handle_.get(); })); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); content::WebContents* active_contents = tab->GetContents(); content::WebContentsTester::For(active_contents) ->NavigateAndCommit(GURL("https://example.com/new_path")); mock_searchbox_page_.FlushForTesting(); } TEST_F(ContextualSearchboxHandlerTestTabsTest, ActiveTabNavigationObserverNotifiesOnTitleSet) { tabs::TabInterface* tab = AddTab(GURL("https://example.com")); mock_searchbox_page_.receiver_.reset(); handler_ = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting([&]() { return contextual_session_handle_.get(); })); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); content::WebContents* active_contents = tab->GetContents(); active_contents->UpdateTitleForEntry( active_contents->GetController().GetActiveEntry(), u"New Title"); mock_searchbox_page_.FlushForTesting(); } TEST_F(ContextualSearchboxHandlerTestTabsTest, ActiveTabNavigationObserverUpdatesOnActiveTabChanged) { tabs::TabInterface* tab1 = AddTab(GURL("https://example.com/tab1")); tabs::TabInterface* tab2 = AddTab(GURL("https://example.com/tab2")); mock_searchbox_page_.receiver_.reset(); handler_ = std::make_unique( mojo::PendingReceiver(), mock_searchbox_page_.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting([&]() { return contextual_session_handle_.get(); })); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(0); content::WebContentsTester::For(tab1->GetContents()) ->NavigateAndCommit(GURL("https://example.com/tab1_navigated")); mock_searchbox_page_.FlushForTesting(); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); content::WebContentsTester::For(tab2->GetContents()) ->NavigateAndCommit(GURL("https://example.com/tab2_navigated")); mock_searchbox_page_.FlushForTesting(); ON_CALL(*tab_list_, GetActiveTab()).WillByDefault(testing::Return(tab1)); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); handler().OnActiveTabChanged(*tab_list(), tab1); mock_searchbox_page_.FlushForTesting(); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(1); content::WebContentsTester::For(tab1->GetContents()) ->NavigateAndCommit(GURL("https://example.com/tab1_navigated_again")); mock_searchbox_page_.FlushForTesting(); EXPECT_CALL(mock_searchbox_page_, OnTabStripChanged).Times(0); content::WebContentsTester::For(tab2->GetContents()) ->NavigateAndCommit(GURL("https://example.com/tab2_navigated_again")); mock_searchbox_page_.FlushForTesting(); } // TODO(b:466469292): Figure out how to null-ify the session handle so we can // test the handler behaves correctly in that case. TEST_F(ContextualSearchboxHandlerTestTabsTest, DISABLED_TabStripModelObserverIsNotAddedWithNullSession) { // Create a handler with a null session handle. // Use a new MockSearchboxPage for the new handler. testing::NiceMock local_mock_searchbox_page; auto handler_with_null_session = std::make_unique( mojo::PendingReceiver(), local_mock_searchbox_page.BindAndGetRemote(), profile(), web_contents(), std::make_unique(), base::BindLambdaForTesting( []() -> contextual_search::ContextualSearchSessionHandle* { return nullptr; })); // The observer should not be added, so OnTabStripChanged should not be // called. EXPECT_CALL(local_mock_searchbox_page, OnTabStripChanged).Times(0); handler_with_null_session->OnTabAdded(*tab_list(), nullptr, 0); local_mock_searchbox_page.FlushForTesting(); } TEST_F(ContextualSearchboxHandlerTestTabsTest, TabWithDuplicateTitleClickedMetric) { // Add tabs with duplicate titles. tabs::TabInterface* tab_a1 = AddTab(GURL("https://a1.com")); content::WebContentsTester::For(tab_a1->GetContents())->SetTitle(u"Title A"); tabs::TabInterface* tab_b1 = AddTab(GURL("https://b1.com")); content::WebContentsTester::For(tab_b1->GetContents())->SetTitle(u"Title B"); tabs::TabInterface* tab_a2 = AddTab(GURL("https://a2.com")); content::WebContentsTester::For(tab_a2->GetContents())->SetTitle(u"Title A"); // Mock tab upload flow. MockTabContextualizationController* controller_a1 = static_cast( lens::TabContextualizationController::From(tab_a1)); EXPECT_CALL(*controller_a1, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); MockTabContextualizationController* controller_b1 = static_cast( lens::TabContextualizationController::From(tab_b1)); EXPECT_CALL(*controller_b1, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(2); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(2); auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .WillRepeatedly(testing::Invoke( metrics_recorder_ptr, &MockContextualSearchMetricsRecorder::NotifySessionStateChangedBase)); // Click on a tab with a duplicate title. base::test::TestFuture> future1; handler().NotifySessionStarted(); handler().AddTabContext(tab_a1->GetHandle().raw_value(), false, searchbox::mojom::TabAttachmentSource::kContextMenu, future1.GetCallback()); ASSERT_TRUE(future1.Wait()); EXPECT_TRUE(future1.Get().has_value()); // Click on a tab with a unique title. base::test::TestFuture> future2; handler().AddTabContext(tab_b1->GetHandle().raw_value(), false, searchbox::mojom::TabAttachmentSource::kContextMenu, future2.GetCallback()); ASSERT_TRUE(future2.Wait()); // End the session to log the metrics. handler().NotifySessionStateChanged(SessionState::kSessionAbandoned); histogram_tester().ExpectUniqueSample( "ContextualSearch.TabWithDuplicateTitleClicked.V2.NewTabPage", 1, 1); histogram_tester().ExpectUniqueSample( "ContextualSearch.TabContextAdded.V2.NewTabPage", 2, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, TabWithDuplicateTitleClickedMetric_NoDuplicates) { // Add tabs with unique titles. tabs::TabInterface* tab_a1 = AddTab(GURL("https://a1.com")); content::WebContentsTester::For(tab_a1->GetContents())->SetTitle(u"Title A"); tabs::TabInterface* tab_b1 = AddTab(GURL("https://b1.com")); content::WebContentsTester::For(tab_b1->GetContents())->SetTitle(u"Title B"); // Mock the call to GetPageContext. MockTabContextualizationController* controller_a1 = static_cast( lens::TabContextualizationController::From(tab_a1)); EXPECT_CALL(*controller_a1, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(1); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); auto* metrics_recorder_ptr = GetMetricsRecorderPtr(); ASSERT_THAT(metrics_recorder_ptr, testing::NotNull()); EXPECT_CALL(*metrics_recorder_ptr, NotifySessionStateChanged(testing::_)) .WillRepeatedly(testing::Invoke( metrics_recorder_ptr, &MockContextualSearchMetricsRecorder::NotifySessionStateChangedBase)); // Click on a tab with a unique title. base::test::TestFuture> future; handler().NotifySessionStarted(); handler().AddTabContext(tab_a1->GetHandle().raw_value(), false, searchbox::mojom::TabAttachmentSource::kContextMenu, future.GetCallback()); ASSERT_TRUE(future.Wait()); EXPECT_TRUE(future.Get().has_value()); // End the session to log the metrics. EXPECT_CALL(*GetMetricsRecorderPtr(), NotifySessionStateChanged(SessionState::kSessionAbandoned)) .WillOnce([this](SessionState session_state) { GetMetricsRecorderPtr()->NotifySessionStateChangedBase(session_state); }); handler().NotifySessionStateChanged(SessionState::kSessionAbandoned); histogram_tester().ExpectUniqueSample( "ContextualSearch.TabWithDuplicateTitleClicked.V2.NewTabPage", 0, 1); histogram_tester().ExpectUniqueSample( "ContextualSearch.TabContextAdded.V2.NewTabPage", 1, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, TabContextRecencyRankingMetric) { // Add tabs with unique titles. tabs::TabInterface* tab_a1 = AddTab(GURL("https://a1.com")); content::WebContentsTester::For(tab_a1->GetContents())->SetTitle(u"Title A"); AddTab(GURL("https://b1.com")); content::WebContentsTester::For(all_tabs_.back()->GetContents()) ->SetTitle(u"Title B"); // Mock the call to GetPageContext. MockTabContextualizationController* controller_a1 = static_cast( lens::TabContextualizationController::From(tab_a1)); EXPECT_CALL(*controller_a1, GetPageContext(testing::_)) .WillOnce([](lens::TabContextualizationController::GetPageContextCallback callback) { std::move(callback).Run(std::make_unique()); }); EXPECT_CALL(query_controller(), StartFileUploadFlow(testing::_, testing::NotNull(), testing::_)) .Times(1); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); // Click on the first tab. base::test::TestFuture> future; handler().AddTabContext(tab_a1->GetHandle().raw_value(), false, searchbox::mojom::TabAttachmentSource::kContextMenu, future.GetCallback()); ASSERT_TRUE(future.Wait()); histogram_tester().ExpectUniqueSample( "ContextualSearch.AddedTabContextRecencyRanking.NewTabPage", 1, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, GetRecentTabs) { base::FieldTrialParams params; params[ntp_composebox::kContextMenuMaxTabSuggestions.name] = "2"; base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeatureWithParameters( ntp_composebox::kNtpComposebox, params); // Add only 1 valid tab, and ensure it is the only one returned. auto* about_blank_tab = AddTab(GURL("about:blank")); AddTab(GURL("chrome://webui-is-ignored")); { // Add only 1 valid tab, and ensure it is the only one returned. base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); ASSERT_EQ(tabs.size(), 1u); EXPECT_EQ(tabs[0]->tab_id, about_blank_tab->GetHandle().raw_value()); } auto* contextual_tasks_tab = AddTab(GURL(chrome::kChromeUIContextualTasksURL)); // Case A: Test when the contextual tasks tab is the ACTIVE tab. ON_CALL(*tab_list(), GetActiveTab()) .WillByDefault(testing::Return(contextual_tasks_tab)); { // Verify that the active contextual tasks tab is successfully filtered out, // and only the regular blank tab is returned. base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); ASSERT_EQ(tabs.size(), 1u); EXPECT_EQ(tabs[0]->tab_id, about_blank_tab->GetHandle().raw_value()); } // Case B: Test when the contextual tasks tab is an INACTIVE tab. ON_CALL(*tab_list(), GetActiveTab()) .WillByDefault(testing::Return(about_blank_tab)); { // Verify that even when the contextual tasks tab is inactive, // it is still correctly filtered out. base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); ASSERT_EQ(tabs.size(), 1u); EXPECT_EQ(tabs[0]->tab_id, about_blank_tab->GetHandle().raw_value()); } AddTab(GURL("https://www.google.com")); auto* youtube_tab = AddTab(GURL("https://www.youtube.com")); auto* gmail_tab = AddTab(GURL("https://www.gmail.com")); { // Add more tabs, and ensure no more than the max allowed tabs are returned. base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); ASSERT_EQ(tabs.size(), 2u); EXPECT_EQ(tabs[0]->tab_id, gmail_tab->GetHandle().raw_value()); EXPECT_EQ(tabs[1]->tab_id, youtube_tab->GetHandle().raw_value()); } content::WebContentsTester::For(about_blank_tab->GetContents()) ->SetLastActiveTimeTicks(IncrementTimeTicksAndGet()); { // Activate an older tab, and ensure it is returned first. base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); EXPECT_EQ(tabs[0]->tab_id, about_blank_tab->GetHandle().raw_value()); EXPECT_EQ(tabs[1]->tab_id, gmail_tab->GetHandle().raw_value()); } } TEST_F(ContextualSearchboxHandlerTestTabsTest, GetRecentTabs_IncludesUnloadedTab) { auto* loaded_tab = AddTab(GURL("https://www.google.com")); auto unloaded_mock_tab = std::make_unique(); tabs::TabInterface* unloaded_tab = unloaded_mock_tab.get(); ui::UnownedUserDataHost unloaded_user_data_host; ON_CALL(*unloaded_mock_tab, GetUnownedUserDataHost()) .WillByDefault(testing::ReturnRef(unloaded_user_data_host)); ON_CALL(testing::Const(*unloaded_mock_tab), GetUnownedUserDataHost()) .WillByDefault(testing::ReturnRef(unloaded_user_data_host)); ON_CALL(*unloaded_mock_tab, GetTabFeatures()) .WillByDefault(testing::Return(nullptr)); ON_CALL(testing::Const(*unloaded_mock_tab), GetTabFeatures()) .WillByDefault(testing::Return(nullptr)); ON_CALL(*unloaded_mock_tab, GetContents) .WillByDefault(testing::Return(nullptr)); ON_CALL(*unloaded_mock_tab, GetURL) .WillByDefault(testing::Return(GURL("https://www.unloaded.com"))); ON_CALL(*unloaded_mock_tab, GetTitle) .WillByDefault(testing::Return(u"Unloaded Tab")); ON_CALL(*unloaded_mock_tab, GetLastActiveTime) .WillByDefault(testing::Return(base::Time::Now() - base::Days(10))); ON_CALL(*unloaded_mock_tab, GetTabHandle).WillByDefault(testing::Return(999)); all_tabs_.push_back(unloaded_tab); ON_CALL(*tab_list(), GetAllTabs()).WillByDefault(testing::Return(all_tabs_)); base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); ASSERT_EQ(tabs.size(), 2u); EXPECT_EQ(tabs[0]->tab_id, loaded_tab->GetHandle().raw_value()); EXPECT_EQ(tabs[1]->tab_id, unloaded_tab->GetHandle().raw_value()); EXPECT_EQ(tabs[1]->url, GURL("https://www.unloaded.com")); EXPECT_EQ(tabs[1]->title, "Unloaded Tab"); all_tabs_.pop_back(); } TEST_F(ContextualSearchboxHandlerTestTabsTest, GetRecentTabs_UsesServerLimit) { base::FieldTrialParams params; params[ntp_composebox::kContextMenuMaxTabSuggestions.name] = "2"; base::test::ScopedFeatureList feature_list; feature_list.InitAndEnableFeatureWithParameters( ntp_composebox::kNtpComposebox, params); // Add 5 valid tabs. AddTab(GURL("https://www.google.com")); AddTab(GURL("https://www.youtube.com")); AddTab(GURL("https://www.gmail.com")); auto* example_tab = AddTab(GURL("https://www.example.com")); auto* chromium_tab = AddTab(GURL("https://www.chromium.org")); // Initially, it should use the feature param limit. base::test::TestFuture> future1; handler().GetRecentTabs(future1.GetCallback()); auto tabs = future1.Take(); ASSERT_EQ(tabs.size(), 2u); EXPECT_EQ(tabs[0]->tab_id, chromium_tab->GetHandle().raw_value()); EXPECT_EQ(tabs[1]->tab_id, example_tab->GetHandle().raw_value()); // Now set a server-side limit of 1. omnibox::InputState state; state.max_inputs_by_type[omnibox::InputType::INPUT_TYPE_BROWSER_TAB] = 1; handler().input_state_model()->set_state_for_testing(state); base::test::TestFuture> future2; handler().GetRecentTabs(future2.GetCallback()); tabs = future2.Take(); ASSERT_EQ(tabs.size(), 1u); EXPECT_EQ(tabs[0]->tab_id, chromium_tab->GetHandle().raw_value()); // Fallback to feature param limit if not in map. state.max_inputs_by_type.erase(omnibox::InputType::INPUT_TYPE_BROWSER_TAB); handler().input_state_model()->set_state_for_testing(state); base::test::TestFuture> future3; handler().GetRecentTabs(future3.GetCallback()); tabs = future3.Take(); ASSERT_EQ(tabs.size(), 2u); EXPECT_EQ(tabs[0]->tab_id, chromium_tab->GetHandle().raw_value()); EXPECT_EQ(tabs[1]->tab_id, example_tab->GetHandle().raw_value()); } class ContextualSearchboxHandlerSignedInTestTabsTest : public ContextualSearchboxHandlerTestTabsTest { public: ContextualSearchboxHandlerSignedInTestTabsTest() : get_variations_ids_provider_( variations::VariationsIdsProvider::Mode::kUseSignedInState) {} private: variations::test::ScopedVariationsIdsProvider get_variations_ids_provider_; }; TEST_F(ContextualSearchboxHandlerSignedInTestTabsTest, GetRecentTabs_SetsShowInRecentTabChip) { // Disable modules so that NTP does not try to initialize modules as // test set up does not support this flow. base::test::ScopedFeatureList feature_list; feature_list.InitWithFeatures(/*enabled_features=*/{}, /*disabled_features=*/{ ntp_features::kNtpCalendarModule, ntp_features::kNtpOutlookCalendarModule, ntp_features::kNtpDriveModule, }); // Add a regular tab, a google search tab, and another regular tab. auto* example_tab = AddTab(GURL("https://www.example.com")); auto* search_tab = AddTab(GURL("https://www.google.com/search?q=test")); auto* chromium_tab = AddTab(GURL("https://www.chromium.org")); // Navigate to NTP AddTab(GURL("chrome://newtab")); // Get the recent tabs. base::test::TestFuture> future; handler().GetRecentTabs(future.GetCallback()); auto tabs = future.Take(); // Expect all three non chrome WebUI tabs to be returned. ASSERT_EQ(tabs.size(), 3u); EXPECT_EQ(tabs[0]->tab_id, chromium_tab->GetHandle().raw_value()); EXPECT_FALSE(tabs[0]->show_in_current_tab_chip); EXPECT_TRUE(tabs[0]->show_in_previous_tab_chip); EXPECT_EQ(tabs[1]->tab_id, search_tab->GetHandle().raw_value()); EXPECT_FALSE(tabs[1]->show_in_previous_tab_chip); EXPECT_EQ(tabs[2]->tab_id, example_tab->GetHandle().raw_value()); EXPECT_TRUE(tabs[2]->show_in_previous_tab_chip); } TEST_F(ContextualSearchboxHandlerTestTabsTest, DuplicateTabsShownMetric) { // Add tabs with duplicate titles. AddTab(GURL("https://a1.com")); content::WebContentsTester::For(all_tabs_[0]->GetContents()) ->SetTitle(u"Title A"); AddTab(GURL("https://b1.com")); content::WebContentsTester::For(all_tabs_[1]->GetContents()) ->SetTitle(u"Title B"); AddTab(GURL("https://a2.com")); content::WebContentsTester::For(all_tabs_[2]->GetContents()) ->SetTitle(u"Title A"); AddTab(GURL("https://c1.com")); content::WebContentsTester::For(all_tabs_[3]->GetContents()) ->SetTitle(u"Title C"); AddTab(GURL("https://a3.com")); content::WebContentsTester::For(all_tabs_[4]->GetContents()) ->SetTitle(u"Title A"); AddTab(GURL("https://b2.com")); content::WebContentsTester::For(all_tabs_[5]->GetContents()) ->SetTitle(u"Title B"); AddTab(GURL("https://b3.com")); content::WebContentsTester::For(all_tabs_[6]->GetContents()) ->SetTitle(u"Title B"); base::test::TestFuture> tab_info_future; handler().GetRecentTabs(tab_info_future.GetCallback()); auto tabs = tab_info_future.Take(); // Even though there are more duplicates above, only three tabs are kept // by `GetRecentTabs` with the default configuration, so there can only // be one title with more than one instance in the list. histogram_tester().ExpectUniqueSample( "ContextualSearch.DuplicateTabTitlesShownCount.NewTabPage", 1, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, ActiveTabsCountMetric) { AddTab(GURL("https://a1.com")); AddTab(GURL("https://b1.com")); AddTab(GURL("https://a2.com")); base::test::TestFuture> tab_info_future; handler().GetRecentTabs(tab_info_future.GetCallback()); auto tabs = tab_info_future.Take(); histogram_tester().ExpectUniqueSample( "ContextualSearch.ActiveTabsCountOnContextMenuOpen.NewTabPage", 3, 1); } TEST_F(ContextualSearchboxHandlerTestTabsTest, GetTabPreview_InvalidTab) { base::test::TestFuture&> future; handler().GetTabPreview(12345, future.GetCallback()); std::optional preview = future.Get(); ASSERT_FALSE(preview.has_value()); } TEST_F(ContextualSearchboxHandlerTestTabsTest, GetTabPreview_CaptureFails) { tabs::TabInterface* tab = AddTab(GURL("https://a1.com")); MockTabContextualizationController* controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*controller, CaptureScreenshot(testing::_, testing::_)) .WillOnce( [](std::optional image_options, lens::TabContextualizationController::CaptureScreenshotCallback callback) { std::move(callback).Run(SkBitmap()); }); base::test::TestFuture&> future; handler().GetTabPreview(tab->GetHandle().raw_value(), future.GetCallback()); std::optional preview = future.Get(); ASSERT_FALSE(preview.has_value()); } TEST_F(ContextualSearchboxHandlerTestTabsTest, GetTabPreview_Success) { tabs::TabInterface* tab = AddTab(GURL("https://a1.com")); SkBitmap bitmap; bitmap.allocN32Pixels(1, 1); bitmap.eraseColor(SK_ColorRED); MockTabContextualizationController* controller = static_cast( lens::TabContextualizationController::From(tab)); EXPECT_CALL(*controller, CaptureScreenshot(testing::_, testing::_)) .WillOnce( [&bitmap]( std::optional image_options, lens::TabContextualizationController::CaptureScreenshotCallback callback) { std::move(callback).Run(bitmap); }); base::test::TestFuture&> future; handler().GetTabPreview(tab->GetHandle().raw_value(), future.GetCallback()); std::optional preview = future.Get(); ASSERT_TRUE(preview.has_value()); EXPECT_EQ(preview.value(), webui::GetBitmapDataUrl(bitmap)); } TEST_F(ContextualSearchboxHandlerTest, IncompatibleToolDisablesSmartTabSharing) { base::test::ScopedFeatureList local_feature_list; local_feature_list.InitWithFeaturesAndParameters( {{contextual_tasks::kContextualTasksContext, {{"ContextualTasksContextSmartTabSharing", "true"}}}, {contextual_tasks::kContextualTasksForceEntryPointEligibility, {}}}, {}); // Enable STS. EXPECT_CALL(mock_searchbox_page_, UpdateSmartTabSharingActive(true)).Times(1); handler().SetSmartTabSharingActive(true); mock_searchbox_page_.FlushForTesting(); EXPECT_TRUE(handler().IsSmartTabSharingActive()); // Select Canvas. EXPECT_CALL(mock_searchbox_page_, UpdateSmartTabSharingActive(false)) .Times(1); handler().SetActiveToolMode(omnibox::TOOL_MODE_CANVAS, /*is_set_by_aim=*/false); mock_searchbox_page_.FlushForTesting(); // STS should now be effectively inactive. EXPECT_FALSE(handler().IsSmartTabSharingActive()); // Select Unspecified again. EXPECT_CALL(mock_searchbox_page_, UpdateSmartTabSharingActive(true)).Times(1); handler().SetActiveToolMode(omnibox::TOOL_MODE_UNSPECIFIED, /*is_set_by_aim=*/false); mock_searchbox_page_.FlushForTesting(); // STS should be active again. EXPECT_TRUE(handler().IsSmartTabSharingActive()); } TEST_F(ContextualSearchboxHandlerTest, ResetInputStateModelResetsSmartTabSharingActiveState) { base::test::ScopedFeatureList local_feature_list; local_feature_list.InitWithFeaturesAndParameters( {{contextual_tasks::kContextualTasksContext, {{"ContextualTasksContextSmartTabSharing", "true"}}}, {contextual_tasks::kContextualTasksForceEntryPointEligibility, {}}}, {}); // Enable STS on current handler. handler().SetSmartTabSharingActive(true); EXPECT_TRUE(handler().IsSmartTabSharingActive()); ASSERT_TRUE(handler().input_state_model()); EXPECT_TRUE(handler().input_state_model()->IsSmartTabSharingActive()); // Simulate starting a new session (new model) and resetting input state model. query_controller_ = nullptr; metrics_recorder_ = nullptr; auto query_controller_config_params = std::make_unique< contextual_search::ContextualSearchContextController::ConfigParams>(); auto query_controller_ptr = std::make_unique( /*identity_manager=*/nullptr, url_loader_factory(), version_info::Channel::UNKNOWN, "en-US", template_url_service(), fake_variations_client(), std::move(query_controller_config_params)); auto metrics_recorder_ptr = std::make_unique(); query_controller_ = query_controller_ptr.get(); metrics_recorder_ = metrics_recorder_ptr.get(); contextual_session_handle_ = service_->CreateSessionForTesting( std::move(query_controller_ptr), std::move(metrics_recorder_ptr)); contextual_session_handle_->CheckSearchContentSharingSettings( profile()->GetPrefs()); handler().ResetInputStateModel(); handler().NotifySessionStarted(); // Trigger model recreation. handler().GetValidInputStateForTesting(); // ResetInputStateModel() should have cleared the active state back to default (false). ASSERT_TRUE(handler().input_state_model()); EXPECT_FALSE(handler().input_state_model()->IsSmartTabSharingActive()); EXPECT_FALSE(handler().IsSmartTabSharingActive()); } class ContextualSearchboxHandlerContextUploadStatusTest : public ContextualSearchboxHandlerTest, public testing::WithParamInterface< composebox_query::mojom::ContextUploadStatus> {}; TEST_P(ContextualSearchboxHandlerContextUploadStatusTest, OnContextUploadStatusChanged) { contextual_search::ContextUploadStatus status; EXPECT_CALL(mock_searchbox_page_, OnContextualInputStatusChanged) .Times(1) .WillOnce( [&status]( const base::UnguessableToken& context_token, contextual_search::ContextUploadStatus context_upload_status, std::optional context_upload_error_type) { status = context_upload_status; }); EXPECT_CALL(mock_searchbox_page_, OnInputStateChanged).Times(1); const auto expected_status = GetParam(); contextual_search::ContextUploadStatus status_cpp = mojo::EnumTraits< composebox_query::mojom::ContextUploadStatus, contextual_search::ContextUploadStatus>::FromMojom(expected_status); base::UnguessableToken token = base::UnguessableToken::Create(); handler().OnContextUploadStatusChanged(token, lens::MimeType::kPdf, status_cpp, std::nullopt); mock_searchbox_page_.FlushForTesting(); EXPECT_EQ(status_cpp, status); } INSTANTIATE_TEST_SUITE_P( All, ContextualSearchboxHandlerContextUploadStatusTest, testing::Values( composebox_query::mojom::ContextUploadStatus::kNotUploaded, composebox_query::mojom::ContextUploadStatus::kProcessing, composebox_query::mojom::ContextUploadStatus::kValidationFailed, composebox_query::mojom::ContextUploadStatus::kUploadStarted, composebox_query::mojom::ContextUploadStatus::kUploadSuccessful, composebox_query::mojom::ContextUploadStatus::kUploadFailed, composebox_query::mojom::ContextUploadStatus::kUploadExpired, composebox_query::mojom::ContextUploadStatus::kUploadReplaced)); #if !BUILDFLAG(IS_ANDROID) class FakeDesktopCapturer : public webrtc::DesktopCapturer { public: explicit FakeDesktopCapturer( webrtc::DesktopSize size = webrtc::DesktopSize(1, 1), Result result = Result::SUCCESS) : size_(size), result_(result) {} ~FakeDesktopCapturer() override = default; void Start(Callback* callback) override { callback_ = callback; } void CaptureFrame() override { if (result_ != Result::SUCCESS) { callback_->OnCaptureResult(result_, nullptr); return; } auto frame = std::make_unique(size_); frame->SetFrameDataToBlack(); callback_->OnCaptureResult(Result::SUCCESS, std::move(frame)); } bool GetSourceList(SourceList* sources) override { return true; } bool SelectSource(SourceId id) override { last_selected_source_id_ = id; return true; } std::optional last_selected_source_id() const { return last_selected_source_id_; } private: webrtc::DesktopSize size_; Result result_; std::optional last_selected_source_id_; raw_ptr callback_ = nullptr; }; class MockScreenshareDelegate : public ContextualSearchboxHandler::ScreenshareDelegate { public: MOCK_METHOD(void, ShowScreenshotMenu, (const gfx::Rect&, base::WeakPtr), (override)); MOCK_METHOD(void, OnScreensharePickerOpened, (), (override)); MOCK_METHOD(void, OnScreensharePickerClosed, (), (override)); }; TEST_F(ContextualSearchboxHandlerTest, ShowScreenshotMenu_ForwardsToDelegate) { MockScreenshareDelegate delegate; EXPECT_CALL(delegate, ShowScreenshotMenu(gfx::Rect(1, 2, 3, 4), testing::_)); handler().set_screenshare_delegate(&delegate); handler().ShowScreenshotMenu(gfx::Rect(1, 2, 3, 4)); } TEST_F(ContextualSearchboxHandlerTest, ShowScreenshotMenu_NoDelegate_NotifiesClosed) { handler().set_screenshare_delegate(nullptr); EXPECT_CALL(mock_searchbox_page_, OnScreenshotMenuClosed()); handler().ShowScreenshotMenu(gfx::Rect(1, 2, 3, 4)); mock_searchbox_page_.FlushForTesting(); } TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_Success) { SetupScreenshotUploadConfig(); MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); FakeDesktopMediaPickerFactory picker_factory; handler().set_desktop_media_picker_factory_for_testing(&picker_factory); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique(webrtc::DesktopSize(4000, 2000))); FakeDesktopMediaPickerFactory::TestFlags test_flags; test_flags.expect_screens = true; test_flags.expect_windows = true; test_flags.picker_result = content::DesktopMediaID(content::DesktopMediaID::TYPE_WINDOW, 42); picker_factory.SetTestFlags(base::span_from_ref(test_flags)); std::unique_ptr captured_input_data; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, std::unique_ptr input_data, std::optional image_options) { captured_input_data = std::move(input_data); EXPECT_TRUE(image_options.has_value()); }); base::UnguessableToken callback_token; EXPECT_CALL(mock_searchbox_page_, AddFileContext) .WillOnce([&](const base::UnguessableToken& token, searchbox::mojom::SelectedFileInfoPtr file_info) { callback_token = token; EXPECT_EQ(file_info->file_name, "Screenshot.png"); EXPECT_EQ(file_info->mime_type, "image/png"); EXPECT_TRUE(file_info->image_data_url.has_value()); if (file_info->image_data_url.has_value()) { EXPECT_TRUE(base::StartsWith(*file_info->image_data_url, "data:image/png;base64,")); // Verify thumbnail dimensions constraint (<= 120px, aspect ratio // preserved 120x60). std::string base64_payload = file_info->image_data_url->substr( std::string_view("data:image/png;base64,").length()); std::optional> thumb_bytes = base::Base64Decode(base64_payload); EXPECT_TRUE(thumb_bytes.has_value()); if (thumb_bytes) { SkBitmap thumb_bitmap = gfx::PNGCodec::Decode(*thumb_bytes); EXPECT_FALSE(thumb_bitmap.isNull()); EXPECT_EQ(thumb_bitmap.width(), 120); EXPECT_EQ(thumb_bitmap.height(), 60); } } }); base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/false, future.GetCallback()); EXPECT_TRUE(future.Get().has_value()); mock_searchbox_page_.FlushForTesting(); auto uploaded_tokens = handler().GetUploadedContextTokens(); ASSERT_EQ(uploaded_tokens.size(), 1u); EXPECT_EQ(uploaded_tokens[0], callback_token); ASSERT_TRUE(captured_input_data); EXPECT_EQ(captured_input_data->file_name, "Screenshot.png"); EXPECT_EQ(captured_input_data->primary_content_type, lens::MimeType::kImage); EXPECT_EQ(captured_input_data->mime_type_string, "image/png"); // Verify oversized image downscaling constraint (<= 2048px, aspect ratio // preserved 2048x1024). ASSERT_TRUE(captured_input_data->context_input.has_value()); ASSERT_FALSE(captured_input_data->context_input->empty()); const auto& file_data = (*captured_input_data->context_input)[0].bytes_; SkBitmap main_bitmap = gfx::PNGCodec::Decode(file_data); EXPECT_FALSE(main_bitmap.isNull()); EXPECT_EQ(main_bitmap.width(), 2048); EXPECT_EQ(main_bitmap.height(), 1024); handler().set_desktop_media_picker_factory_for_testing(nullptr); handler().set_screenshare_delegate(nullptr); } TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_SmallImageNotResized) { profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); scoped_config().config.mutable_composebox()->set_max_num_files(5); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_max_size_bytes(1024 * 1024); scoped_config() .config.mutable_composebox() ->mutable_image_upload() ->set_mime_types_allowed("image/png"); FakeDesktopMediaPickerFactory picker_factory; handler().set_desktop_media_picker_factory_for_testing(&picker_factory); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique(webrtc::DesktopSize(50, 30))); FakeDesktopMediaPickerFactory::TestFlags test_flags; test_flags.expect_screens = true; test_flags.expect_windows = true; test_flags.picker_result = content::DesktopMediaID(content::DesktopMediaID::TYPE_WINDOW, 42); picker_factory.SetTestFlags(base::span_from_ref(test_flags)); std::unique_ptr captured_input_data; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, std::unique_ptr input_data, std::optional image_options) { captured_input_data = std::move(input_data); }); EXPECT_CALL(mock_searchbox_page_, AddFileContext) .WillOnce([&](const base::UnguessableToken& token, searchbox::mojom::SelectedFileInfoPtr file_info) { EXPECT_TRUE(file_info->image_data_url.has_value()); if (file_info->image_data_url.has_value()) { std::string base64_payload = file_info->image_data_url->substr( std::string_view("data:image/png;base64,").length()); std::optional> thumb_bytes = base::Base64Decode(base64_payload); EXPECT_TRUE(thumb_bytes.has_value()); if (thumb_bytes) { SkBitmap thumb_bitmap = gfx::PNGCodec::Decode(*thumb_bytes); EXPECT_FALSE(thumb_bitmap.isNull()); EXPECT_EQ(thumb_bitmap.width(), 50); EXPECT_EQ(thumb_bitmap.height(), 30); } } }); base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/false, future.GetCallback()); EXPECT_TRUE(future.Get().has_value()); mock_searchbox_page_.FlushForTesting(); ASSERT_TRUE(captured_input_data); ASSERT_TRUE(captured_input_data->context_input.has_value()); ASSERT_FALSE(captured_input_data->context_input->empty()); const auto& file_data = (*captured_input_data->context_input)[0].bytes_; SkBitmap main_bitmap = gfx::PNGCodec::Decode(file_data); EXPECT_FALSE(main_bitmap.isNull()); EXPECT_EQ(main_bitmap.width(), 50); EXPECT_EQ(main_bitmap.height(), 30); handler().set_desktop_media_picker_factory_for_testing(nullptr); handler().set_screenshare_delegate(nullptr); } TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_Cancelled) { MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); FakeDesktopMediaPickerFactory picker_factory; handler().set_desktop_media_picker_factory_for_testing(&picker_factory); FakeDesktopMediaPickerFactory::TestFlags test_flags; test_flags.expect_screens = true; test_flags.expect_windows = true; test_flags.picker_result = content::DesktopMediaID(); picker_factory.SetTestFlags(base::span_from_ref(test_flags)); base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/true, future.GetCallback()); EXPECT_FALSE(future.Get().has_value()); handler().set_desktop_media_picker_factory_for_testing(nullptr); handler().set_screenshare_delegate(nullptr); } // Tests that CaptureRegionScreenshot successfully captures a screen/desktop // region, generates thumbnail and PNG data, triggers the file upload flow to // Lens, and notifies the WebUI with the attached file context token. TEST_F(ContextualSearchboxHandlerTest, CaptureRegionScreenshot_Success) { MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); scoped_config().config.mutable_composebox()->set_max_num_files(5); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_max_size_bytes(1024 * 1024); scoped_config() .config.mutable_composebox() ->mutable_image_upload() ->set_mime_types_allowed("image/png"); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique(webrtc::DesktopSize(3840, 2160))); std::unique_ptr captured_input_data; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, std::unique_ptr input_data, std::optional image_options) { captured_input_data = std::move(input_data); EXPECT_TRUE(image_options.has_value()); }); base::UnguessableToken callback_token; EXPECT_CALL(mock_searchbox_page_, AddFileContext) .WillOnce([&](const base::UnguessableToken& token, searchbox::mojom::SelectedFileInfoPtr file_info) { callback_token = token; EXPECT_EQ(file_info->file_name, "Screenshot.png"); EXPECT_EQ(file_info->mime_type, "image/png"); EXPECT_TRUE(file_info->image_data_url.has_value()); }); base::test::TestFuture&> future; handler().CaptureRegionScreenshot(future.GetCallback()); EXPECT_TRUE(future.Get().has_value()); mock_searchbox_page_.FlushForTesting(); auto uploaded_tokens = handler().GetUploadedContextTokens(); ASSERT_EQ(uploaded_tokens.size(), 1u); EXPECT_EQ(uploaded_tokens[0], callback_token); ASSERT_TRUE(captured_input_data); EXPECT_EQ(captured_input_data->file_name, "Screenshot.png"); EXPECT_EQ(captured_input_data->primary_content_type, lens::MimeType::kImage); EXPECT_EQ(captured_input_data->mime_type_string, "image/png"); handler().set_screenshare_delegate(nullptr); } // Tests that concurrent calls to CaptureRegionScreenshot while a capture is // already in progress are rejected, returning an empty token for subsequent // calls. TEST_F(ContextualSearchboxHandlerTest, CaptureRegionScreenshot_AlreadyCapturing) { MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); profile()->GetPrefs()->SetInteger( contextual_search::kSearchContentSharingSettings, static_cast( contextual_search::SearchContentSharingSettingsValue::kEnabled)); scoped_config().config.mutable_composebox()->set_max_num_files(5); scoped_config() .config.mutable_composebox() ->mutable_attachment_upload() ->set_max_size_bytes(1024 * 1024); scoped_config() .config.mutable_composebox() ->mutable_image_upload() ->set_mime_types_allowed("image/png"); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique(webrtc::DesktopSize(100, 100))); EXPECT_CALL(query_controller(), StartFileUploadFlow); EXPECT_CALL(mock_searchbox_page_, AddFileContext); base::test::TestFuture&> future1; base::test::TestFuture&> future2; handler().CaptureRegionScreenshot(future1.GetCallback()); handler().CaptureRegionScreenshot(future2.GetCallback()); EXPECT_FALSE(future2.Get().has_value()); EXPECT_TRUE(future1.Get().has_value()); handler().set_screenshare_delegate(nullptr); } // Tests that CaptureRegionScreenshot handles an empty or failed capturer result // gracefully by closing the picker delegate and returning an empty token. TEST_F(ContextualSearchboxHandlerTest, CaptureRegionScreenshot_EmptyBitmap) { MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique( webrtc::DesktopSize(1, 1), webrtc::DesktopCapturer::Result::ERROR_PERMANENT)); base::test::TestFuture&> future; handler().CaptureRegionScreenshot(future.GetCallback()); EXPECT_FALSE(future.Get().has_value()); handler().set_screenshare_delegate(nullptr); } #if BUILDFLAG(IS_MAC) using content::desktop_capture::ScopedNativePickerForTesting; // Tests that StartScreenshare on macOS 14+ successfully opens the native // ScreenCaptureKit picker (SCContentSharingPicker), captures the selected // source, and attaches the screenshot to the WebUI. TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_NativePicker_Success) { if (base::mac::MacOSMajorVersion() < 14) { GTEST_SKIP() << "Native picker only supported on macOS 14+"; } base::test::ScopedFeatureList scoped_feature_list; scoped_feature_list.InitAndEnableFeature(media::kUseSCContentSharingPicker); SetupScreenshotUploadConfig(); MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique()); ScopedNativePickerForTesting scoped_picker( ScopedNativePickerForTesting::Action::kSelectSource); std::unique_ptr captured_input_data; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, std::unique_ptr input_data, std::optional image_options) { captured_input_data = std::move(input_data); EXPECT_TRUE(image_options.has_value()); }); base::UnguessableToken callback_token; EXPECT_CALL(mock_searchbox_page_, AddFileContext) .WillOnce([&](const base::UnguessableToken& token, searchbox::mojom::SelectedFileInfoPtr file_info) { callback_token = token; EXPECT_EQ(file_info->file_name, "Screenshot.png"); EXPECT_EQ(file_info->mime_type, "image/png"); EXPECT_TRUE(file_info->image_data_url.has_value()); if (file_info->image_data_url.has_value()) { EXPECT_TRUE(base::StartsWith(*file_info->image_data_url, "data:image/png;base64,")); } }); base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/false, future.GetCallback()); EXPECT_TRUE(future.Get().has_value()); mock_searchbox_page_.FlushForTesting(); auto uploaded_tokens = handler().GetUploadedContextTokens(); ASSERT_EQ(uploaded_tokens.size(), 1u); EXPECT_EQ(uploaded_tokens[0], callback_token); handler().set_desktop_media_picker_factory_for_testing(nullptr); handler().set_screenshare_delegate(nullptr); } // Tests that when the user cancels the macOS 14+ native ScreenCaptureKit // picker, StartScreenshare terminates cleanly without error and returns an // empty token. TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_NativePicker_Cancelled) { if (base::mac::MacOSMajorVersion() < 14) { GTEST_SKIP() << "Native picker only supported on macOS 14+"; } base::test::ScopedFeatureList scoped_feature_list; scoped_feature_list.InitAndEnableFeature(media::kUseSCContentSharingPicker); MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); ScopedNativePickerForTesting scoped_picker( ScopedNativePickerForTesting::Action::kCancel); base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/false, future.GetCallback()); EXPECT_FALSE(future.Get().has_value()); handler().set_screenshare_delegate(nullptr); } // Tests that if the macOS 14+ native ScreenCaptureKit picker encounters an // error, StartScreenshare gracefully falls back to Chrome's default desktop // media picker dialog. TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_NativePicker_Error_FallsBackToDefaultPicker) { if (base::mac::MacOSMajorVersion() < 14) { GTEST_SKIP() << "Native picker only supported on macOS 14+"; } base::test::ScopedFeatureList scoped_feature_list; scoped_feature_list.InitAndEnableFeature(media::kUseSCContentSharingPicker); SetupScreenshotUploadConfig(); MockScreenshareDelegate delegate; EXPECT_CALL(delegate, OnScreensharePickerOpened()); EXPECT_CALL(delegate, OnScreensharePickerClosed()); handler().set_screenshare_delegate(&delegate); FakeDesktopMediaPickerFactory picker_factory; handler().set_desktop_media_picker_factory_for_testing(&picker_factory); FakeDesktopMediaPickerFactory::TestFlags test_flags; test_flags.expect_screens = true; test_flags.expect_windows = true; test_flags.picker_result = content::DesktopMediaID(content::DesktopMediaID::TYPE_WINDOW, 42); picker_factory.SetTestFlags(base::span_from_ref(test_flags)); content::desktop_capture::ScopedDesktopCapturerForTesting scoped_capturer( std::make_unique()); ScopedNativePickerForTesting scoped_picker( ScopedNativePickerForTesting::Action::kError); std::unique_ptr captured_input_data; EXPECT_CALL(query_controller(), StartFileUploadFlow) .WillOnce([&](const base::UnguessableToken& token, std::unique_ptr input_data, std::optional image_options) { captured_input_data = std::move(input_data); EXPECT_TRUE(image_options.has_value()); }); base::UnguessableToken callback_token; EXPECT_CALL(mock_searchbox_page_, AddFileContext) .WillOnce([&](const base::UnguessableToken& token, searchbox::mojom::SelectedFileInfoPtr file_info) { callback_token = token; EXPECT_EQ(file_info->file_name, "Screenshot.png"); EXPECT_EQ(file_info->mime_type, "image/png"); EXPECT_TRUE(file_info->image_data_url.has_value()); if (file_info->image_data_url.has_value()) { EXPECT_TRUE(base::StartsWith(*file_info->image_data_url, "data:image/png;base64,")); } }); base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/false, future.GetCallback()); EXPECT_TRUE(future.Get().has_value()); mock_searchbox_page_.FlushForTesting(); auto uploaded_tokens = handler().GetUploadedContextTokens(); ASSERT_EQ(uploaded_tokens.size(), 1u); EXPECT_EQ(uploaded_tokens[0], callback_token); handler().set_desktop_media_picker_factory_for_testing(nullptr); handler().set_screenshare_delegate(nullptr); } #endif #else TEST_F(ContextualSearchboxHandlerTest, StartScreenshare_AndroidAlwaysFails) { base::test::TestFuture&> future; handler().StartScreenshare(/*prefer_entire_screen=*/false, future.GetCallback()); EXPECT_FALSE(future.Get().has_value()); } // Tests that CaptureRegionScreenshot on Android immediately fails and returns // an empty token since desktop screen capture is unsupported on Android. TEST_F(ContextualSearchboxHandlerTest, CaptureRegionScreenshot_AndroidAlwaysFails) { base::test::TestFuture&> future; handler().CaptureRegionScreenshot(future.GetCallback()); EXPECT_FALSE(future.Get().has_value()); } #endif