Add utility for consumers to override the model file that gets returned to them

Bug: 1146151
Change-Id: Ib10da995af2d6d7b91ac08e06460544a525f4e7e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2567184
Commit-Queue: Sophie Chang <sophiechang@chromium.org>
Reviewed-by: Michael Crouse <mcrouse@chromium.org>
Cr-Commit-Position: refs/heads/master@{#833165}

chrome/browser/optimization_guide/optimization_guide_keyed_service.cc

@@ -338,3 +338,12 @@ void OptimizationGuideKeyedService::OverrideTargetDecisionForTesting(
         optimization_target, optimization_guide_decision);
   }
 }
+
+void OptimizationGuideKeyedService::OverrideTargetModelFileForTesting(
+    optimization_guide::proto::OptimizationTarget optimization_target,
+    const base::FilePath& file_path) {
+  if (prediction_manager_) {
+    prediction_manager_->OverrideTargetModelFileForTesting(optimization_target,
+                                                           file_path);
+  }
+}

chrome/browser/optimization_guide/optimization_guide_keyed_service.h

@@ -96,6 +96,12 @@ class OptimizationGuideKeyedService
       optimization_guide::OptimizationGuideDecision
           optimization_guide_decision);
 
+  // Override the model file sent to observers of |optimization_target|. For
+  // testing purposes only.
+  void OverrideTargetModelFileForTesting(
+      optimization_guide::proto::OptimizationTarget optimization_target,
+      const base::FilePath& file_path);
+
  private:
   friend class ChromeBrowsingDataRemoverDelegate;
   friend class HintsFetcherBrowserTest;

chrome/browser/optimization_guide/prediction/prediction_manager.cc

@@ -1201,4 +1201,22 @@ void PredictionManager::OverrideTargetDecisionForTesting(
       optimization_target, CreatePredictionModel(*prediction_model));
 }
 
+void PredictionManager::OverrideTargetModelFileForTesting(
+    proto::OptimizationTarget optimization_target,
+    const base::FilePath& file_path) {
+  proto::PredictionModel prediction_model;
+  prediction_model.mutable_model_info()->set_version(1);
+  prediction_model.mutable_model_info()->set_optimization_target(
+      optimization_target);
+  SetFilePathInPredictionModel(file_path, &prediction_model);
+  std::unique_ptr<PredictionModelFile> prediction_model_file =
+      PredictionModelFile::Create(prediction_model);
+  DCHECK(prediction_model_file);
+
+  optimization_target_prediction_model_file_map_.insert_or_assign(
+      optimization_target, std::move(prediction_model_file));
+
+  NotifyObserversOfNewModelPath(optimization_target, file_path);
+}
+
 }  // namespace optimization_guide

chrome/browser/optimization_guide/prediction/prediction_manager.h

@@ -183,6 +183,12 @@ class PredictionManager
       proto::OptimizationTarget optimization_target,
       OptimizationGuideDecision optimization_guide_decision);
 
+  // Override the model file returned to observers for |optimization_target|.
+  // For testing purposes only.
+  void OverrideTargetModelFileForTesting(
+      proto::OptimizationTarget optimization_target,
+      const base::FilePath& file_path);
+
   // PredictionModelDownloadObserver:
   void OnModelReady(const proto::PredictionModel& model) override;