ranker_model_loader_impl_unittest.cc 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345
  1. // Copyright 2017 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/assist_ranker/ranker_model_loader_impl.h"
  5. #include <initializer_list>
  6. #include <memory>
  7. #include <vector>
  8. #include "base/bind.h"
  9. #include "base/containers/circular_deque.h"
  10. #include "base/files/file_util.h"
  11. #include "base/files/scoped_temp_dir.h"
  12. #include "base/memory/ref_counted.h"
  13. #include "base/test/task_environment.h"
  14. #include "base/threading/thread_task_runner_handle.h"
  15. #include "components/assist_ranker/proto/ranker_model.pb.h"
  16. #include "components/assist_ranker/proto/translate_ranker_model.pb.h"
  17. #include "components/assist_ranker/ranker_model.h"
  18. #include "services/network/public/cpp/shared_url_loader_factory.h"
  19. #include "services/network/public/cpp/weak_wrapper_shared_url_loader_factory.h"
  20. #include "services/network/public/mojom/url_response_head.mojom.h"
  21. #include "services/network/test/test_url_loader_factory.h"
  22. #include "testing/gtest/include/gtest/gtest.h"
  23. namespace {
  24. using assist_ranker::RankerModel;
  25. using assist_ranker::RankerModelLoaderImpl;
  26. using assist_ranker::RankerModelStatus;
  27. const char kInvalidModelData[] = "not a valid model";
  28. const int kInvalidModelSize = sizeof(kInvalidModelData) - 1;
  29. class RankerModelLoaderImplTest : public ::testing::Test {
  30. public:
  31. RankerModelLoaderImplTest(const RankerModelLoaderImplTest&) = delete;
  32. RankerModelLoaderImplTest& operator=(const RankerModelLoaderImplTest&) =
  33. delete;
  34. protected:
  35. RankerModelLoaderImplTest();
  36. void SetUp() override;
  37. // Returns a copy of |model|.
  38. static std::unique_ptr<RankerModel> Clone(const RankerModel& model);
  39. // Returns true if |m1| and |m2| are identical.
  40. static bool IsEqual(const RankerModel& m1, const RankerModel& m2);
  41. // Returns true if |m1| and |m2| are identical modulo metadata.
  42. static bool IsEquivalent(const RankerModel& m1, const RankerModel& m2);
  43. // Helper method to drive the loader for |model_path| and |model_url|.
  44. bool DoLoaderTest(const base::FilePath& model_path, const GURL& model_url);
  45. // Initialize the "remote" model data used for testing.
  46. void InitRemoteModels();
  47. // Initialize the "local" model data used for testing.
  48. void InitLocalModels();
  49. // Helper method used by InitRemoteModels() and InitLocalModels().
  50. void InitModel(const GURL& model_url,
  51. const base::Time& last_modified,
  52. const base::TimeDelta& cache_duration,
  53. RankerModel* model);
  54. // Save |model| to |model_path|. Used by InitRemoteModels() and
  55. // InitLocalModels()
  56. void SaveModel(const RankerModel& model, const base::FilePath& model_path);
  57. // Implements RankerModelLoaderImpl's ValidateModelCallback interface.
  58. RankerModelStatus ValidateModel(const RankerModel& model);
  59. // Implements RankerModelLoaderImpl's OnModelAvailableCallback interface.
  60. void OnModelAvailable(std::unique_ptr<RankerModel> model);
  61. // Sets up the task scheduling/task-runner environment for each test.
  62. base::test::TaskEnvironment task_environment_;
  63. // Override the default URL loader to return custom responses for tests.
  64. network::TestURLLoaderFactory test_loader_factory_;
  65. scoped_refptr<network::SharedURLLoaderFactory> test_shared_loader_factory_;
  66. // Temporary directory for model files.
  67. base::ScopedTempDir scoped_temp_dir_;
  68. // A queue of responses to return from Validate(). If empty, validate will
  69. // return 'OK'.
  70. base::circular_deque<RankerModelStatus> validate_model_response_;
  71. // A cached to remember the model validation calls.
  72. std::vector<std::unique_ptr<RankerModel>> validated_models_;
  73. // A cache to remember the OnModelAvailable calls.
  74. std::vector<std::unique_ptr<RankerModel>> available_models_;
  75. // Cached model file paths.
  76. base::FilePath local_model_path_;
  77. base::FilePath expired_model_path_;
  78. base::FilePath invalid_model_path_;
  79. // Model URLS.
  80. GURL remote_model_url_;
  81. GURL invalid_model_url_;
  82. GURL failed_model_url_;
  83. // Model Data.
  84. RankerModel remote_model_;
  85. RankerModel local_model_;
  86. RankerModel expired_model_;
  87. };
  88. RankerModelLoaderImplTest::RankerModelLoaderImplTest() {
  89. test_shared_loader_factory_ =
  90. base::MakeRefCounted<network::WeakWrapperSharedURLLoaderFactory>(
  91. &test_loader_factory_);
  92. }
  93. void RankerModelLoaderImplTest::SetUp() {
  94. ASSERT_TRUE(scoped_temp_dir_.CreateUniqueTempDir());
  95. const auto& temp_dir_path = scoped_temp_dir_.GetPath();
  96. // Setup the model file paths.
  97. local_model_path_ = temp_dir_path.AppendASCII("local_model.bin");
  98. expired_model_path_ = temp_dir_path.AppendASCII("expired_model.bin");
  99. invalid_model_path_ = temp_dir_path.AppendASCII("invalid_model.bin");
  100. // Setup the model URLs.
  101. remote_model_url_ = GURL("https://some.url.net/good.model.bin");
  102. invalid_model_url_ = GURL("https://some.url.net/bad.model.bin");
  103. failed_model_url_ = GURL("https://some.url.net/fail");
  104. // Initialize the model data.
  105. ASSERT_NO_FATAL_FAILURE(InitRemoteModels());
  106. ASSERT_NO_FATAL_FAILURE(InitLocalModels());
  107. }
  108. // static
  109. std::unique_ptr<RankerModel> RankerModelLoaderImplTest::Clone(
  110. const RankerModel& model) {
  111. auto copy = std::make_unique<RankerModel>();
  112. *copy->mutable_proto() = model.proto();
  113. return copy;
  114. }
  115. // static
  116. bool RankerModelLoaderImplTest::IsEqual(const RankerModel& m1,
  117. const RankerModel& m2) {
  118. return m1.SerializeAsString() == m2.SerializeAsString();
  119. }
  120. // static
  121. bool RankerModelLoaderImplTest::IsEquivalent(const RankerModel& m1,
  122. const RankerModel& m2) {
  123. auto copy_m1 = Clone(m1);
  124. copy_m1->mutable_proto()->mutable_metadata()->Clear();
  125. auto copy_m2 = Clone(m2);
  126. copy_m2->mutable_proto()->mutable_metadata()->Clear();
  127. return IsEqual(*copy_m1, *copy_m2);
  128. }
  129. bool RankerModelLoaderImplTest::DoLoaderTest(const base::FilePath& model_path,
  130. const GURL& model_url) {
  131. auto loader = std::make_unique<RankerModelLoaderImpl>(
  132. base::BindRepeating(&RankerModelLoaderImplTest::ValidateModel,
  133. base::Unretained(this)),
  134. base::BindRepeating(&RankerModelLoaderImplTest::OnModelAvailable,
  135. base::Unretained(this)),
  136. test_shared_loader_factory_, model_path, model_url,
  137. "RankerModelLoaderImplTest");
  138. loader->NotifyOfRankerActivity();
  139. task_environment_.RunUntilIdle();
  140. return true;
  141. }
  142. void RankerModelLoaderImplTest::InitRemoteModels() {
  143. InitModel(remote_model_url_, base::Time(), base::TimeDelta(), &remote_model_);
  144. test_loader_factory_.AddResponse(remote_model_url_.spec(),
  145. remote_model_.SerializeAsString());
  146. test_loader_factory_.AddResponse(invalid_model_url_.spec(),
  147. kInvalidModelData);
  148. test_loader_factory_.AddResponse(
  149. failed_model_url_, network::mojom::URLResponseHead::New(), "",
  150. network::URLLoaderCompletionStatus(net::HTTP_INTERNAL_SERVER_ERROR));
  151. }
  152. void RankerModelLoaderImplTest::InitLocalModels() {
  153. InitModel(remote_model_url_, base::Time::Now(), base::Days(30),
  154. &local_model_);
  155. InitModel(remote_model_url_, base::Time::Now() - base::Days(60),
  156. base::Days(30), &expired_model_);
  157. SaveModel(local_model_, local_model_path_);
  158. SaveModel(expired_model_, expired_model_path_);
  159. ASSERT_EQ(base::WriteFile(invalid_model_path_, kInvalidModelData,
  160. kInvalidModelSize),
  161. kInvalidModelSize);
  162. }
  163. void RankerModelLoaderImplTest::InitModel(const GURL& model_url,
  164. const base::Time& last_modified,
  165. const base::TimeDelta& cache_duration,
  166. RankerModel* model) {
  167. ASSERT_TRUE(model != nullptr);
  168. model->mutable_proto()->Clear();
  169. auto* metadata = model->mutable_proto()->mutable_metadata();
  170. if (!model_url.is_empty())
  171. metadata->set_source(model_url.spec());
  172. if (!last_modified.is_null()) {
  173. auto last_modified_sec = (last_modified - base::Time()).InSeconds();
  174. metadata->set_last_modified_sec(last_modified_sec);
  175. }
  176. if (!cache_duration.is_zero())
  177. metadata->set_cache_duration_sec(cache_duration.InSeconds());
  178. auto* translate = model->mutable_proto()->mutable_translate();
  179. translate->set_version(1);
  180. auto* logit = translate->mutable_translate_logistic_regression_model();
  181. logit->set_bias(0.1f);
  182. logit->set_accept_ratio_weight(0.2f);
  183. logit->set_decline_ratio_weight(0.3f);
  184. logit->set_ignore_ratio_weight(0.4f);
  185. }
  186. void RankerModelLoaderImplTest::SaveModel(const RankerModel& model,
  187. const base::FilePath& model_path) {
  188. std::string model_str = model.SerializeAsString();
  189. ASSERT_EQ(base::WriteFile(model_path, model_str.data(), model_str.size()),
  190. static_cast<int>(model_str.size()));
  191. }
  192. RankerModelStatus RankerModelLoaderImplTest::ValidateModel(
  193. const RankerModel& model) {
  194. validated_models_.push_back(Clone(model));
  195. RankerModelStatus response = RankerModelStatus::OK;
  196. if (!validate_model_response_.empty()) {
  197. response = validate_model_response_.front();
  198. validate_model_response_.pop_front();
  199. }
  200. return response;
  201. }
  202. void RankerModelLoaderImplTest::OnModelAvailable(
  203. std::unique_ptr<RankerModel> model) {
  204. available_models_.push_back(std::move(model));
  205. }
  206. } // namespace
  207. TEST_F(RankerModelLoaderImplTest, NoLocalOrRemoteModel) {
  208. ASSERT_TRUE(DoLoaderTest(base::FilePath(), GURL()));
  209. EXPECT_EQ(0U, validated_models_.size());
  210. EXPECT_EQ(0U, available_models_.size());
  211. }
  212. TEST_F(RankerModelLoaderImplTest, BadLocalAndRemoteModel) {
  213. ASSERT_TRUE(DoLoaderTest(invalid_model_path_, invalid_model_url_));
  214. EXPECT_EQ(0U, validated_models_.size());
  215. EXPECT_EQ(0U, available_models_.size());
  216. }
  217. TEST_F(RankerModelLoaderImplTest, LoadFromFileOnly) {
  218. EXPECT_TRUE(DoLoaderTest(local_model_path_, GURL()));
  219. ASSERT_EQ(1U, validated_models_.size());
  220. ASSERT_EQ(1U, available_models_.size());
  221. EXPECT_TRUE(IsEqual(*validated_models_[0], local_model_));
  222. EXPECT_TRUE(IsEqual(*available_models_[0], local_model_));
  223. }
  224. TEST_F(RankerModelLoaderImplTest, LoadFromFileSkipsDownload) {
  225. ASSERT_TRUE(DoLoaderTest(local_model_path_, remote_model_url_));
  226. ASSERT_EQ(1U, validated_models_.size());
  227. ASSERT_EQ(1U, available_models_.size());
  228. EXPECT_TRUE(IsEqual(*validated_models_[0], local_model_));
  229. EXPECT_TRUE(IsEqual(*available_models_[0], local_model_));
  230. }
  231. TEST_F(RankerModelLoaderImplTest, LoadFromFileAndBadUrl) {
  232. ASSERT_TRUE(DoLoaderTest(local_model_path_, invalid_model_url_));
  233. ASSERT_EQ(1U, validated_models_.size());
  234. ASSERT_EQ(1U, available_models_.size());
  235. EXPECT_TRUE(IsEqual(*validated_models_[0], local_model_));
  236. EXPECT_TRUE(IsEqual(*available_models_[0], local_model_));
  237. }
  238. TEST_F(RankerModelLoaderImplTest, LoadFromURLOnly) {
  239. ASSERT_TRUE(DoLoaderTest(base::FilePath(), remote_model_url_));
  240. ASSERT_EQ(1U, validated_models_.size());
  241. ASSERT_EQ(1U, available_models_.size());
  242. EXPECT_TRUE(IsEquivalent(*validated_models_[0], remote_model_));
  243. EXPECT_TRUE(IsEquivalent(*available_models_[0], remote_model_));
  244. }
  245. TEST_F(RankerModelLoaderImplTest, LoadFromExpiredFileTriggersDownload) {
  246. ASSERT_TRUE(DoLoaderTest(expired_model_path_, remote_model_url_));
  247. ASSERT_EQ(2U, validated_models_.size());
  248. ASSERT_EQ(2U, available_models_.size());
  249. EXPECT_TRUE(IsEquivalent(*validated_models_[0], local_model_));
  250. EXPECT_TRUE(IsEquivalent(*available_models_[0], local_model_));
  251. EXPECT_TRUE(IsEquivalent(*validated_models_[1], remote_model_));
  252. EXPECT_TRUE(IsEquivalent(*available_models_[1], remote_model_));
  253. }
  254. TEST_F(RankerModelLoaderImplTest, LoadFromBadFileTriggersDownload) {
  255. ASSERT_TRUE(DoLoaderTest(invalid_model_path_, remote_model_url_));
  256. ASSERT_EQ(1U, validated_models_.size());
  257. ASSERT_EQ(1U, available_models_.size());
  258. EXPECT_TRUE(IsEquivalent(*validated_models_[0], remote_model_));
  259. EXPECT_TRUE(IsEquivalent(*available_models_[0], remote_model_));
  260. }
  261. TEST_F(RankerModelLoaderImplTest, IncompatibleCachedFileTriggersDownload) {
  262. validate_model_response_.push_back(RankerModelStatus::INCOMPATIBLE);
  263. ASSERT_TRUE(DoLoaderTest(local_model_path_, remote_model_url_));
  264. ASSERT_EQ(2U, validated_models_.size());
  265. ASSERT_EQ(1U, available_models_.size());
  266. EXPECT_TRUE(IsEquivalent(*validated_models_[0], local_model_));
  267. EXPECT_TRUE(IsEquivalent(*validated_models_[1], remote_model_));
  268. EXPECT_TRUE(IsEquivalent(*available_models_[0], remote_model_));
  269. }
  270. TEST_F(RankerModelLoaderImplTest, IncompatibleDownloadedFileKeepsExpired) {
  271. validate_model_response_.push_back(RankerModelStatus::OK);
  272. validate_model_response_.push_back(RankerModelStatus::INCOMPATIBLE);
  273. ASSERT_TRUE(DoLoaderTest(expired_model_path_, remote_model_url_));
  274. ASSERT_EQ(2U, validated_models_.size());
  275. ASSERT_EQ(1U, available_models_.size());
  276. EXPECT_TRUE(IsEquivalent(*validated_models_[0], local_model_));
  277. EXPECT_TRUE(IsEquivalent(*validated_models_[1], remote_model_));
  278. EXPECT_TRUE(IsEquivalent(*available_models_[0], local_model_));
  279. }