ukm_data_manager_test_utils.h 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. // Copyright 2022 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. #ifndef CHROME_BROWSER_SEGMENTATION_PLATFORM_UKM_DATA_MANAGER_TEST_UTILS_H_
  5. #define CHROME_BROWSER_SEGMENTATION_PLATFORM_UKM_DATA_MANAGER_TEST_UTILS_H_
  6. #include <set>
  7. #include <string>
  8. #include "base/callback.h"
  9. #include "base/memory/weak_ptr.h"
  10. #include "components/segmentation_platform/internal/execution/mock_model_provider.h"
  11. #include "components/segmentation_platform/public/proto/model_metadata.pb.h"
  12. #include "components/segmentation_platform/public/proto/segmentation_platform.pb.h"
  13. #include "components/ukm/test_ukm_recorder.h"
  14. class GURL;
  15. namespace history {
  16. class HistoryService;
  17. }
  18. namespace segmentation_platform {
  19. // Utility used for testing UKM based engine.
  20. class UkmDataManagerTestUtils {
  21. public:
  22. explicit UkmDataManagerTestUtils(ukm::TestUkmRecorder* ukm_recorder);
  23. ~UkmDataManagerTestUtils();
  24. UkmDataManagerTestUtils(UkmDataManagerTestUtils&) = delete;
  25. UkmDataManagerTestUtils& operator=(UkmDataManagerTestUtils&) = delete;
  26. // Must be called before the first profile initialization, sets up default
  27. // model overrides for the given `default_overrides`
  28. void PreProfileInit(const std::set<proto::SegmentId>& default_overrides);
  29. // Waits for platform to initialize and request default model for
  30. // `segment_id`, and then returns the provided `metadata` to the platform.
  31. void WaitForModelRequestAndUpdateWith(
  32. proto::SegmentId segment_id,
  33. const proto::SegmentationModelMetadata& metadata);
  34. // Creates a sample page load UKM based model metadata, with a simple SQL
  35. // feature with `query`.
  36. proto::SegmentationModelMetadata GetSamplePageLoadMetadata(
  37. const std::string& query);
  38. // Records a page load and 2 valid UKM metrics associated with it. May record
  39. // other UKM metrics that are unrelated to the metadata provided by
  40. // GetSamplePageLoadMetadata().
  41. void RecordPageLoadUkm(const GURL& url, base::Time history_timestamp);
  42. // Returns whether the `url` is part of the UKM database.
  43. bool IsUrlInDatabase(const GURL& url);
  44. // Returns the model provider override for the `segment_id`.
  45. MockModelProvider* GetDefaultOverride(proto::SegmentId segment_id);
  46. // History service is needed for validating test URLs written to database.
  47. void set_history_service(history::HistoryService* history_service) {
  48. history_service_ = history_service;
  49. }
  50. private:
  51. void StoreModelUpdateCallback(
  52. proto::SegmentId segment_id,
  53. const ModelProvider::ModelUpdatedCallback& callback);
  54. const raw_ptr<ukm::TestUkmRecorder> ukm_recorder_;
  55. int source_id_counter_ = 1;
  56. raw_ptr<history::HistoryService> history_service_;
  57. std::map<proto::SegmentId, MockModelProvider*> default_overrides_;
  58. std::map<proto::SegmentId, std::vector<ModelProvider::ModelUpdatedCallback>>
  59. callbacks_;
  60. base::WeakPtrFactory<UkmDataManagerTestUtils> weak_factory_{this};
  61. };
  62. } // namespace segmentation_platform
  63. #endif // CHROME_BROWSER_SEGMENTATION_PLATFORM_UKM_DATA_MANAGER_TEST_UTILS_H_