prediction_model_download_manager_unittest.cc 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655
  1. // Copyright 2020 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "components/optimization_guide/core/prediction_model_download_manager.h"
  5. #include "base/command_line.h"
  6. #include "base/files/file_util.h"
  7. #include "base/files/scoped_temp_dir.h"
  8. #include "base/path_service.h"
  9. #include "base/sequence_checker.h"
  10. #include "base/strings/utf_string_conversions.h"
  11. #include "base/test/gmock_move_support.h"
  12. #include "base/test/metrics/histogram_tester.h"
  13. #include "base/test/scoped_feature_list.h"
  14. #include "base/test/task_environment.h"
  15. #include "build/build_config.h"
  16. #include "components/download/public/background_service/test/mock_download_service.h"
  17. #include "components/optimization_guide/core/model_util.h"
  18. #include "components/optimization_guide/core/optimization_guide_enums.h"
  19. #include "components/optimization_guide/core/optimization_guide_features.h"
  20. #include "components/optimization_guide/core/optimization_guide_switches.h"
  21. #include "components/optimization_guide/core/optimization_guide_util.h"
  22. #include "components/optimization_guide/core/prediction_model_download_observer.h"
  23. #include "components/services/unzip/in_process_unzipper.h"
  24. #include "testing/gtest/include/gtest/gtest.h"
  25. #include "third_party/abseil-cpp/absl/types/optional.h"
  26. #include "third_party/zlib/google/zip.h"
  27. #if !BUILDFLAG(IS_IOS)
  28. #include "components/services/unzip/content/unzip_service.h" // nogncheck
  29. #endif
  30. namespace optimization_guide {
  31. using ::testing::_;
  32. using ::testing::Eq;
  33. class TestPredictionModelDownloadObserver
  34. : public PredictionModelDownloadObserver {
  35. public:
  36. TestPredictionModelDownloadObserver() = default;
  37. ~TestPredictionModelDownloadObserver() override = default;
  38. void OnModelReady(const proto::PredictionModel& model) override {
  39. last_ready_model_ = model;
  40. }
  41. void OnModelDownloadStarted(
  42. proto::OptimizationTarget optimization_target) override {}
  43. void OnModelDownloadFailed(
  44. proto::OptimizationTarget optimization_target) override {}
  45. absl::optional<proto::PredictionModel> last_ready_model() const {
  46. return last_ready_model_;
  47. }
  48. private:
  49. absl::optional<proto::PredictionModel> last_ready_model_;
  50. };
  51. enum class PredictionModelDownloadFileStatus {
  52. kVerifiedCrxWithGoodModelFiles,
  53. kVerifiedCrxWithNoFiles,
  54. kVerifiedCrxWithInvalidPublisher,
  55. kVerifiedCrxWithBadModelInfoFile,
  56. kVerifiedCrxWithInvalidModelInfo,
  57. kVerfiedCrxWithValidModelInfoNoModelFile,
  58. kUnverifiedFile,
  59. };
  60. class PredictionModelDownloadManagerTest : public testing::Test {
  61. public:
  62. PredictionModelDownloadManagerTest() = default;
  63. ~PredictionModelDownloadManagerTest() override = default;
  64. void SetUp() override {
  65. ASSERT_TRUE(temp_download_dir_.CreateUniqueTempDir());
  66. ASSERT_TRUE(temp_models_dir_.CreateUniqueTempDir());
  67. mock_download_service_ =
  68. std::make_unique<download::test::MockDownloadService>();
  69. download_manager_ = std::make_unique<PredictionModelDownloadManager>(
  70. mock_download_service_.get(), temp_models_dir_.GetPath(),
  71. task_environment_.GetMainThreadTaskRunner());
  72. #if !BUILDFLAG(IS_IOS)
  73. unzip::SetUnzipperLaunchOverrideForTesting(
  74. base::BindRepeating(&unzip::LaunchInProcessUnzipper));
  75. #endif
  76. }
  77. void TearDown() override {
  78. download_manager_.reset();
  79. mock_download_service_ = nullptr;
  80. }
  81. PredictionModelDownloadManager* download_manager() {
  82. return download_manager_.get();
  83. }
  84. download::test::MockDownloadService* download_service() {
  85. return mock_download_service_.get();
  86. }
  87. base::FilePath models_dir() const { return temp_models_dir_.GetPath(); }
  88. protected:
  89. void SetDownloadServiceReady(
  90. const std::set<std::string>& pending_guids,
  91. const std::map<std::string, PredictionModelDownloadFileStatus>&
  92. successful_guids) {
  93. std::map<std::string, base::FilePath> success_map;
  94. for (const auto& guid_and_status : successful_guids) {
  95. success_map.emplace(
  96. guid_and_status.first,
  97. GetFilePathForDownloadFileStatus(guid_and_status.second));
  98. }
  99. download_manager()->OnDownloadServiceReady(pending_guids, success_map);
  100. }
  101. void SetDownloadServiceUnavailable() {
  102. download_manager()->OnDownloadServiceUnavailable();
  103. }
  104. void SetDownloadSucceeded(const std::string& guid,
  105. PredictionModelDownloadFileStatus file_status) {
  106. WriteFileForStatus(file_status);
  107. download_manager()->OnDownloadSucceeded(
  108. proto::OptimizationTarget::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD, guid,
  109. GetFilePathForDownloadFileStatus(file_status));
  110. }
  111. void SetDownloadFailed(const std::string& guid) {
  112. download_manager()->OnDownloadFailed(
  113. proto::OptimizationTarget::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD, guid);
  114. }
  115. void RunUntilIdle() { task_environment_.RunUntilIdle(); }
  116. base::FilePath GetFilePathForDownloadFileStatus(
  117. PredictionModelDownloadFileStatus file_status) {
  118. base::FilePath path;
  119. base::PathService::Get(base::DIR_SOURCE_ROOT, &path);
  120. switch (file_status) {
  121. case PredictionModelDownloadFileStatus::kUnverifiedFile:
  122. return temp_download_dir_.GetPath().AppendASCII("unverified.crx3");
  123. case PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidPublisher:
  124. return temp_download_dir_.GetPath().AppendASCII(
  125. "invalidpublisher.crx3");
  126. case PredictionModelDownloadFileStatus::kVerifiedCrxWithNoFiles:
  127. return temp_download_dir_.GetPath().AppendASCII("nofiles.crx3");
  128. case PredictionModelDownloadFileStatus::kVerifiedCrxWithBadModelInfoFile:
  129. return temp_download_dir_.GetPath().AppendASCII("badmodelinfo.crx3");
  130. case PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidModelInfo:
  131. return temp_download_dir_.GetPath().AppendASCII(
  132. "invalidmodelinfo.crx3");
  133. case PredictionModelDownloadFileStatus::
  134. kVerfiedCrxWithValidModelInfoNoModelFile:
  135. return temp_download_dir_.GetPath().AppendASCII("nomodel.crx3");
  136. case PredictionModelDownloadFileStatus::kVerifiedCrxWithGoodModelFiles:
  137. return temp_download_dir_.GetPath().AppendASCII("good.crx3");
  138. }
  139. }
  140. void TurnOffDownloadVerification() {
  141. base::CommandLine::ForCurrentProcess()->AppendSwitch(
  142. switches::kDisableModelDownloadVerificationForTesting);
  143. }
  144. // Retries until the path has been deleted or until all handles to |path| have
  145. // been closed. Returns whether |path| has been deleted.
  146. //
  147. // See crbug/1156112#c1 for suggested mitigation steps.
  148. bool HasPathBeenDeleted(const base::FilePath& path) {
  149. while (true) {
  150. RunUntilIdle();
  151. bool path_exists = base::PathExists(path);
  152. if (!path_exists)
  153. return true;
  154. // Retry if the last file error is access denied since it's likely that
  155. // the file is in the process of being deleted.
  156. }
  157. }
  158. private:
  159. void WriteFileForStatus(PredictionModelDownloadFileStatus status) {
  160. base::FilePath source_root_dir;
  161. base::PathService::Get(base::DIR_SOURCE_ROOT, &source_root_dir);
  162. if (status == PredictionModelDownloadFileStatus::
  163. kVerifiedCrxWithInvalidPublisher ||
  164. status == PredictionModelDownloadFileStatus::kUnverifiedFile) {
  165. base::FilePath crx_file_source_dir =
  166. source_root_dir.AppendASCII("components")
  167. .AppendASCII("test")
  168. .AppendASCII("data")
  169. .AppendASCII("crx_file");
  170. std::string crx_file =
  171. status == PredictionModelDownloadFileStatus::kUnverifiedFile
  172. ? "unsigned.crx3"
  173. : "valid_publisher.crx3"; // Despite name, wrong publisher.
  174. ASSERT_TRUE(base::CopyFile(crx_file_source_dir.AppendASCII(crx_file),
  175. GetFilePathForDownloadFileStatus(status)));
  176. return;
  177. }
  178. if (status == PredictionModelDownloadFileStatus::kVerifiedCrxWithNoFiles) {
  179. base::FilePath invalid_crx_model =
  180. source_root_dir.AppendASCII("components")
  181. .AppendASCII("test")
  182. .AppendASCII("data")
  183. .AppendASCII("optimization_guide")
  184. .AppendASCII("invalid_model.crx3");
  185. ASSERT_TRUE(base::CopyFile(invalid_crx_model,
  186. GetFilePathForDownloadFileStatus(status)));
  187. return;
  188. }
  189. base::FilePath zip_dir =
  190. temp_download_dir_.GetPath().AppendASCII("zip_dir");
  191. ASSERT_TRUE(base::CreateDirectory(zip_dir));
  192. if (status ==
  193. PredictionModelDownloadFileStatus::kVerifiedCrxWithBadModelInfoFile) {
  194. base::WriteFile(zip_dir.AppendASCII("model-info.pb"), "boo", 3);
  195. } else {
  196. proto::ModelInfo model_info;
  197. model_info.set_optimization_target(
  198. proto::OptimizationTarget::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
  199. model_info.set_version(123);
  200. if (status ==
  201. PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidModelInfo) {
  202. model_info.clear_version();
  203. }
  204. std::string serialized_model_info;
  205. ASSERT_TRUE(model_info.SerializeToString(&serialized_model_info));
  206. ASSERT_EQ(static_cast<int32_t>(serialized_model_info.length()),
  207. base::WriteFile(zip_dir.AppendASCII("model-info.pb"),
  208. serialized_model_info.data(),
  209. serialized_model_info.length()));
  210. if (status ==
  211. PredictionModelDownloadFileStatus::kVerifiedCrxWithGoodModelFiles) {
  212. base::WriteFile(zip_dir.AppendASCII("model.tflite"), "model", 5);
  213. }
  214. }
  215. ASSERT_TRUE(
  216. zip::Zip(zip_dir, GetFilePathForDownloadFileStatus(status), true));
  217. }
  218. base::test::TaskEnvironment task_environment_;
  219. base::ScopedTempDir temp_download_dir_;
  220. base::ScopedTempDir temp_models_dir_;
  221. std::unique_ptr<download::test::MockDownloadService> mock_download_service_;
  222. std::unique_ptr<PredictionModelDownloadManager> download_manager_;
  223. };
  224. TEST_F(PredictionModelDownloadManagerTest, DownloadServiceReadyPersistsGuids) {
  225. base::HistogramTester histogram_tester;
  226. SetDownloadServiceReady(
  227. {"pending1", "pending2", "pending3"},
  228. {{"success1", PredictionModelDownloadFileStatus::kUnverifiedFile},
  229. {"success2", PredictionModelDownloadFileStatus::kUnverifiedFile},
  230. {"success3", PredictionModelDownloadFileStatus::kUnverifiedFile}});
  231. RunUntilIdle();
  232. // Should only persist and thus cancel the pending ones.
  233. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending1")));
  234. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending2")));
  235. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending3")));
  236. download_manager()->CancelAllPendingDownloads();
  237. // The successful downloads should not trigger us to do anything with them.
  238. histogram_tester.ExpectTotalCount(
  239. "OptimizationGuide.PredictionModelDownloadManager.DownloadSucceeded", 0);
  240. }
  241. TEST_F(PredictionModelDownloadManagerTest, StartDownloadRestrictedDownloading) {
  242. base::HistogramTester histogram_tester;
  243. base::test::ScopedFeatureList features;
  244. features.InitWithFeaturesAndParameters(
  245. {
  246. {optimization_guide::features::kOptimizationGuideModelDownloading,
  247. {{"unrestricted_model_downloading", "false"}}},
  248. },
  249. /*disabled_features=*/{});
  250. download::DownloadParams download_params;
  251. EXPECT_CALL(*download_service(), StartDownload_(_))
  252. .WillOnce(MoveArg<0>(&download_params));
  253. download_manager()->StartDownload(
  254. GURL("someurl"), proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
  255. // Validate parameters - basically that we attach the correct client, just do
  256. // a passthrough of the URL, and attach the API key.
  257. EXPECT_EQ(download_params.client,
  258. download::DownloadClient::OPTIMIZATION_GUIDE_PREDICTION_MODELS);
  259. EXPECT_EQ(download_params.request_params.url, GURL("someurl"));
  260. EXPECT_EQ(download_params.request_params.method, "GET");
  261. EXPECT_TRUE(download_params.request_params.request_headers.HasHeader(
  262. "X-Goog-Api-Key"));
  263. EXPECT_FALSE(download_params.request_params.require_safety_checks);
  264. EXPECT_EQ(download_params.scheduling_params.priority,
  265. download::SchedulingParams::Priority::NORMAL);
  266. EXPECT_EQ(
  267. download_params.scheduling_params.battery_requirements,
  268. download::SchedulingParams::BatteryRequirements::BATTERY_INSENSITIVE);
  269. EXPECT_EQ(download_params.scheduling_params.network_requirements,
  270. download::SchedulingParams::NetworkRequirements::NONE);
  271. histogram_tester.ExpectUniqueSample(
  272. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  273. PredictionModelDownloadManager::PredictionModelDownloadState::kRequested,
  274. 1);
  275. histogram_tester.ExpectTotalCount(
  276. "OptimizationGuide.PredictionModelDownloadManager.DownloadStartLatency."
  277. "PainfulPageLoad",
  278. 0);
  279. // Now invoke start callback.
  280. std::move(download_params.callback)
  281. .Run("someguid", download::DownloadParams::StartResult::ACCEPTED);
  282. histogram_tester.ExpectTotalCount(
  283. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  284. 2);
  285. histogram_tester.ExpectBucketCount(
  286. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  287. PredictionModelDownloadManager::PredictionModelDownloadState::kStarted,
  288. 1);
  289. histogram_tester.ExpectTotalCount(
  290. "OptimizationGuide.PredictionModelDownloadManager.DownloadStartLatency."
  291. "PainfulPageLoad",
  292. 1);
  293. // Now cancel all downloads to ensure that callback persisted pending GUID.
  294. EXPECT_CALL(*download_service(), CancelDownload(Eq("someguid")));
  295. download_manager()->CancelAllPendingDownloads();
  296. }
  297. TEST_F(PredictionModelDownloadManagerTest,
  298. StartDownloadUnrestrictedDownloading) {
  299. base::HistogramTester histogram_tester;
  300. base::test::ScopedFeatureList features;
  301. features.InitWithFeaturesAndParameters(
  302. {
  303. {optimization_guide::features::kOptimizationGuideModelDownloading,
  304. {{"unrestricted_model_downloading", "true"}}},
  305. },
  306. /*disabled_features=*/{});
  307. download::DownloadParams download_params;
  308. EXPECT_CALL(*download_service(), StartDownload_(_))
  309. .WillOnce(MoveArg<0>(&download_params));
  310. download_manager()->StartDownload(
  311. GURL("someurl"), proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
  312. // Validate parameters - basically that we attach the correct client, just do
  313. // a passthrough of the URL, and attach the API key.
  314. EXPECT_EQ(download_params.client,
  315. download::DownloadClient::OPTIMIZATION_GUIDE_PREDICTION_MODELS);
  316. EXPECT_EQ(download_params.request_params.url, GURL("someurl"));
  317. EXPECT_EQ(download_params.request_params.method, "GET");
  318. EXPECT_TRUE(download_params.request_params.request_headers.HasHeader(
  319. "X-Goog-Api-Key"));
  320. EXPECT_FALSE(download_params.request_params.require_safety_checks);
  321. EXPECT_EQ(download_params.scheduling_params.priority,
  322. download::SchedulingParams::Priority::HIGH);
  323. EXPECT_EQ(
  324. download_params.scheduling_params.battery_requirements,
  325. download::SchedulingParams::BatteryRequirements::BATTERY_INSENSITIVE);
  326. EXPECT_EQ(download_params.scheduling_params.network_requirements,
  327. download::SchedulingParams::NetworkRequirements::NONE);
  328. histogram_tester.ExpectUniqueSample(
  329. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  330. PredictionModelDownloadManager::PredictionModelDownloadState::kRequested,
  331. 1);
  332. histogram_tester.ExpectTotalCount(
  333. "OptimizationGuide.PredictionModelDownloadManager.DownloadStartLatency."
  334. "PainfulPageLoad",
  335. 0);
  336. // Now invoke start callback.
  337. std::move(download_params.callback)
  338. .Run("someguid", download::DownloadParams::StartResult::ACCEPTED);
  339. histogram_tester.ExpectTotalCount(
  340. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  341. 2);
  342. histogram_tester.ExpectBucketCount(
  343. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  344. PredictionModelDownloadManager::PredictionModelDownloadState::kStarted,
  345. 1);
  346. histogram_tester.ExpectTotalCount(
  347. "OptimizationGuide.PredictionModelDownloadManager.DownloadStartLatency."
  348. "PainfulPageLoad",
  349. 1);
  350. // Now cancel all downloads to ensure that callback persisted pending GUID.
  351. EXPECT_CALL(*download_service(), CancelDownload(Eq("someguid")));
  352. download_manager()->CancelAllPendingDownloads();
  353. }
  354. TEST_F(PredictionModelDownloadManagerTest, StartDownloadFailedToSchedule) {
  355. base::HistogramTester histogram_tester;
  356. download::DownloadParams download_params;
  357. EXPECT_CALL(*download_service(), StartDownload_(_))
  358. .WillOnce(MoveArg<0>(&download_params));
  359. download_manager()->StartDownload(
  360. GURL("someurl"), proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
  361. // Now invoke start callback.
  362. std::move(download_params.callback)
  363. .Run("someguid", download::DownloadParams::StartResult::INTERNAL_ERROR);
  364. histogram_tester.ExpectUniqueSample(
  365. "OptimizationGuide.PredictionModelDownloadManager.State.PainfulPageLoad",
  366. PredictionModelDownloadManager::PredictionModelDownloadState::kRequested,
  367. 1);
  368. // Now cancel all downloads to ensure that bad GUID was not accepted.
  369. EXPECT_CALL(*download_service(), CancelDownload(_)).Times(0);
  370. download_manager()->CancelAllPendingDownloads();
  371. histogram_tester.ExpectTotalCount(
  372. "OptimizationGuide.PredictionModelDownloadManager.DownloadStartLatency."
  373. "PainfulPageLoad",
  374. 0);
  375. }
  376. TEST_F(PredictionModelDownloadManagerTest, IsAvailableForDownloads) {
  377. EXPECT_TRUE(download_manager()->IsAvailableForDownloads());
  378. SetDownloadServiceUnavailable();
  379. EXPECT_FALSE(download_manager()->IsAvailableForDownloads());
  380. }
  381. TEST_F(PredictionModelDownloadManagerTest,
  382. SuccessfulDownloadShouldNoLongerBeTracked) {
  383. base::HistogramTester histogram_tester;
  384. SetDownloadServiceReady({"pending1", "pending2", "pending3"},
  385. /*successful_guids=*/{});
  386. SetDownloadSucceeded("pending1",
  387. PredictionModelDownloadFileStatus::kUnverifiedFile);
  388. RunUntilIdle();
  389. // Should only persist and thus cancel the pending ones.
  390. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending2")));
  391. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending3")));
  392. download_manager()->CancelAllPendingDownloads();
  393. histogram_tester.ExpectUniqueSample(
  394. "OptimizationGuide.PredictionModelDownloadManager.DownloadSucceeded",
  395. true, 1);
  396. }
  397. TEST_F(PredictionModelDownloadManagerTest,
  398. FailedDownloadShouldNoLongerBeTracked) {
  399. base::HistogramTester histogram_tester;
  400. SetDownloadServiceReady({"pending1", "pending2", "pending3"},
  401. /*successful_guids=*/{});
  402. SetDownloadFailed("pending2");
  403. // Should only persist and thus cancel the pending ones.
  404. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending1")));
  405. EXPECT_CALL(*download_service(), CancelDownload(Eq("pending3")));
  406. download_manager()->CancelAllPendingDownloads();
  407. histogram_tester.ExpectUniqueSample(
  408. "OptimizationGuide.PredictionModelDownloadManager.DownloadSucceeded",
  409. false, 1);
  410. }
  411. TEST_F(PredictionModelDownloadManagerTest, UnverifiedFileShouldDeleteTempFile) {
  412. base::HistogramTester histogram_tester;
  413. TestPredictionModelDownloadObserver observer;
  414. download_manager()->AddObserver(&observer);
  415. SetDownloadSucceeded("model",
  416. PredictionModelDownloadFileStatus::kUnverifiedFile);
  417. RunUntilIdle();
  418. EXPECT_FALSE(observer.last_ready_model().has_value());
  419. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  420. PredictionModelDownloadFileStatus::kUnverifiedFile)));
  421. histogram_tester.ExpectUniqueSample(
  422. "OptimizationGuide.PredictionModelDownloadManager."
  423. "DownloadStatus",
  424. PredictionModelDownloadStatus::kFailedCrxVerification, 1);
  425. }
  426. TEST_F(PredictionModelDownloadManagerTest,
  427. VerifiedCrxWithInvalidPublisherShouldDeleteTempFile) {
  428. base::HistogramTester histogram_tester;
  429. TestPredictionModelDownloadObserver observer;
  430. download_manager()->AddObserver(&observer);
  431. SetDownloadSucceeded(
  432. "model",
  433. PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidPublisher);
  434. RunUntilIdle();
  435. EXPECT_FALSE(observer.last_ready_model().has_value());
  436. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  437. PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidPublisher)));
  438. histogram_tester.ExpectUniqueSample(
  439. "OptimizationGuide.PredictionModelDownloadManager."
  440. "DownloadStatus",
  441. PredictionModelDownloadStatus::kFailedCrxInvalidPublisher, 1);
  442. }
  443. TEST_F(PredictionModelDownloadManagerTest,
  444. VerifiedCrxWithNoFilesShouldDeleteTempFile) {
  445. base::HistogramTester histogram_tester;
  446. TestPredictionModelDownloadObserver observer;
  447. download_manager()->AddObserver(&observer);
  448. SetDownloadSucceeded(
  449. "model", PredictionModelDownloadFileStatus::kVerifiedCrxWithNoFiles);
  450. RunUntilIdle();
  451. EXPECT_FALSE(observer.last_ready_model().has_value());
  452. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  453. PredictionModelDownloadFileStatus::kVerifiedCrxWithNoFiles)));
  454. histogram_tester.ExpectUniqueSample(
  455. "OptimizationGuide.PredictionModelDownloadManager."
  456. "DownloadStatus",
  457. PredictionModelDownloadStatus::kFailedModelInfoFileRead, 1);
  458. }
  459. TEST_F(PredictionModelDownloadManagerTest,
  460. VerifiedCrxWithBadModelInfoFileShouldDeleteTempFile) {
  461. base::HistogramTester histogram_tester;
  462. TestPredictionModelDownloadObserver observer;
  463. download_manager()->AddObserver(&observer);
  464. TurnOffDownloadVerification();
  465. SetDownloadSucceeded(
  466. "model",
  467. PredictionModelDownloadFileStatus::kVerifiedCrxWithBadModelInfoFile);
  468. RunUntilIdle();
  469. EXPECT_FALSE(observer.last_ready_model().has_value());
  470. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  471. PredictionModelDownloadFileStatus::kVerifiedCrxWithBadModelInfoFile)));
  472. histogram_tester.ExpectUniqueSample(
  473. "OptimizationGuide.PredictionModelDownloadManager."
  474. "DownloadStatus",
  475. PredictionModelDownloadStatus::kFailedModelInfoParsing, 1);
  476. }
  477. TEST_F(PredictionModelDownloadManagerTest,
  478. VerifiedCrxWithInvalidModelInfoShouldDeleteTempFile) {
  479. base::HistogramTester histogram_tester;
  480. TestPredictionModelDownloadObserver observer;
  481. download_manager()->AddObserver(&observer);
  482. TurnOffDownloadVerification();
  483. SetDownloadSucceeded(
  484. "model",
  485. PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidModelInfo);
  486. RunUntilIdle();
  487. EXPECT_FALSE(observer.last_ready_model().has_value());
  488. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  489. PredictionModelDownloadFileStatus::kVerifiedCrxWithInvalidModelInfo)));
  490. histogram_tester.ExpectUniqueSample(
  491. "OptimizationGuide.PredictionModelDownloadManager."
  492. "DownloadStatus",
  493. PredictionModelDownloadStatus::kFailedModelInfoInvalid, 1);
  494. }
  495. TEST_F(PredictionModelDownloadManagerTest,
  496. VerifiedCrxWithValidModelInfoFileButNoModelFileShouldDeleteTempFile) {
  497. base::HistogramTester histogram_tester;
  498. TestPredictionModelDownloadObserver observer;
  499. download_manager()->AddObserver(&observer);
  500. TurnOffDownloadVerification();
  501. SetDownloadSucceeded("model", PredictionModelDownloadFileStatus::
  502. kVerfiedCrxWithValidModelInfoNoModelFile);
  503. RunUntilIdle();
  504. EXPECT_FALSE(observer.last_ready_model().has_value());
  505. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  506. PredictionModelDownloadFileStatus::
  507. kVerfiedCrxWithValidModelInfoNoModelFile)));
  508. histogram_tester.ExpectUniqueSample(
  509. "OptimizationGuide.PredictionModelDownloadManager."
  510. "DownloadStatus",
  511. PredictionModelDownloadStatus::kFailedModelFileOtherError, 1);
  512. // The error code for ReplaceFile varies by platform for this test, only
  513. // care that the error code is recorded.
  514. histogram_tester.ExpectTotalCount(
  515. "OptimizationGuide.PredictionModelDownloadManager.ReplaceFileError."
  516. "PainfulPageLoad",
  517. 1);
  518. }
  519. TEST_F(
  520. PredictionModelDownloadManagerTest,
  521. VerifiedCrxWithGoodModelFilesShouldDeleteDownloadFileButHaveContentExtracted) {
  522. base::HistogramTester histogram_tester;
  523. TestPredictionModelDownloadObserver observer;
  524. download_manager()->AddObserver(&observer);
  525. TurnOffDownloadVerification();
  526. SetDownloadSucceeded(
  527. "modelfile",
  528. PredictionModelDownloadFileStatus::kVerifiedCrxWithGoodModelFiles);
  529. RunUntilIdle();
  530. ASSERT_TRUE(observer.last_ready_model().has_value());
  531. EXPECT_EQ(observer.last_ready_model()->model_info().optimization_target(),
  532. proto::OptimizationTarget::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
  533. EXPECT_EQ(observer.last_ready_model()->model_info().version(), 123);
  534. // Make sure that there is a file path is written to the model file.
  535. EXPECT_EQ(StringToFilePath(
  536. observer.last_ready_model().value().model().download_url())
  537. .value()
  538. .DirName()
  539. .DirName(),
  540. models_dir());
  541. EXPECT_EQ(StringToFilePath(
  542. observer.last_ready_model().value().model().download_url())
  543. .value()
  544. .BaseName()
  545. .value(),
  546. FILE_PATH_LITERAL("model.tflite"));
  547. // Downloaded file should still be deleted.
  548. EXPECT_TRUE(HasPathBeenDeleted(GetFilePathForDownloadFileStatus(
  549. PredictionModelDownloadFileStatus::kVerifiedCrxWithGoodModelFiles)));
  550. histogram_tester.ExpectUniqueSample(
  551. "OptimizationGuide.PredictionModelDownloadManager."
  552. "DownloadStatus",
  553. PredictionModelDownloadStatus::kSuccess, 1);
  554. }
  555. } // namespace optimization_guide