1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
|
// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "chrome/browser/segmentation_platform/ukm_data_manager_test_utils.h"
#include "base/run_loop.h"
#include "chrome/browser/history/history_service_factory.h"
#include "chrome/browser/segmentation_platform/segmentation_platform_service_factory.h"
#include "chrome/browser/segmentation_platform/ukm_database_client.h"
#include "components/history/core/browser/history_service.h"
#include "components/segmentation_platform/embedder/model_provider_factory_impl.h"
#include "components/segmentation_platform/internal/database/ukm_database.h"
#include "components/segmentation_platform/internal/execution/mock_model_provider.h"
#include "components/segmentation_platform/internal/metadata/metadata_writer.h"
#include "components/segmentation_platform/internal/segmentation_platform_service_impl.h"
#include "components/segmentation_platform/internal/signals/ukm_observer.h"
#include "components/segmentation_platform/internal/ukm_data_manager.h"
#include "components/segmentation_platform/public/proto/segmentation_platform.pb.h"
#include "services/metrics/public/cpp/ukm_builders.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"
namespace segmentation_platform {
namespace {
using ::segmentation_platform::proto::SegmentId;
using ::testing::Return;
using ::ukm::builders::PageLoad;
// Returns a sample UKM entry.
ukm::mojom::UkmEntryPtr GetSamplePageLoadEntry(ukm::SourceId source_id) {
ukm::mojom::UkmEntryPtr entry = ukm::mojom::UkmEntry::New();
entry->source_id = source_id;
entry->event_hash = PageLoad::kEntryNameHash;
entry->metrics[PageLoad::kCpuTimeNameHash] = 10;
entry->metrics[PageLoad::kIsNewBookmarkNameHash] = 20;
entry->metrics[PageLoad::kIsNTPCustomLinkNameHash] = 30;
return entry;
}
// Runs the given query and returns the result as float value. See
// RunReadOnlyQueries() for more info.
std::optional<float> RunQueryAndGetResult(UkmDatabase* database,
UkmDatabase::CustomSqlQuery&& query) {
std::optional<float> output;
UkmDatabase::QueryList queries;
queries.emplace(0, std::move(query));
base::RunLoop wait_for_query;
database->RunReadOnlyQueries(
std::move(queries),
base::BindOnce(
[](base::OnceClosure quit, std::optional<float>* output, bool success,
processing::IndexedTensors tensor) {
if (success) {
EXPECT_EQ(1u, tensor.size());
EXPECT_EQ(1u, tensor.at(0).size());
*output = tensor.at(0)[0].float_val;
}
std::move(quit).Run();
},
wait_for_query.QuitClosure(), &output));
wait_for_query.Run();
return output;
}
} // namespace
UkmDataManagerTestUtils::UkmDataManagerTestUtils(
ukm::TestUkmRecorder* ukm_recorder,
bool owned_db_client)
: ukm_recorder_(ukm_recorder) {
if (owned_db_client) {
owned_db_client_ = std::make_unique<UkmDatabaseClient>();
ukm_database_client_ = owned_db_client_.get();
} else {
ukm_database_client_ = &UkmDatabaseClientHolder::GetClientInstance(nullptr);
}
}
UkmDataManagerTestUtils::~UkmDataManagerTestUtils() {
#if !BUILDFLAG(IS_ANDROID)
// The client should be torn down after profile is destroyed. On Android
// browser tests the profile is never destroyed, so do not tear down the
// client.
ukm_database_client_->TearDownForTesting();
#endif
ukm_database_client_ = nullptr;
}
void UkmDataManagerTestUtils::PreProfileInit(
const std::map<SegmentId, proto::SegmentationModelMetadata>&
default_overrides) {
// Set test recorder before UkmObserver is created.
ukm_database_client_->set_ukm_recorder_for_testing(ukm_recorder_);
for (const auto& segment : default_overrides) {
auto provider = std::make_unique<MockDefaultModelProvider>(segment.first,
segment.second);
default_overrides_[segment.first] = provider.get();
// Default model must be overridden before the platform is created:
TestDefaultModelOverride::GetInstance().SetModelForTesting(
segment.first, std::move(provider));
}
if (owned_db_client_) {
owned_db_client_->PreProfileInit(/*in_memory_database=*/true);
}
}
void UkmDataManagerTestUtils::SetupForProfile(Profile* profile) {
UkmDatabaseClientHolder::SetUkmClientForTesting(profile,
ukm_database_client_.get());
CHECK_EQ(ukm_database_client_.get(),
&UkmDatabaseClientHolder::GetClientInstance(profile));
history_service_ = HistoryServiceFactory::GetForProfile(
profile, ServiceAccessType::EXPLICIT_ACCESS);
// Create the platform to kick off initialization.
segmentation_platform::SegmentationPlatformServiceFactory::GetForProfile(
profile);
}
void UkmDataManagerTestUtils::WillDestroyProfile(Profile* profile) {
UkmDatabaseClientHolder::SetUkmClientForTesting(profile, nullptr);
}
void UkmDataManagerTestUtils::WaitForUkmObserverRegistration() {
UkmObserver* observer = ukm_database_client_->ukm_observer_for_testing();
while (!observer->is_started_for_testing()) {
base::RunLoop().RunUntilIdle();
}
}
proto::SegmentationModelMetadata
UkmDataManagerTestUtils::GetSamplePageLoadMetadata(const std::string& query) {
proto::SegmentationModelMetadata metadata;
MetadataWriter writer(&metadata);
writer.AddOutputConfigForBinaryClassifier(
/*threshold=*/0.5f,
/*positive_label=*/"Show",
/*negative_label=*/"NotShow");
metadata.set_time_unit(proto::TimeUnit::DAY);
metadata.set_bucket_duration(42u);
auto* feature = metadata.add_input_features();
auto* sql_feature = feature->mutable_sql_feature();
sql_feature->set_sql(query);
auto* ukm_event = sql_feature->mutable_signal_filter()->add_ukm_events();
ukm_event->set_event_hash(PageLoad::kEntryNameHash);
ukm_event->add_metric_hash_filter(PageLoad::kCpuTimeNameHash);
ukm_event->add_metric_hash_filter(PageLoad::kIsNewBookmarkNameHash);
return metadata;
}
void UkmDataManagerTestUtils::RecordPageLoadUkm(const GURL& url,
base::Time history_timestamp) {
UkmObserver* observer = ukm_database_client_->ukm_observer_for_testing();
// Ensure that the observer is started before recording metrics.
ASSERT_TRUE(observer->is_started_for_testing());
// Ensure that OTR profiles are not started in the test.
ASSERT_FALSE(observer->is_paused_for_testing());
ukm_recorder_->AddEntry(GetSamplePageLoadEntry(source_id_counter_));
ukm_recorder_->UpdateSourceURL(source_id_counter_, url);
source_id_counter_++;
// Without a history service the recorded URLs will not be written to
// database.
ASSERT_TRUE(history_service_);
history_service_->AddPage(url, history_timestamp,
history::VisitSource::SOURCE_BROWSED);
}
bool UkmDataManagerTestUtils::IsUrlInDatabase(const GURL& url) {
UkmDatabase::CustomSqlQuery query("SELECT 1 FROM urls WHERE url=?",
{processing::ProcessedValue(url.spec())});
std::optional<float> result = RunQueryAndGetResult(
ukm_database_client_->GetUkmDataManager()->GetUkmDatabase(),
std::move(query));
return !!result;
}
MockDefaultModelProvider* UkmDataManagerTestUtils::GetDefaultOverride(
proto::SegmentId segment_id) {
return default_overrides_[segment_id];
}
} // namespace segmentation_platform
|
