Enable browser testing of ShouldTargetNavigation

Adds a function OverrideTargetDecisionForTesting for
OptimizationGuideKeyedService. Clients can use this to control the
desired prediction outcome in tests.

Bug: 1076075
Change-Id: I56ea2c7c1170c0b305b8f0870e22664460fd0806
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2233300Reviewed-by: Sophie Chang <sophiechang@chromium.org>
Reviewed-by: Michael Crouse <mcrouse@chromium.org>
Commit-Queue: Jia Li <jiakl@google.com>
Cr-Commit-Position: refs/heads/master@{#776327}

chrome/browser/BUILD.gn

@@ -1017,6 +1017,8 @@ static_library("browser") {
     "optimization_guide/optimization_guide_permissions_util.h",
     "optimization_guide/optimization_guide_session_statistic.cc",
     "optimization_guide/optimization_guide_session_statistic.h",
+    "optimization_guide/optimization_guide_test_util.cc",
+    "optimization_guide/optimization_guide_test_util.h",
     "optimization_guide/optimization_guide_top_host_provider.cc",
     "optimization_guide/optimization_guide_top_host_provider.h",
     "optimization_guide/optimization_guide_util.cc",

chrome/browser/optimization_guide/optimization_guide_hints_manager.cc

@@ -13,6 +13,7 @@
 #include "base/metrics/histogram_functions.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/metrics/histogram_macros_local.h"
+#include "base/notreached.h"
 #include "base/rand_util.h"
 #include "base/sequenced_task_runner.h"
 #include "base/task/post_task.h"
@@ -43,6 +44,7 @@
 #include "components/optimization_guide/optimization_guide_store.h"
 #include "components/optimization_guide/optimization_guide_switches.h"
 #include "components/optimization_guide/optimization_metadata.h"
+#include "components/optimization_guide/proto/models.pb.h"
 #include "components/optimization_guide/top_host_provider.h"
 #include "components/prefs/pref_service.h"
 #include "components/prefs/scoped_user_pref_update.h"
@@ -1304,3 +1306,28 @@ void OptimizationGuideHintsManager::AddHintForTesting(
   hint_cache_->AddHintForTesting(url, std::move(hint));
   PrepareToInvokeRegisteredCallbacks(url);
 }
+
+void OptimizationGuideHintsManager::OverrideTargetDecisionForTesting(
+    optimization_guide::proto::OptimizationTarget optimization_target,
+    optimization_guide::OptimizationGuideDecision optimization_guide_decision) {
+  if (optimization_target !=
+      optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD) {
+    return;
+  }
+
+  // Manipulate ECTs to effectively change the target decision.
+  switch (optimization_guide_decision) {
+    case optimization_guide::OptimizationGuideDecision::kTrue:
+      current_effective_connection_type_ =
+          net::EffectiveConnectionType::EFFECTIVE_CONNECTION_TYPE_SLOW_2G;
+      break;
+    case optimization_guide::OptimizationGuideDecision::kFalse:
+      current_effective_connection_type_ =
+          net::EffectiveConnectionType::EFFECTIVE_CONNECTION_TYPE_4G;
+      break;
+    case optimization_guide::OptimizationGuideDecision::kUnknown:
+      // No way to override for |kUnknown|. Should not be used in tests.
+      NOTREACHED();
+      break;
+  }
+}
\ No newline at end of file

chrome/browser/optimization_guide/optimization_guide_hints_manager.h

@@ -168,6 +168,13 @@ class OptimizationGuideHintsManager
       optimization_guide::proto::OptimizationType optimization_type,
       const base::Optional<optimization_guide::OptimizationMetadata>& metadata);
 
+  // Override the decision returned by |ShouldTargetNavigation|
+  // for |optimization_target|. For testing purposes only.
+  void OverrideTargetDecisionForTesting(
+      optimization_guide::proto::OptimizationTarget optimization_target,
+      optimization_guide::OptimizationGuideDecision
+          optimization_guide_decision);
+
  private:
   FRIEND_TEST_ALL_PREFIXES(OptimizationGuideHintsManagerTest, IsGoogleURL);
   FRIEND_TEST_ALL_PREFIXES(OptimizationGuideHintsManagerFetchingTest,

chrome/browser/optimization_guide/optimization_guide_keyed_service.cc

@@ -20,6 +20,7 @@
 #include "components/leveldb_proto/public/proto_database_provider.h"
 #include "components/optimization_guide/command_line_top_host_provider.h"
 #include "components/optimization_guide/hints_processing_util.h"
+#include "components/optimization_guide/optimization_guide_decider.h"
 #include "components/optimization_guide/optimization_guide_features.h"
 #include "components/optimization_guide/optimization_guide_service.h"
 #include "components/optimization_guide/top_host_provider.h"
@@ -276,3 +277,16 @@ void OptimizationGuideKeyedService::UpdateSessionFCP(base::TimeDelta fcp) {
   if (prediction_manager_)
     prediction_manager_->UpdateFCPSessionStatistics(fcp);
 }
+
+void OptimizationGuideKeyedService::OverrideTargetDecisionForTesting(
+    optimization_guide::proto::OptimizationTarget optimization_target,
+    optimization_guide::OptimizationGuideDecision optimization_guide_decision) {
+  if (prediction_manager_) {
+    prediction_manager_->OverrideTargetDecisionForTesting(
+        optimization_target, optimization_guide_decision);
+  } else {
+    DCHECK(hints_manager_);
+    hints_manager_->OverrideTargetDecisionForTesting(
+        optimization_target, optimization_guide_decision);
+  }
+}

chrome/browser/optimization_guide/optimization_guide_keyed_service.h

@@ -78,6 +78,13 @@ class OptimizationGuideKeyedService
       optimization_guide::proto::OptimizationType optimization_type,
       const base::Optional<optimization_guide::OptimizationMetadata>& metadata);
 
+  // Override the decision returned by |ShouldTargetNavigation|
+  // for |optimization_target|. For testing purposes only.
+  void OverrideTargetDecisionForTesting(
+      optimization_guide::proto::OptimizationTarget optimization_target,
+      optimization_guide::OptimizationGuideDecision
+          optimization_guide_decision);
+
  private:
   friend class ChromeBrowsingDataRemoverDelegate;
   friend class HintsFetcherBrowserTest;

chrome/browser/optimization_guide/optimization_guide_test_util.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/optimization_guide/optimization_guide_test_util.h"
+
+std::unique_ptr<optimization_guide::proto::PredictionModel>
+GetMinimalDecisionTreePredictionModel(double threshold, double weight) {
+  auto prediction_model =
+      std::make_unique<optimization_guide::proto::PredictionModel>();
+  prediction_model->mutable_model()->mutable_threshold()->set_value(threshold);
+  optimization_guide::proto::DecisionTree decision_tree_model;
+  decision_tree_model.set_weight(weight);
+
+  optimization_guide::proto::TreeNode* tree_node =
+      decision_tree_model.add_nodes();
+  tree_node->mutable_node_id()->set_value(0);
+
+  *prediction_model->mutable_model()->mutable_decision_tree() =
+      decision_tree_model;
+
+  return prediction_model;
+}
+
+std::unique_ptr<optimization_guide::proto::PredictionModel>
+GetSingleLeafDecisionTreePredictionModel(double threshold,
+                                         double weight,
+                                         double leaf_value) {
+  auto prediction_model =
+      GetMinimalDecisionTreePredictionModel(threshold, weight);
+  prediction_model->mutable_model()
+      ->mutable_decision_tree()
+      ->mutable_nodes(0)
+      ->mutable_leaf()
+      ->mutable_vector()
+      ->add_value()
+      ->set_double_value(leaf_value);
+  return prediction_model;
+}
\ No newline at end of file

chrome/browser/optimization_guide/optimization_guide_test_util.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_OPTIMIZATION_GUIDE_OPTIMIZATION_GUIDE_TEST_UTIL_H_
+#define CHROME_BROWSER_OPTIMIZATION_GUIDE_OPTIMIZATION_GUIDE_TEST_UTIL_H_
+
+#include <memory>
+
+#include "components/optimization_guide/proto/models.pb.h"
+
+// Returns a decision tree model with |threshold|, |weight|, and a single
+// uninitialized node.
+std::unique_ptr<optimization_guide::proto::PredictionModel>
+GetMinimalDecisionTreePredictionModel(double threshold, double weight);
+
+// Returns a decision tree model with |threshold|, |weight|, and a single
+// leaf node with |leaf_value|.
+std::unique_ptr<optimization_guide::proto::PredictionModel>
+GetSingleLeafDecisionTreePredictionModel(double threshold,
+                                         double weight,
+                                         double leaf_value);
+
+#endif  // CHROME_BROWSER_OPTIMIZATION_GUIDE_OPTIMIZATION_GUIDE_TEST_UTIL_H_
\ No newline at end of file

chrome/browser/optimization_guide/prediction/prediction_manager.cc

@@ -4,6 +4,7 @@
 
 #include "chrome/browser/optimization_guide/prediction/prediction_manager.h"
 
+#include <memory>
 #include <utility>
 
 #include "base/containers/flat_map.h"
@@ -22,6 +23,7 @@
 #include "chrome/browser/optimization_guide/optimization_guide_navigation_data.h"
 #include "chrome/browser/optimization_guide/optimization_guide_permissions_util.h"
 #include "chrome/browser/optimization_guide/optimization_guide_session_statistic.h"
+#include "chrome/browser/optimization_guide/optimization_guide_test_util.h"
 #include "chrome/browser/optimization_guide/optimization_guide_util.h"
 #include "chrome/browser/optimization_guide/prediction/prediction_model.h"
 #include "chrome/browser/optimization_guide/prediction/prediction_model_fetcher.h"
@@ -887,4 +889,41 @@ PredictionManager::GetHostModelFeaturesForHost(const std::string& host) const {
   return it->second;
 }
 
+void PredictionManager::OverrideTargetDecisionForTesting(
+    proto::OptimizationTarget optimization_target,
+    OptimizationGuideDecision optimization_guide_decision) {
+  auto it = optimization_target_prediction_model_map_.find(optimization_target);
+  if (it != optimization_target_prediction_model_map_.end())
+    optimization_target_prediction_model_map_.erase(it);
+
+  // No model for |kUnknown|. This will make |ShouldTargetNavigation|
+  // return an |OptimizationTargetDecision::kModelNotAvailableOnClient|,
+  // which in turn yields an |OptimizationGuideDecision::kUnknown| in
+  // |OptimizationGuideKeyedService::ShouldTargetNavigation|.
+  if (optimization_guide_decision == OptimizationGuideDecision::kUnknown)
+    return;
+
+  // Construct a simple model that will return the provided
+  // |optimization_guide_decision|.
+  const double threshold = 5.0;
+  const double weight = 1.0;
+  double leaf_value =
+      (optimization_guide_decision == OptimizationGuideDecision::kTrue)
+          ? threshold + 1.0  // Value is greater than |threshold| to get |kTrue|
+          : threshold - 1.0;  // Value is less than |threshold| to get |kFalse|
+
+  std::unique_ptr<proto::PredictionModel> prediction_model =
+      GetSingleLeafDecisionTreePredictionModel(threshold, weight,
+                                               leaf_value / weight);
+
+  proto::ModelInfo* model_info = prediction_model->mutable_model_info();
+
+  model_info->set_version(1);
+  model_info->set_optimization_target(optimization_target);
+  model_info->add_supported_model_types(proto::MODEL_TYPE_DECISION_TREE);
+
+  optimization_target_prediction_model_map_.emplace(
+      optimization_target, CreatePredictionModel(*prediction_model));
+}
+
 }  // namespace optimization_guide

chrome/browser/optimization_guide/prediction/prediction_manager.h

@@ -130,6 +130,12 @@ class PredictionManager
   // from the models and features store.
   void ClearHostModelFeatures();
 
+  // Override the decision returned by |ShouldTargetNavigation|
+  // for |optimization_target|. For testing purposes only.
+  void OverrideTargetDecisionForTesting(
+      proto::OptimizationTarget optimization_target,
+      OptimizationGuideDecision optimization_guide_decision);
+
  protected:
   // Return the prediction model for the optimization target used by this
   // PredictionManager for testing.

chrome/browser/optimization_guide/prediction/prediction_manager_browsertest.cc

@@ -13,6 +13,7 @@
 #include "chrome/browser/optimization_guide/optimization_guide_keyed_service.h"
 #include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
 #include "chrome/browser/optimization_guide/optimization_guide_session_statistic.h"
+#include "chrome/browser/optimization_guide/optimization_guide_test_util.h"
 #include "chrome/browser/optimization_guide/prediction/prediction_manager.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/ui/browser.h"
@@ -71,16 +72,14 @@ void RetryForHistogramUntilCountReached(
 std::unique_ptr<optimization_guide::proto::PredictionModel>
 GetValidDecisionTreePredictionModel() {
   std::unique_ptr<optimization_guide::proto::PredictionModel> prediction_model =
-      std::make_unique<optimization_guide::proto::PredictionModel>();
-  prediction_model->mutable_model()->mutable_threshold()->set_value(5.0);
+      GetMinimalDecisionTreePredictionModel(/* threshold= */ 5.0,
+                                            /* weight= */ 2.0);
 
-  optimization_guide::proto::DecisionTree decision_tree_model =
-      optimization_guide::proto::DecisionTree();
-  decision_tree_model.set_weight(2.0);
+  optimization_guide::proto::DecisionTree* decision_tree_model =
+      prediction_model->mutable_model()->mutable_decision_tree();
 
   optimization_guide::proto::TreeNode* tree_node =
-      decision_tree_model.add_nodes();
-  tree_node->mutable_node_id()->set_value(0);
+      decision_tree_model->mutable_nodes(0);
   tree_node->mutable_binary_node()->mutable_left_child_id()->set_value(1);
   tree_node->mutable_binary_node()->mutable_right_child_id()->set_value(2);
   tree_node->mutable_binary_node()
@@ -96,18 +95,16 @@ GetValidDecisionTreePredictionModel() {
       ->mutable_threshold()
       ->set_float_value(1.0);
 
-  tree_node = decision_tree_model.add_nodes();
+  tree_node = decision_tree_model->add_nodes();
   tree_node->mutable_node_id()->set_value(1);
   tree_node->mutable_leaf()->mutable_vector()->add_value()->set_double_value(
       2.);
 
-  tree_node = decision_tree_model.add_nodes();
+  tree_node = decision_tree_model->add_nodes();
   tree_node->mutable_node_id()->set_value(2);
   tree_node->mutable_leaf()->mutable_vector()->add_value()->set_double_value(
       4.);
 
-  *prediction_model->mutable_model()->mutable_decision_tree() =
-      decision_tree_model;
   return prediction_model;
 }
 

chrome/browser/previews/defer_all_script_browsertest.cc

@@ -3,6 +3,7 @@
 // found in the LICENSE file.
 
 #include <map>
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -14,6 +15,8 @@
 #include "base/test/scoped_feature_list.h"
 #include "build/build_config.h"
 #include "chrome/browser/browser_process.h"
+#include "chrome/browser/optimization_guide/optimization_guide_keyed_service.h"
+#include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
 #include "chrome/browser/previews/previews_service.h"
 #include "chrome/browser/previews/previews_service_factory.h"
 #include "chrome/browser/previews/previews_test_util.h"
@@ -53,6 +56,10 @@ class DeferAllScriptBrowserTest : public InProcessBrowserTest {
         {previews::features::kPreviews,
          previews::features::kDeferAllScriptPreviews,
          optimization_guide::features::kOptimizationHints,
+         // TODO(crbug/1021364): Remove the following two features after the
+         // model rollout
+         optimization_guide::features::kRemoteOptimizationGuideFetching,
+         optimization_guide::features::kOptimizationTargetPrediction,
          features::kBackForwardCache},
         {});
   }
@@ -61,11 +68,12 @@ class DeferAllScriptBrowserTest : public InProcessBrowserTest {
 
   void SetUpOnMainThread() override {
     host_resolver()->AddRule("*", "127.0.0.1");
-    g_browser_process->network_quality_tracker()
-        ->ReportEffectiveConnectionTypeForTesting(
-            net::EFFECTIVE_CONNECTION_TYPE_2G);
-    https_server_.reset(
-        new net::EmbeddedTestServer(net::EmbeddedTestServer::TYPE_HTTPS));
+    OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+        ->OverrideTargetDecisionForTesting(
+            optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+            optimization_guide::OptimizationGuideDecision::kTrue);
+    https_server_ = std::make_unique<net::EmbeddedTestServer>(
+        net::EmbeddedTestServer::TYPE_HTTPS);
     https_server_->ServeFilesFromSourceDirectory("chrome/test/data/previews");
     ASSERT_TRUE(https_server_->Start());
 
@@ -556,11 +564,12 @@ IN_PROC_BROWSER_TEST_F(DeferAllScriptBrowserTest,
   histogram_tester.ExpectBucketCount("Previews.PreviewShown.DeferAllScript",
                                      true, 1);
 
-  // Now adjust the network triggering condition to not choose preview for a
+  // Override the target decision to |kFalse| to not trigger a preview for the
   // new decision.
-  g_browser_process->network_quality_tracker()
-      ->ReportEffectiveConnectionTypeForTesting(
-          net::EFFECTIVE_CONNECTION_TYPE_4G);
+  OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+      ->OverrideTargetDecisionForTesting(
+          optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+          optimization_guide::OptimizationGuideDecision::kFalse);
 
   // Navigate to another host on same tab (to cause previous navigation
   // to be saved in BackForward cache).
@@ -620,11 +629,12 @@ IN_PROC_BROWSER_TEST_F(
   // Wait for initial page load to complete.
   content::WaitForLoadStop(web_contents());
 
-  // Adjust the network triggering condition to not choose preview for this
-  // navigation.
-  g_browser_process->network_quality_tracker()
-      ->ReportEffectiveConnectionTypeForTesting(
-          net::EFFECTIVE_CONNECTION_TYPE_4G);
+  // Override the target decision to |kFalse| to choose not to trigger a
+  // preview this navigation.
+  OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+      ->OverrideTargetDecisionForTesting(
+          optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+          optimization_guide::OptimizationGuideDecision::kFalse);
 
   // Navigate to DeferAllScript url.
   ui_test_utils::NavigateToURL(browser(), url);
@@ -637,11 +647,12 @@ IN_PROC_BROWSER_TEST_F(
       "Previews.EligibilityReason.DeferAllScript",
       static_cast<int>(previews::PreviewsEligibilityReason::COMMITTED), 0);
 
-  // Now adjust the network triggering condition to allow a preview for a
+  // Now override the model decision to |kTrue| to allow a preview for a
   // new decision.
-  g_browser_process->network_quality_tracker()
-      ->ReportEffectiveConnectionTypeForTesting(
-          net::EFFECTIVE_CONNECTION_TYPE_2G);
+  OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+      ->OverrideTargetDecisionForTesting(
+          optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+          optimization_guide::OptimizationGuideDecision::kTrue);
 
   // Navigate to another host on same tab (to cause previous navigation
   // to be saved in BackForward cache).

chrome/browser/previews/previews_browsertest.cc

@@ -12,6 +12,8 @@
 #include "build/build_config.h"
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/metrics/subprocess_metrics_provider.h"
+#include "chrome/browser/optimization_guide/optimization_guide_keyed_service.h"
+#include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
 #include "chrome/browser/previews/previews_service_factory.h"
 #include "chrome/browser/previews/previews_test_util.h"
 #include "chrome/browser/previews/previews_ui_tab_helper.h"
@@ -26,6 +28,7 @@
 #include "components/optimization_guide/optimization_guide_features.h"
 #include "components/optimization_guide/optimization_guide_service.h"
 #include "components/optimization_guide/proto/hints.pb.h"
+#include "components/optimization_guide/proto/models.pb.h"
 #include "components/optimization_guide/test_hints_component_creator.h"
 #include "components/previews/core/previews_features.h"
 #include "components/previews/core/previews_switches.h"
@@ -46,10 +49,6 @@ class PreviewsBrowserTest : public InProcessBrowserTest {
   ~PreviewsBrowserTest() override = default;
 
   void SetUpOnMainThread() override {
-    g_browser_process->network_quality_tracker()
-        ->ReportEffectiveConnectionTypeForTesting(
-            net::EFFECTIVE_CONNECTION_TYPE_2G);
-
     // Set up https server with resource monitor.
     https_server_.reset(
         new net::EmbeddedTestServer(net::EmbeddedTestServer::TYPE_HTTPS));
@@ -70,8 +69,8 @@ class PreviewsBrowserTest : public InProcessBrowserTest {
     ASSERT_EQ(https_hint_setup_url_.host(), https_url_.host());
 
     // Set up http server with resource monitor and redirect handler.
-    http_server_.reset(
-        new net::EmbeddedTestServer(net::EmbeddedTestServer::TYPE_HTTP));
+    http_server_ = std::make_unique<net::EmbeddedTestServer>(
+        net::EmbeddedTestServer::TYPE_HTTP);
     http_server_->ServeFilesFromSourceDirectory("chrome/test/data/previews");
     http_server_->RegisterRequestMonitor(base::BindRepeating(
         &PreviewsBrowserTest::MonitorResourceRequest, base::Unretained(this)));
@@ -207,6 +206,14 @@ class PreviewsNoScriptBrowserTest : public ::testing::WithParamInterface<bool>,
     PreviewsBrowserTest::SetUp();
   }
 
+  void SetUpOnMainThread() override {
+    OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+        ->OverrideTargetDecisionForTesting(
+            optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+            optimization_guide::OptimizationGuideDecision::kTrue);
+    PreviewsBrowserTest::SetUpOnMainThread();
+  }
+
   void SetUpCommandLine(base::CommandLine* cmd) override {
     cmd->AppendSwitch("enable-spdy-proxy-auth");
     cmd->AppendSwitch("optimization-guide-disable-installer");
@@ -412,10 +419,12 @@ IN_PROC_BROWSER_TEST_P(
 IN_PROC_BROWSER_TEST_P(PreviewsNoScriptBrowserTest,
                        DISABLE_ON_WIN_MAC_CHROMEOS(
                            NoScriptPreviewsEnabledShouldSkipPreviewCheck)) {
-  // Set ECT to 4G so that the Preview should not be shown in the regular case.
-  g_browser_process->network_quality_tracker()
-      ->ReportEffectiveConnectionTypeForTesting(
-          net::EFFECTIVE_CONNECTION_TYPE_4G);
+  // Override the decision to |kFalse| so that the Preview should not be shown
+  // in the regular case.
+  OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+      ->OverrideTargetDecisionForTesting(
+          optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+          optimization_guide::OptimizationGuideDecision::kFalse);
 
   GURL url = https_url();