Add initial browser tests for downloading models

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

chrome/browser/optimization_guide/prediction/prediction_manager_browsertest.cc

@@ -12,15 +12,19 @@
 #include "base/test/scoped_feature_list.h"
 #include "build/build_config.h"
 #include "chrome/browser/browser_process.h"
+#include "chrome/browser/download/download_service_factory.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/optimization_guide/optimization_guide_session_statistic.h"
 #include "chrome/browser/optimization_guide/prediction/prediction_manager.h"
 #include "chrome/browser/profiles/profile.h"
+#include "chrome/browser/profiles/profile_key.h"
 #include "chrome/browser/ui/browser.h"
 #include "chrome/test/base/in_process_browser_test.h"
 #include "chrome/test/base/ui_test_utils.h"
 #include "components/data_reduction_proxy/core/common/data_reduction_proxy_switches.h"
+#include "components/download/public/background_service/download_service.h"
+#include "components/download/public/background_service/logger.h"
 #include "components/metrics/content/subprocess_metrics_provider.h"
 #include "components/optimization_guide/optimization_guide_constants.h"
 #include "components/optimization_guide/optimization_guide_features.h"
@@ -40,6 +44,9 @@
 
 namespace {
 
+const char kOptimizationGuidePredictionModelsClient[] =
+    "OptimizationGuidePredictionModels";
+
 // Fetch and calculate the total number of samples from all the bins for
 // |histogram_name|. Note: from some browertests run (such as chromeos) there
 // might be two profiles created, and this will return the total sample count
@@ -181,7 +188,9 @@ enum class PredictionModelsFetcherRemoteResponseType {
   kSuccessfulWithModelsAndFeatures = 0,
   kSuccessfulWithFeaturesAndNoModels = 1,
   kSuccessfulWithModelsAndNoFeatures = 2,
-  kUnsuccessful = 3,
+  kSuccessfulWithValidModelFile = 3,
+  kSuccessfulWithInvalidModelFile = 4,
+  kUnsuccessful = 5,
 };
 
 // A WebContentsObserver that asks whether an optimization target can be
@@ -241,7 +250,8 @@ class PredictionManagerBrowserTestBase : public InProcessBrowserTest {
 
     models_server_ = std::make_unique<net::EmbeddedTestServer>(
         net::EmbeddedTestServer::TYPE_HTTPS);
-    models_server_->ServeFilesFromSourceDirectory("chrome/test/data/previews");
+    models_server_->ServeFilesFromSourceDirectory(
+        "chrome/test/data/optimization_guide");
     models_server_->RegisterRequestHandler(base::BindRepeating(
         &PredictionManagerBrowserTestBase::HandleGetModelsRequest,
         base::Unretained(this)));
@@ -259,6 +269,7 @@ class PredictionManagerBrowserTestBase : public InProcessBrowserTest {
     ASSERT_TRUE(https_server_->Start());
     https_url_with_content_ = https_server_->GetURL("/english_page.html");
     https_url_without_content_ = https_server_->GetURL("/empty.html");
+    model_file_url_ = models_server_->GetURL("/unsignedmodel.crx3");
 
     // Set up an OptimizationGuideKeyedService consumer.
     consumer_ = std::make_unique<OptimizationGuideConsumerWebContentsObserver>(
@@ -355,6 +366,9 @@ class PredictionManagerBrowserTestBase : public InProcessBrowserTest {
  private:
   std::unique_ptr<net::test_server::HttpResponse> HandleGetModelsRequest(
       const net::test_server::HttpRequest& request) {
+    if (request.GetURL() == model_file_url_)
+      return std::unique_ptr<net::test_server::HttpResponse>();
+
     std::unique_ptr<net::test_server::BasicHttpResponse> response;
 
     response = std::make_unique<net::test_server::BasicHttpResponse>();
@@ -388,6 +402,14 @@ class PredictionManagerBrowserTestBase : public InProcessBrowserTest {
     } else if (response_type_ == PredictionModelsFetcherRemoteResponseType::
                                      kSuccessfulWithModelsAndNoFeatures) {
       get_models_response->clear_host_model_features();
+    } else if (response_type_ == PredictionModelsFetcherRemoteResponseType::
+                                     kSuccessfulWithInvalidModelFile) {
+      get_models_response->mutable_models(0)->mutable_model()->set_download_url(
+          https_url_with_content_.spec());
+    } else if (response_type_ == PredictionModelsFetcherRemoteResponseType::
+                                     kSuccessfulWithValidModelFile) {
+      get_models_response->mutable_models(0)->mutable_model()->set_download_url(
+          model_file_url_.spec());
     } else if (response_type_ ==
                PredictionModelsFetcherRemoteResponseType::kUnsuccessful) {
       response->set_code(net::HTTP_NOT_FOUND);
@@ -399,6 +421,7 @@ class PredictionManagerBrowserTestBase : public InProcessBrowserTest {
     return std::move(response);
   }
 
+  GURL model_file_url_;
   GURL https_url_with_content_, https_url_without_content_;
   std::unique_ptr<net::EmbeddedTestServer> https_server_;
   std::unique_ptr<net::EmbeddedTestServer> models_server_;
@@ -409,7 +432,6 @@ class PredictionManagerBrowserTestBase : public InProcessBrowserTest {
   base::flat_set<uint32_t> expected_field_trial_name_hashes_;
 };
 
-// Parametrized on whether the ML Service path is enabled.
 class PredictionManagerBrowserTest : public PredictionManagerBrowserTestBase {
  public:
   PredictionManagerBrowserTest() = default;
@@ -730,4 +752,187 @@ IN_PROC_BROWSER_TEST_F(
   run_loop->Run();
 }
 
+// Implementation of a download system logger that provides the ability to wait
+// for certain events to happen, notably added and progressing downloads.
+class DownloadServiceObserver : public download::Logger::Observer {
+ public:
+  using DownloadCompletedCallback = base::OnceCallback<void()>;
+
+  DownloadServiceObserver() = default;
+  ~DownloadServiceObserver() override = default;
+
+  // Sets |callback| to be invoked when a download has completed.
+  void set_download_completed_callback(DownloadCompletedCallback callback) {
+    download_completed_callback_ = std::move(callback);
+  }
+
+  // download::Logger::Observer implementation:
+  void OnServiceStatusChanged(const base::Value& service_status) override {}
+  void OnServiceDownloadsAvailable(
+      const base::Value& service_downloads) override {}
+  void OnServiceDownloadFailed(const base::Value& service_download) override {}
+  void OnServiceRequestMade(const base::Value& service_request) override {}
+  void OnServiceDownloadChanged(const base::Value& service_download) override {
+    const std::string& client = service_download.FindKey("client")->GetString();
+    const std::string& state = service_download.FindKey("state")->GetString();
+
+    if (client != kOptimizationGuidePredictionModelsClient)
+      return;
+
+    if (state == "COMPLETE" && download_completed_callback_)
+      std::move(download_completed_callback_).Run();
+  }
+
+ private:
+  DownloadCompletedCallback download_completed_callback_;
+};
+
+class ModelFileObserver : public OptimizationTargetModelObserver {
+ public:
+  using ModelFileReceivedCallback =
+      base::OnceCallback<void(proto::OptimizationTarget,
+                              const base::FilePath&)>;
+
+  ModelFileObserver() = default;
+  ~ModelFileObserver() override = default;
+
+  void set_model_file_received_callback(ModelFileReceivedCallback callback) {
+    file_received_callback_ = std::move(callback);
+  }
+
+  void OnModelFileUpdated(proto::OptimizationTarget optimization_target,
+                          const base::FilePath& file_path) override {
+    if (file_received_callback_)
+      std::move(file_received_callback_).Run(optimization_target, file_path);
+  }
+
+ private:
+  ModelFileReceivedCallback file_received_callback_;
+};
+
+class PredictionManagerModelDownloadingBrowserTest
+    : public PredictionManagerBrowserTest {
+ public:
+  PredictionManagerModelDownloadingBrowserTest() = default;
+  ~PredictionManagerModelDownloadingBrowserTest() override = default;
+
+  void SetUpOnMainThread() override {
+    download_service_observer_ = std::make_unique<DownloadServiceObserver>();
+    DownloadServiceFactory::GetForKey(browser()->profile()->GetProfileKey())
+        ->GetLogger()
+        ->AddObserver(download_service_observer_.get());
+
+    model_file_observer_ = std::make_unique<ModelFileObserver>();
+
+    PredictionManagerBrowserTest::SetUpOnMainThread();
+  }
+
+  void TearDownOnMainThread() override {
+    PredictionManagerBrowserTest::TearDownOnMainThread();
+
+    DownloadServiceFactory::GetForKey(browser()->profile()->GetProfileKey())
+        ->GetLogger()
+        ->RemoveObserver(download_service_observer_.get());
+  }
+
+  DownloadServiceObserver* download_observer() {
+    return download_service_observer_.get();
+  }
+
+  ModelFileObserver* model_file_observer() {
+    return model_file_observer_.get();
+  }
+
+  void RegisterModelFileObserverWithKeyedService() {
+    OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
+        ->AddObserverForOptimizationTargetModel(
+            proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD,
+            model_file_observer_.get());
+  }
+
+ private:
+  void InitializeFeatureList() override {
+    scoped_feature_list_.InitWithFeaturesAndParameters(
+        {
+            {features::kOptimizationHints, {}},
+            {features::kRemoteOptimizationGuideFetching, {}},
+            {features::kOptimizationTargetPrediction, {}},
+            {features::kOptimizationGuideModelDownloading,
+             {{"unrestricted_model_downloading", "true"}}},
+        },
+        {});
+  }
+
+  std::unique_ptr<DownloadServiceObserver> download_service_observer_;
+  std::unique_ptr<ModelFileObserver> model_file_observer_;
+};
+
+IN_PROC_BROWSER_TEST_F(
+    PredictionManagerModelDownloadingBrowserTest,
+    DISABLE_ON_WIN_MAC_CHROMEOS(
+        TestDownloadUrlAcceptedByDownloadServiceButInvalid)) {
+  base::HistogramTester histogram_tester;
+
+  SetResponseType(PredictionModelsFetcherRemoteResponseType::
+                      kSuccessfulWithInvalidModelFile);
+
+  std::unique_ptr<base::RunLoop> completed_run_loop =
+      std::make_unique<base::RunLoop>();
+  download_observer()->set_download_completed_callback(
+      base::BindOnce([](base::RunLoop* run_loop) { run_loop->Quit(); },
+                     completed_run_loop.get()));
+  // Registering should initiate the fetch and receive a response with a model
+  // containing a download URL and then subsequently downloaded.
+  RegisterModelFileObserverWithKeyedService();
+
+  // Wait until the download has completed.
+  completed_run_loop->Run();
+
+  histogram_tester.ExpectUniqueSample(
+      "OptimizationGuide.PredictionModelDownloadManager.DownloadStatus",
+      PredictionModelDownloadStatus::kFailedCrxVerification, 1);
+  // An unverified file should not notify us that it's ready.
+  histogram_tester.ExpectTotalCount(
+      "OptimizationGuide.PredictionModelUpdateVersion.PainfulPageLoad", 0);
+}
+
+IN_PROC_BROWSER_TEST_F(
+    PredictionManagerModelDownloadingBrowserTest,
+    DISABLE_ON_WIN_MAC_CHROMEOS(TestSuccessfulModelFileFlow)) {
+  // TODO(crbug/1146151): Remove this switch once we can produce a signed model
+  // file.
+  base::CommandLine::ForCurrentProcess()->AppendSwitch(
+      switches::kDisableModelDownloadVerificationForTesting);
+
+  base::HistogramTester histogram_tester;
+
+  SetResponseType(
+      PredictionModelsFetcherRemoteResponseType::kSuccessfulWithValidModelFile);
+
+  std::unique_ptr<base::RunLoop> run_loop = std::make_unique<base::RunLoop>();
+  model_file_observer()->set_model_file_received_callback(base::BindOnce(
+      [](base::RunLoop* run_loop, proto::OptimizationTarget optimization_target,
+         const base::FilePath& file_path) {
+        EXPECT_EQ(optimization_target,
+                  proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
+        run_loop->Quit();
+      },
+      run_loop.get()));
+
+  // Registering should initiate the fetch and receive a response with a model
+  // containing a download URL and then subsequently downloaded.
+  RegisterModelFileObserverWithKeyedService();
+
+  // Wait until the observer receives the file.
+  run_loop->Run();
+
+  histogram_tester.ExpectUniqueSample(
+      "OptimizationGuide.PredictionModelDownloadManager.DownloadStatus",
+      PredictionModelDownloadStatus::kSuccess, 1);
+  histogram_tester.ExpectUniqueSample(
+      "OptimizationGuide.PredictionModelUpdateVersion.PainfulPageLoad", 123, 1);
+  histogram_tester.ExpectUniqueSample(
+      "OptimizationGuide.PredictionModelLoadedVersion.PainfulPageLoad", 123, 1);
+}
+
 }  // namespace optimization_guide

chrome/browser/optimization_guide/prediction/prediction_model_download_manager.cc

@@ -24,6 +24,7 @@
 #include "components/download/public/background_service/download_service.h"
 #include "components/optimization_guide/optimization_guide_enums.h"
 #include "components/optimization_guide/optimization_guide_features.h"
+#include "components/optimization_guide/optimization_guide_switches.h"
 #include "components/optimization_guide/optimization_guide_util.h"
 #include "components/services/unzip/content/unzip_service.h"
 #include "components/services/unzip/public/cpp/unzip.h"
@@ -214,7 +215,7 @@ PredictionModelDownloadManager::ProcessDownload(
     const base::FilePath& file_path) {
   DCHECK(background_task_runner_->RunsTasksInCurrentSequence());
 
-  if (should_verify_download_) {
+  if (!switches::ShouldSkipModelDownloadVerificationForTesting()) {
     // Verify that the |file_path| contains a file signed with a key we trust.
     crx_file::VerifierResult verifier_result = crx_file::Verify(
         file_path, crx_file::VerifierFormat::CRX3_WITH_PUBLISHER_PROOF,
@@ -269,7 +270,7 @@ void PredictionModelDownloadManager::OnDownloadUnzipped(
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
   // Clean up original download file when this function finishes.
-  base::SequencedTaskRunnerHandle::Get()->PostTask(
+  background_task_runner_->PostTask(
       FROM_HERE,
       base::BindOnce(base::GetDeleteFileCallback(), original_file_path));
 
@@ -351,8 +352,4 @@ void PredictionModelDownloadManager::NotifyModelReady(
     observer.OnModelReady(*model);
 }
 
-void PredictionModelDownloadManager::TurnOffVerificationForTesting() {
-  should_verify_download_ = false;
-}
-
 }  // namespace optimization_guide

chrome/browser/optimization_guide/prediction/prediction_model_download_manager.h

@@ -116,9 +116,6 @@ class PredictionModelDownloadManager {
   // Must be invoked on the UI thread.
   void NotifyModelReady(const base::Optional<proto::PredictionModel>& model);
 
-  // Turns off CRX3 verification for testing.
-  void TurnOffVerificationForTesting();
-
   // The set of GUIDs that are still pending download.
   std::set<std::string> pending_download_guids_;
 

chrome/browser/optimization_guide/prediction/prediction_model_download_manager_unittest.cc

@@ -21,6 +21,7 @@
 #include "components/download/public/background_service/test/mock_download_service.h"
 #include "components/optimization_guide/optimization_guide_enums.h"
 #include "components/optimization_guide/optimization_guide_features.h"
+#include "components/optimization_guide/optimization_guide_switches.h"
 #include "components/optimization_guide/optimization_guide_util.h"
 #include "components/services/unzip/content/unzip_service.h"
 #include "components/services/unzip/in_process_unzipper.h"
@@ -158,7 +159,8 @@ class PredictionModelDownloadManagerTest
   }
 
   void TurnOffDownloadVerification() {
-    download_manager_->TurnOffVerificationForTesting();
+    base::CommandLine::ForCurrentProcess()->AppendSwitch(
+        switches::kDisableModelDownloadVerificationForTesting);
   }
 
  private:

chrome/common/chrome_paths.cc

@@ -565,6 +565,7 @@ bool PathProvider(int key, base::FilePath* result) {
       if (!base::PathService::Get(chrome::DIR_USER_DATA, &cur))
         return false;
       cur = cur.Append(FILE_PATH_LITERAL("OptimizationGuidePredictionModels"));
+      create_dir = true;
       break;
 
     default:

components/optimization_guide/optimization_guide_switches.cc

@@ -66,6 +66,9 @@ const char kDisableFetchingHintsAtNavigationStartForTesting[] =
 const char kDisableCheckingUserPermissionsForTesting[] =
     "disable-checking-optimization-guide-user-permissions";
 
+const char kDisableModelDownloadVerificationForTesting[] =
+    "disable-model-download-verification";
+
 bool IsHintComponentProcessingDisabled() {
   return base::CommandLine::ForCurrentProcess()->HasSwitch(kHintsProtoOverride);
 }
@@ -149,5 +152,10 @@ bool ShouldOverrideCheckingUserPermissionsToFetchHintsForTesting() {
   return command_line->HasSwitch(kDisableCheckingUserPermissionsForTesting);
 }
 
+bool ShouldSkipModelDownloadVerificationForTesting() {
+  base::CommandLine* command_line = base::CommandLine::ForCurrentProcess();
+  return command_line->HasSwitch(kDisableModelDownloadVerificationForTesting);
+}
+
 }  // namespace switches
 }  // namespace optimization_guide

components/optimization_guide/optimization_guide_switches.h

@@ -29,6 +29,7 @@ extern const char kPurgeHintsStore[];
 extern const char kPurgeModelAndFeaturesStore[];
 extern const char kDisableFetchingHintsAtNavigationStartForTesting[];
 extern const char kDisableCheckingUserPermissionsForTesting[];
+extern const char kDisableModelDownloadVerificationForTesting[];
 
 // Returns whether the hint component should be processed.
 // Available hint components are only processed if a proto override isn't being
@@ -72,6 +73,9 @@ bool DisableFetchingHintsAtNavigationStartForTesting();
 // tests.
 bool ShouldOverrideCheckingUserPermissionsToFetchHintsForTesting();
 
+// Returns true if the verification of model downloads should be skipped.
+bool ShouldSkipModelDownloadVerificationForTesting();
+
 }  // namespace switches
 }  // namespace optimization_guide