test_model_info_builder.cc 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  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/test_model_info_builder.h"
  5. #include "components/optimization_guide/core/model_util.h"
  6. #include "components/optimization_guide/core/optimization_guide_test_util.h"
  7. namespace optimization_guide {
  8. TestModelInfoBuilder::TestModelInfoBuilder() {
  9. // Valid (dummy values) by default.
  10. model_.mutable_model()->set_download_url(kTestAbsoluteFilePath);
  11. model_.mutable_model_info()->set_version(123);
  12. }
  13. TestModelInfoBuilder::~TestModelInfoBuilder() = default;
  14. TestModelInfoBuilder& TestModelInfoBuilder::SetModelFilePath(
  15. const base::FilePath& file_path) {
  16. model_.mutable_model()->set_download_url(FilePathToString(file_path));
  17. return *this;
  18. }
  19. TestModelInfoBuilder& TestModelInfoBuilder::SetAdditionalFiles(
  20. const base::flat_set<base::FilePath>& additional_files) {
  21. for (const base::FilePath& file_path : additional_files) {
  22. model_.mutable_model_info()->add_additional_files()->set_file_path(
  23. FilePathToString(file_path));
  24. }
  25. return *this;
  26. }
  27. TestModelInfoBuilder& TestModelInfoBuilder::SetVersion(int64_t version) {
  28. model_.mutable_model_info()->set_version(version);
  29. return *this;
  30. }
  31. TestModelInfoBuilder& TestModelInfoBuilder::SetModelMetadata(
  32. absl::optional<proto::Any> model_metadata) {
  33. if (!model_metadata) {
  34. model_.mutable_model_info()->clear_model_metadata();
  35. return *this;
  36. }
  37. *model_.mutable_model_info()->mutable_model_metadata() =
  38. model_metadata.value();
  39. return *this;
  40. }
  41. std::unique_ptr<ModelInfo> TestModelInfoBuilder::Build() {
  42. return ModelInfo::Create(model_);
  43. }
  44. } // namespace optimization_guide