[omnibox][search aggregator] Use RichSuggestionImage for aggregator keyword matches

Currently, aggregator keyword matches rely on favicons to fetch the
policy set favicon URL. However, favicons rely on user navigation
to fetch the icon through the history service. As a result,
when a user sets the favicon URL for search aggregator policy,
the icon may never be fetched/used.

RichSuggestionImage requests an image using the bitmap_fetcher_service.
By using RichSuggestionImage, the icon is fetched using the favicon URL
set in the EnterpriseSearchAggregatorSettings policy.

Before: http://screencast/cast/NDk2NTk4MjA0Mzg5Nzg1NnxiOGRlYTU4NS1iNA
After: http://screencast/cast/NTQ3Mjc4MjE3ODU4MjUyOHxiMzZkNzViZS1kNA

Bug: b:380509200
Change-Id: Ie0a3903d973112381f12a92c86aae91bf0d5f83d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6173646
Reviewed-by: manuk hovanesian <manukh@chromium.org>
Commit-Queue: Alex Chen <alexwchen@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1413837}
diff --git a/chrome/browser/ui/omnibox/chrome_omnibox_client.cc b/chrome/browser/ui/omnibox/chrome_omnibox_client.cc
index 4215e6e3..29e826a 100644
--- a/chrome/browser/ui/omnibox/chrome_omnibox_client.cc
+++ b/chrome/browser/ui/omnibox/chrome_omnibox_client.cc
@@ -350,12 +350,22 @@
   int result_index = -1;
   for (const AutocompleteMatch& match : result) {
     ++result_index;
-    if (match.ImageUrl().is_empty()) {
-      continue;
+    if (!match.ImageUrl().is_empty()) {
+      request_ids_.push_back(bitmap_fetcher_service->RequestImage(
+          match.ImageUrl(), base::BindOnce(on_bitmap_fetched, result_index)));
+    } else if (AutocompleteMatch::IsFeaturedEnterpriseSearchType(match.type)) {
+      // Request the policy favicon for featured search aggregator matches.
+      // `IsFeaturedEnterpriseSearchType()` also includes site search matches.
+      // We should only fetch the bitmap for enterprise search aggregators.
+      const TemplateURL* turl =
+          match.GetTemplateURL(GetTemplateURLService(), false);
+      if (turl && turl->policy_origin() ==
+                      TemplateURLData::PolicyOrigin::kSearchAggregator) {
+        request_ids_.push_back(bitmap_fetcher_service->RequestImage(
+            turl->favicon_url(),
+            base::BindOnce(on_bitmap_fetched, result_index)));
+      }
     }
-
-    request_ids_.push_back(bitmap_fetcher_service->RequestImage(
-        match.ImageUrl(), base::BindOnce(on_bitmap_fetched, result_index)));
   }
 }
 
diff --git a/chrome/browser/ui/views/omnibox/omnibox_popup_view_views.cc b/chrome/browser/ui/views/omnibox/omnibox_popup_view_views.cc
index 50d87fa2..bf1eb9b7 100644
--- a/chrome/browser/ui/views/omnibox/omnibox_popup_view_views.cc
+++ b/chrome/browser/ui/views/omnibox/omnibox_popup_view_views.cc
@@ -373,10 +373,15 @@
     result_view->SetVisible(!group_hidden && !row_hidden);
     result_view->UpdateAccessibilityProperties();
 
-    const SkBitmap* bitmap = model()->GetPopupRichSuggestionBitmap(i);
-    if (bitmap) {
-      result_view->SetRichSuggestionImage(
-          gfx::ImageSkia::CreateFrom1xBitmap(*bitmap));
+    // Enterprise search aggregator matches use the bitmap for icon setting in
+    // `OmniboxResultView::ApplyThemeAndRefreshIcons()`. Only set the image here
+    // for other match types.
+    if (!AutocompleteMatch::IsFeaturedEnterpriseSearchType(match.type)) {
+      const SkBitmap* bitmap = model()->GetPopupRichSuggestionBitmap(i);
+      if (bitmap) {
+        result_view->SetRichSuggestionImage(
+            gfx::ImageSkia::CreateFrom1xBitmap(*bitmap));
+      }
     }
   }
   // If we have more views than matches, hide the surplus ones.
diff --git a/components/omnibox/browser/autocomplete_controller.h b/components/omnibox/browser/autocomplete_controller.h
index 2e75d27..3ba07a75 100644
--- a/components/omnibox/browser/autocomplete_controller.h
+++ b/components/omnibox/browser/autocomplete_controller.h
@@ -367,6 +367,9 @@
   FRIEND_TEST_ALL_PREFIXES(OmniboxPopupViewViewsTest,
                            AccessibleSelectionOnResultSelection);
   FRIEND_TEST_ALL_PREFIXES(OmniboxPopupViewViewsTest, AccessibleResultName);
+  FRIEND_TEST_ALL_PREFIXES(
+      OmniboxEditModelPopupTest,
+      GetMatchIconForFeaturedEnterpriseSearchAggregatorUsesDoesNotUseFavicon);
 
   // A minimal representation of the previous `AutocompleteResult`. Used by
   // `UpdateResult()`'s helper methods.
diff --git a/components/omnibox/browser/omnibox_edit_model.cc b/components/omnibox/browser/omnibox_edit_model.cc
index da740f1..001267ef7 100644
--- a/components/omnibox/browser/omnibox_edit_model.cc
+++ b/components/omnibox/browser/omnibox_edit_model.cc
@@ -1751,12 +1751,24 @@
     auto on_icon_fetched =
         base::BindOnce(&OmniboxEditModel::OnFaviconFetched,
                        weak_factory_.GetWeakPtr(), match.destination_url);
-    favicon =
-        (turl && AutocompleteMatch::IsFeaturedEnterpriseSearchType(match.type))
-            ? controller_->client()->GetFaviconForKeywordSearchProvider(
-                  turl, std::move(on_icon_fetched))
-            : controller_->client()->GetFaviconForPageUrl(
-                  match.destination_url, std::move(on_icon_fetched));
+    if (turl && AutocompleteMatch::IsFeaturedEnterpriseSearchType(match.type)) {
+      // `IsFeaturedEnterpriseSearchType()` also includes site search matches.
+      // We should only use the bitmap for enterprise search aggregators.
+      // Otherwise, use the favicon for the keyword search provider.
+      if (turl->policy_origin() ==
+          TemplateURLData::PolicyOrigin::kSearchAggregator) {
+        const SkBitmap* bitmap = GetPopupRichSuggestionBitmap(match.keyword);
+        if (bitmap) {
+          favicon = gfx::Image(gfx::ImageSkia::CreateFrom1xBitmap(*bitmap));
+        }
+      } else {
+        favicon = controller_->client()->GetFaviconForKeywordSearchProvider(
+            turl, std::move(on_icon_fetched));
+      }
+    } else {
+      favicon = controller_->client()->GetFaviconForPageUrl(
+          match.destination_url, std::move(on_icon_fetched));
+    }
 
     // Extension icons are the correct size for non-touch UI but need to be
     // adjusted to be the correct size for touch mode.
@@ -2846,3 +2858,16 @@
     NOTREACHED();
   }
 }
+
+const SkBitmap* OmniboxEditModel::GetPopupRichSuggestionBitmap(
+    const std::u16string& keyword) const {
+  DCHECK(popup_view_);
+
+  for (size_t i = 0; i < autocomplete_controller()->result().size(); ++i) {
+    auto& result_match = autocomplete_controller()->result().match_at(i);
+    if (result_match.keyword == keyword) {
+      return GetPopupRichSuggestionBitmap(i);
+    }
+  }
+  return nullptr;
+}
diff --git a/components/omnibox/browser/omnibox_edit_model.h b/components/omnibox/browser/omnibox_edit_model.h
index a664064..c76d2ba 100644
--- a/components/omnibox/browser/omnibox_edit_model.h
+++ b/components/omnibox/browser/omnibox_edit_model.h
@@ -610,6 +610,11 @@
   // data source, this should not be called when there's no view.
   std::u16string GetText() const;
 
+  // Lookup the bitmap for `match` based on keyword. Returns nullptr if not
+  // found.
+  const SkBitmap* GetPopupRichSuggestionBitmap(
+      const std::u16string& keyword) const;
+
   // Owns this.
   raw_ptr<OmniboxController> controller_;
 
diff --git a/components/omnibox/browser/omnibox_edit_model_unittest.cc b/components/omnibox/browser/omnibox_edit_model_unittest.cc
index 7b0ead00..0286ec13 100644
--- a/components/omnibox/browser/omnibox_edit_model_unittest.cc
+++ b/components/omnibox/browser/omnibox_edit_model_unittest.cc
@@ -21,6 +21,7 @@
 #include "components/omnibox/browser/actions/tab_switch_action.h"
 #include "components/omnibox/browser/autocomplete_controller.h"
 #include "components/omnibox/browser/autocomplete_match.h"
+#include "components/omnibox/browser/fake_autocomplete_controller.h"
 #include "components/omnibox/browser/omnibox_controller.h"
 #include "components/omnibox/browser/omnibox_popup_view.h"
 #include "components/omnibox/browser/omnibox_prefs.h"
@@ -39,8 +40,10 @@
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/metrics_proto/omnibox_event.pb.h"
 #include "third_party/omnibox_proto/answer_type.pb.h"
+#include "third_party/skia/include/core/SkBitmap.h"
 #include "ui/base/window_open_disposition.h"
 #include "ui/gfx/geometry/rect.h"
+#include "ui/gfx/image/image_unittest_util.h"
 
 using metrics::OmniboxEventProto;
 using Selection = OmniboxPopupSelection;
@@ -1323,6 +1326,105 @@
   EXPECT_EQ(FeedbackType::kNone, result->match_at(1)->feedback_type);
 }
 
+#if !BUILDFLAG(IS_ANDROID) && !BUILDFLAG(IS_IOS)
+// Tests the `GetMatchIcon()` method, verifying that a page favicon is used for
+// `URL_WHAT_YOU_TYPED` matches.
+TEST_F(OmniboxEditModelPopupTest,
+       GetMatchIconForUrlWhatYouTypedUsesPageFavicon) {
+  const GURL kUrl("https://foo.com");
+
+  GURL page_url;
+  EXPECT_CALL(*client(), GetFaviconForPageUrl(_, _))
+      .WillOnce(DoAll(SaveArg<0>(&page_url), Return(gfx::Image())));
+  EXPECT_CALL(*client(), GetFaviconForKeywordSearchProvider(_, _)).Times(0);
+
+  AutocompleteMatch match;
+  match.type = AutocompleteMatchType::URL_WHAT_YOU_TYPED;
+  match.destination_url = kUrl;
+
+  gfx::Image image = model()->GetMatchIcon(match, 0);
+  EXPECT_EQ(page_url, kUrl);
+}
+
+// Tests the `GetMatchIcon()` method, verifying that a keyword favicon is used
+// for `FEATURED_ENTERPRISE_SEARCH` matches with `kSiteSearch` policy origin.
+TEST_F(OmniboxEditModelPopupTest,
+       GetMatchIconForFeaturedEnterpriseSiteSearchUsesKeywordFavicon) {
+  SkBitmap bitmap;
+  bitmap.allocN32Pixels(16, 16);
+  bitmap.eraseColor(SK_ColorRED);
+  gfx::Image expected_image =
+      gfx::Image(gfx::ImageSkia::CreateFrom1xBitmap(bitmap));
+
+  EXPECT_CALL(*client(), GetFaviconForPageUrl(_, _)).Times(0);
+  EXPECT_CALL(*client(), GetFaviconForKeywordSearchProvider(_, _))
+      .WillOnce(Return(expected_image));
+
+  TemplateURLData data;
+  data.SetKeyword(u"sitesearch");
+  data.SetURL("https://sitesearch.com");
+  data.featured_by_policy = true;
+  data.policy_origin = TemplateURLData::PolicyOrigin::kSiteSearch;
+  TemplateURL* turl = controller()->client()->GetTemplateURLService()->Add(
+      std::make_unique<TemplateURL>(data));
+  ASSERT_TRUE(turl);
+
+  AutocompleteMatch match;
+  match.type = AutocompleteMatchType::FEATURED_ENTERPRISE_SEARCH;
+  match.destination_url = GURL("https://sitesearch.com");
+  match.keyword = u"sitesearch";
+  match.associated_keyword = std::make_unique<AutocompleteMatch>(match);
+
+  gfx::Image image = model()->GetMatchIcon(match, 0);
+  gfx::test::CheckColors(bitmap.getColor(0, 0),
+                         image.ToSkBitmap()->getColor(0, 0));
+}
+
+// Tests the `GetMatchIcon()` method, verifying that no favicon is used for
+// `FEATURED_ENTERPRISE_SEARCH` matches with `kSearchAggregator` policy origin.
+TEST_F(OmniboxEditModelPopupTest,
+       GetMatchIconForFeaturedEnterpriseSearchAggregatorUsesDoesNotUseFavicon) {
+  SkBitmap bitmap;
+  bitmap.allocN32Pixels(16, 16);
+  bitmap.eraseColor(SK_ColorRED);
+
+  EXPECT_CALL(*client(), GetFaviconForPageUrl(_, _)).Times(0);
+  EXPECT_CALL(*client(), GetFaviconForKeywordSearchProvider(_, _)).Times(0);
+
+  TemplateURLData data;
+  data.SetKeyword(u"searchaggregator");
+  data.SetURL("https://searchaggregator.com");
+  data.featured_by_policy = true;
+  data.policy_origin = TemplateURLData::PolicyOrigin::kSearchAggregator;
+  TemplateURL* turl = controller()->client()->GetTemplateURLService()->Add(
+      std::make_unique<TemplateURL>(data));
+  ASSERT_TRUE(turl);
+
+  // Creates a set of matches.
+  ACMatches matches;
+  AutocompleteMatch search_aggregator_match(
+      nullptr, 1350, false, AutocompleteMatchType::FEATURED_ENTERPRISE_SEARCH);
+  search_aggregator_match.keyword = u"searchaggregator";
+  search_aggregator_match.associated_keyword =
+      std::make_unique<AutocompleteMatch>(search_aggregator_match);
+  matches.push_back(search_aggregator_match);
+  AutocompleteMatch url_match(nullptr, 1000, false,
+                              AutocompleteMatchType::URL_WHAT_YOU_TYPED);
+  url_match.keyword = u"match";
+  matches.push_back(url_match);
+  AutocompleteResult* result =
+      &controller()->autocomplete_controller()->published_result_;
+  result->AppendMatches(matches);
+
+  // Sets the popup rich suggestion bitmap for search aggregator match.
+  model()->SetPopupRichSuggestionBitmap(0, bitmap);
+
+  gfx::Image image = model()->GetMatchIcon(search_aggregator_match, 0);
+  gfx::test::CheckColors(bitmap.getColor(0, 0),
+                         image.ToSkBitmap()->getColor(0, 0));
+}
+#endif
+
 TEST_F(OmniboxEditModelTest, OmniboxEscapeHistogram) {
   // Escape should incrementally revert temporary text, close the popup, clear
   // input, and blur the omnibox.
diff --git a/components/omnibox/browser/test_omnibox_client.cc b/components/omnibox/browser/test_omnibox_client.cc
index 766f08e5..c3b4c96d 100644
--- a/components/omnibox/browser/test_omnibox_client.cc
+++ b/components/omnibox/browser/test_omnibox_client.cc
@@ -93,6 +93,13 @@
   return gfx::Image(gfx::ImageSkia::CreateFrom1xBitmap(bitmap));
 }
 
+gfx::Image TestOmniboxClient::GetSizedIcon(const gfx::Image& icon) const {
+  if (icon.IsEmpty()) {
+    return gfx::Image();
+  }
+  return gfx::Image(*icon.ToImageSkia());
+}
+
 std::u16string TestOmniboxClient::GetFormattedFullURL() const {
   return location_bar_model_.GetFormattedFullURL();
 }
diff --git a/components/omnibox/browser/test_omnibox_client.h b/components/omnibox/browser/test_omnibox_client.h
index f52922f..6fc88d835 100644
--- a/components/omnibox/browser/test_omnibox_client.h
+++ b/components/omnibox/browser/test_omnibox_client.h
@@ -46,6 +46,7 @@
   bool IsUsingFakeHttpsForHttpsUpgradeTesting() const override;
   gfx::Image GetSizedIcon(const gfx::VectorIcon& vector_icon_type,
                           SkColor vector_icon_color) const override;
+  gfx::Image GetSizedIcon(const gfx::Image& icon) const override;
   std::u16string GetFormattedFullURL() const override;
   std::u16string GetURLForDisplay() const override;
   GURL GetNavigationEntryURL() const override;
@@ -60,6 +61,10 @@
               GetFaviconForPageUrl,
               (const GURL& page_url,
                FaviconFetchedCallback on_favicon_fetched));
+  MOCK_METHOD(gfx::Image,
+              GetFaviconForKeywordSearchProvider,
+              (const TemplateURL* template_url,
+               FaviconFetchedCallback on_favicon_fetched));
   MOCK_METHOD(void,
               ShowFeedbackPage,
               (const std::u16string& input_text, const GURL& destination_url));