1 // Copyright 2019 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 
5 #include "chrome/browser/android/usage_stats/usage_stats_database.h"
6 
7 #include <utility>
8 
9 #include "base/bind.h"
10 #include "base/time/time.h"
11 #include "chrome/browser/android/usage_stats/website_event.pb.h"
12 #include "components/leveldb_proto/testing/fake_db.h"
13 #include "testing/gmock/include/gmock/gmock.h"
14 #include "testing/gtest/include/gtest/gtest.h"
15 
16 using leveldb_proto::test::FakeDB;
17 using testing::ElementsAre;
18 
19 namespace usage_stats {
20 
21 namespace {
22 
23 const char kFqdn1[] = "foo.com";
24 const char kFqdn2[] = "bar.org";
25 
26 const char kToken1[] = "token1";
27 const char kToken2[] = "token2";
28 
CreateWebsiteEvent(const std::string & fqdn,int64_t seconds,const WebsiteEvent::EventType & type)29 const WebsiteEvent CreateWebsiteEvent(const std::string& fqdn,
30                                       int64_t seconds,
31                                       const WebsiteEvent::EventType& type) {
32   WebsiteEvent event;
33   event.set_fqdn(fqdn);
34   Timestamp* timestamp = event.mutable_timestamp();
35   timestamp->set_seconds(seconds);
36   event.set_type(type);
37   return event;
38 }
39 
40 MATCHER_P(EqualsWebsiteEvent, other, "") {
41   return other.fqdn() == arg.fqdn() &&
42          other.timestamp().seconds() == arg.timestamp().seconds() &&
43          other.timestamp().nanos() == arg.timestamp().nanos() &&
44          other.type() == arg.type();
45 }
46 
47 }  // namespace
48 
49 class MockUsageStatsDatabase : public UsageStatsDatabase {
50  public:
MockUsageStatsDatabase(std::unique_ptr<ProtoDatabase<WebsiteEvent>> website_event_db,std::unique_ptr<ProtoDatabase<Suspension>> suspension_db,std::unique_ptr<ProtoDatabase<TokenMapping>> token_mapping_db)51   MockUsageStatsDatabase(
52       std::unique_ptr<ProtoDatabase<WebsiteEvent>> website_event_db,
53       std::unique_ptr<ProtoDatabase<Suspension>> suspension_db,
54       std::unique_ptr<ProtoDatabase<TokenMapping>> token_mapping_db)
55       : UsageStatsDatabase(std::move(website_event_db),
56                            std::move(suspension_db),
57                            std::move(token_mapping_db)) {}
58 };
59 
60 class UsageStatsDatabaseTest : public testing::Test {
61  public:
UsageStatsDatabaseTest()62   UsageStatsDatabaseTest() {
63     auto fake_website_event_db =
64         std::make_unique<FakeDB<WebsiteEvent>>(&website_event_store_);
65     auto fake_suspension_db =
66         std::make_unique<FakeDB<Suspension>>(&suspension_store_);
67     auto fake_token_mapping_db =
68         std::make_unique<FakeDB<TokenMapping>>(&token_mapping_store_);
69 
70     // Maintain pointers to the FakeDBs for test callback execution.
71     website_event_db_unowned_ = fake_website_event_db.get();
72     suspension_db_unowned_ = fake_suspension_db.get();
73     token_mapping_db_unowned_ = fake_token_mapping_db.get();
74 
75     usage_stats_database_ = std::make_unique<MockUsageStatsDatabase>(
76         std::move(fake_website_event_db), std::move(fake_suspension_db),
77         std::move(fake_token_mapping_db));
78   }
79 
usage_stats_database()80   UsageStatsDatabase* usage_stats_database() {
81     return usage_stats_database_.get();
82   }
83 
fake_website_event_db()84   FakeDB<WebsiteEvent>* fake_website_event_db() {
85     return website_event_db_unowned_;
86   }
87 
fake_suspension_db()88   FakeDB<Suspension>* fake_suspension_db() { return suspension_db_unowned_; }
89 
fake_token_mapping_db()90   FakeDB<TokenMapping>* fake_token_mapping_db() {
91     return token_mapping_db_unowned_;
92   }
93 
94   MOCK_METHOD1(OnUpdateDone, void(UsageStatsDatabase::Error));
95   MOCK_METHOD2(OnGetEventsDone,
96                void(UsageStatsDatabase::Error, std::vector<WebsiteEvent>));
97   MOCK_METHOD2(OnGetSuspensionsDone,
98                void(UsageStatsDatabase::Error, std::vector<std::string>));
99   MOCK_METHOD2(OnGetTokenMappingsDone,
100                void(UsageStatsDatabase::Error, UsageStatsDatabase::TokenMap));
101 
102  private:
103   std::map<std::string, WebsiteEvent> website_event_store_;
104   std::map<std::string, Suspension> suspension_store_;
105   std::map<std::string, TokenMapping> token_mapping_store_;
106 
107   FakeDB<WebsiteEvent>* website_event_db_unowned_;
108   FakeDB<Suspension>* suspension_db_unowned_;
109   FakeDB<TokenMapping>* token_mapping_db_unowned_;
110 
111   std::unique_ptr<UsageStatsDatabase> usage_stats_database_;
112 
113   DISALLOW_COPY_AND_ASSIGN(UsageStatsDatabaseTest);
114 };
115 
TEST_F(UsageStatsDatabaseTest,Initialization)116 TEST_F(UsageStatsDatabaseTest, Initialization) {
117   ASSERT_NE(nullptr, usage_stats_database());
118   ASSERT_NE(nullptr, fake_website_event_db());
119   ASSERT_NE(nullptr, fake_suspension_db());
120   ASSERT_NE(nullptr, fake_token_mapping_db());
121 
122   // Expect that Init has been called on all ProtoDatabases. If it hasn't, then
123   // InitStatusCallback will cause the test to crash.
124   fake_website_event_db()->InitStatusCallback(
125       leveldb_proto::Enums::InitStatus::kOK);
126   fake_suspension_db()->InitStatusCallback(
127       leveldb_proto::Enums::InitStatus::kOK);
128   fake_token_mapping_db()->InitStatusCallback(
129       leveldb_proto::Enums::InitStatus::kOK);
130 }
131 
132 // Website Event Tests
TEST_F(UsageStatsDatabaseTest,GetAllEventsSuccess)133 TEST_F(UsageStatsDatabaseTest, GetAllEventsSuccess) {
134   fake_website_event_db()->InitStatusCallback(
135       leveldb_proto::Enums::InitStatus::kOK);
136 
137   usage_stats_database()->GetAllEvents(base::BindOnce(
138       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
139 
140   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
141                                      ElementsAre()));
142 
143   fake_website_event_db()->LoadCallback(true);
144 }
145 
TEST_F(UsageStatsDatabaseTest,GetAllEventsFailure)146 TEST_F(UsageStatsDatabaseTest, GetAllEventsFailure) {
147   fake_website_event_db()->InitStatusCallback(
148       leveldb_proto::Enums::InitStatus::kOK);
149 
150   usage_stats_database()->GetAllEvents(base::BindOnce(
151       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
152 
153   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kUnknownError,
154                                      ElementsAre()));
155 
156   fake_website_event_db()->LoadCallback(false);
157 }
158 
TEST_F(UsageStatsDatabaseTest,AddEventsEmpty)159 TEST_F(UsageStatsDatabaseTest, AddEventsEmpty) {
160   fake_website_event_db()->InitStatusCallback(
161       leveldb_proto::Enums::InitStatus::kOK);
162 
163   std::vector<WebsiteEvent> events;
164 
165   usage_stats_database()->AddEvents(
166       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
167                              base::Unretained(this)));
168 
169   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
170 
171   fake_website_event_db()->UpdateCallback(true);
172 }
173 
TEST_F(UsageStatsDatabaseTest,AddAndGetOneEvent)174 TEST_F(UsageStatsDatabaseTest, AddAndGetOneEvent) {
175   fake_website_event_db()->InitStatusCallback(
176       leveldb_proto::Enums::InitStatus::kOK);
177 
178   // Add 1 event.
179   WebsiteEvent event1 =
180       CreateWebsiteEvent(kFqdn1, 1, WebsiteEvent::START_BROWSING);
181   std::vector<WebsiteEvent> events({event1});
182 
183   usage_stats_database()->AddEvents(
184       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
185                              base::Unretained(this)));
186 
187   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
188 
189   fake_website_event_db()->UpdateCallback(true);
190 
191   // Get 1 event.
192   usage_stats_database()->GetAllEvents(base::BindOnce(
193       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
194 
195   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
196                                      ElementsAre(EqualsWebsiteEvent(event1))));
197 
198   fake_website_event_db()->LoadCallback(true);
199 }
200 
TEST_F(UsageStatsDatabaseTest,AddAndQueryEventsInRange)201 TEST_F(UsageStatsDatabaseTest, AddAndQueryEventsInRange) {
202   fake_website_event_db()->InitStatusCallback(
203       leveldb_proto::Enums::InitStatus::kOK);
204 
205   // Add 2 events at time 5s and 10s.
206   WebsiteEvent event1 =
207       CreateWebsiteEvent(kFqdn1, 5, WebsiteEvent::START_BROWSING);
208   WebsiteEvent event2 =
209       CreateWebsiteEvent(kFqdn2, 10, WebsiteEvent::STOP_BROWSING);
210   std::vector<WebsiteEvent> events({event1, event2});
211 
212   usage_stats_database()->AddEvents(
213       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
214                              base::Unretained(this)));
215 
216   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
217 
218   fake_website_event_db()->UpdateCallback(true);
219 
220   // Get events between time 0 (inclusive) and 9 (exclusive).
221   // This test validates the correct lexicographic ordering of timestamps such
222   // that key(0) <= key(5) < key(9) <= key(10).
223   usage_stats_database()->QueryEventsInRange(
224       base::Time::FromDoubleT(0), base::Time::FromDoubleT(9),
225       base::BindOnce(&UsageStatsDatabaseTest::OnGetEventsDone,
226                      base::Unretained(this)));
227 
228   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
229                                      ElementsAre(EqualsWebsiteEvent(event1))));
230 
231   fake_website_event_db()->LoadCallback(true);
232 }
233 
TEST_F(UsageStatsDatabaseTest,AddAndDeleteAllEvents)234 TEST_F(UsageStatsDatabaseTest, AddAndDeleteAllEvents) {
235   fake_website_event_db()->InitStatusCallback(
236       leveldb_proto::Enums::InitStatus::kOK);
237 
238   // Add 1 event.
239   WebsiteEvent event1 =
240       CreateWebsiteEvent(kFqdn1, 1, WebsiteEvent::START_BROWSING);
241   std::vector<WebsiteEvent> events({event1});
242 
243   usage_stats_database()->AddEvents(
244       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
245                              base::Unretained(this)));
246 
247   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
248 
249   fake_website_event_db()->UpdateCallback(true);
250 
251   // Delete all events.
252   usage_stats_database()->DeleteAllEvents(base::BindOnce(
253       &UsageStatsDatabaseTest::OnUpdateDone, base::Unretained(this)));
254 
255   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
256 
257   fake_website_event_db()->UpdateCallback(true);
258 
259   // Get all events (expecting none).
260   usage_stats_database()->GetAllEvents(base::BindOnce(
261       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
262 
263   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
264                                      ElementsAre()));
265 
266   fake_website_event_db()->LoadCallback(true);
267 }
268 
TEST_F(UsageStatsDatabaseTest,AddAndDeleteEventsInRange)269 TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsInRange) {
270   fake_website_event_db()->InitStatusCallback(
271       leveldb_proto::Enums::InitStatus::kOK);
272 
273   // Add 3 events.
274   WebsiteEvent event1 =
275       CreateWebsiteEvent(kFqdn1, 1, WebsiteEvent::START_BROWSING);
276   WebsiteEvent event2 =
277       CreateWebsiteEvent(kFqdn1, 2, WebsiteEvent::START_BROWSING);
278   WebsiteEvent event3 =
279       CreateWebsiteEvent(kFqdn1, 10, WebsiteEvent::START_BROWSING);
280   std::vector<WebsiteEvent> events({event1, event2, event3});
281 
282   usage_stats_database()->AddEvents(
283       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
284                              base::Unretained(this)));
285 
286   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
287 
288   fake_website_event_db()->UpdateCallback(true);
289 
290   // Delete events between time 1 (inclusive) and 10 (exclusive).
291   usage_stats_database()->DeleteEventsInRange(
292       base::Time::FromDoubleT(1), base::Time::FromDoubleT(10),
293       base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
294                      base::Unretained(this)));
295 
296   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
297 
298   fake_website_event_db()->LoadCallback(true);
299   fake_website_event_db()->UpdateCallback(true);
300 
301   // Get 1 remaining event outside range (at time 10).
302   usage_stats_database()->GetAllEvents(base::BindOnce(
303       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
304 
305   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
306                                      ElementsAre(EqualsWebsiteEvent(event3))));
307 
308   fake_website_event_db()->LoadCallback(true);
309 }
310 
TEST_F(UsageStatsDatabaseTest,ExpiryDeletesOldEvents)311 TEST_F(UsageStatsDatabaseTest, ExpiryDeletesOldEvents) {
312   fake_website_event_db()->InitStatusCallback(
313       leveldb_proto::Enums::InitStatus::kOK);
314 
315   // Add 3 events.
316   base::Time now = base::Time::NowFromSystemTime();
317   int64_t now_in_seconds = (int64_t)now.ToDoubleT();
318   WebsiteEvent event1 = CreateWebsiteEvent(kFqdn1, now_in_seconds + 1,
319                                            WebsiteEvent::START_BROWSING);
320   WebsiteEvent event2 = CreateWebsiteEvent(kFqdn1, now_in_seconds + 2,
321                                            WebsiteEvent::START_BROWSING);
322   WebsiteEvent event3 = CreateWebsiteEvent(kFqdn1, now_in_seconds + 10,
323                                            WebsiteEvent::START_BROWSING);
324   std::vector<WebsiteEvent> events({event1, event2, event3});
325 
326   usage_stats_database()->AddEvents(
327       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
328                              base::Unretained(this)));
329   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
330 
331   fake_website_event_db()->UpdateCallback(true);
332 
333   // Advance "now" by 7 days + 9 seconds so that the first two events are > 7
334   // days old.
335   now = now +
336         base::TimeDelta::FromDays(UsageStatsDatabase::EXPIRY_THRESHOLD_DAYS) +
337         base::TimeDelta::FromSeconds(9);
338   usage_stats_database()->ExpireEvents(now);
339 
340   fake_website_event_db()->LoadCallback(true);
341   fake_website_event_db()->UpdateCallback(true);
342 
343   usage_stats_database()->GetAllEvents(base::BindOnce(
344       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
345 
346   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
347                                      ElementsAre(EqualsWebsiteEvent(event3))));
348 
349   fake_website_event_db()->LoadCallback(true);
350 }
351 
TEST_F(UsageStatsDatabaseTest,AddAndDeleteEventsMatchingDomain)352 TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsMatchingDomain) {
353   fake_website_event_db()->InitStatusCallback(
354       leveldb_proto::Enums::InitStatus::kOK);
355 
356   // Add 3 events.
357   WebsiteEvent event1 =
358       CreateWebsiteEvent(kFqdn1, 1, WebsiteEvent::START_BROWSING);
359   WebsiteEvent event2 =
360       CreateWebsiteEvent(kFqdn1, 1, WebsiteEvent::STOP_BROWSING);
361   WebsiteEvent event3 =
362       CreateWebsiteEvent(kFqdn2, 1, WebsiteEvent::START_BROWSING);
363   std::vector<WebsiteEvent> events({event1, event2, event3});
364 
365   usage_stats_database()->AddEvents(
366       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
367                              base::Unretained(this)));
368 
369   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
370 
371   fake_website_event_db()->UpdateCallback(true);
372 
373   // Delete 2 events by FQDN.
374   base::flat_set<std::string> domains({kFqdn1});
375 
376   usage_stats_database()->DeleteEventsWithMatchingDomains(
377       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
378                               base::Unretained(this)));
379 
380   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
381 
382   fake_website_event_db()->UpdateCallback(true);
383 
384   // Get 1 remaining event with non-matching FQDN.
385   usage_stats_database()->GetAllEvents(base::BindOnce(
386       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
387 
388   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
389                                      ElementsAre(EqualsWebsiteEvent(event3))));
390 
391   fake_website_event_db()->LoadCallback(true);
392 }
393 
TEST_F(UsageStatsDatabaseTest,GetAllEventsDeferred)394 TEST_F(UsageStatsDatabaseTest, GetAllEventsDeferred) {
395   // Don't complete the database initialization yet.
396 
397   // Make request to database.
398   usage_stats_database()->GetAllEvents(base::BindOnce(
399       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
400 
401   // Expect callback to be run after initialization succeeds.
402   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
403                                      ElementsAre()));
404 
405   fake_website_event_db()->InitStatusCallback(
406       leveldb_proto::Enums::InitStatus::kOK);
407 
408   fake_website_event_db()->LoadCallback(true);
409 }
410 
411 // Suspension Tests
TEST_F(UsageStatsDatabaseTest,SetSuspensionsSuccess)412 TEST_F(UsageStatsDatabaseTest, SetSuspensionsSuccess) {
413   fake_suspension_db()->InitStatusCallback(
414       leveldb_proto::Enums::InitStatus::kOK);
415 
416   base::flat_set<std::string> domains({kFqdn1, kFqdn2});
417 
418   usage_stats_database()->SetSuspensions(
419       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
420                               base::Unretained(this)));
421 
422   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
423 
424   fake_suspension_db()->UpdateCallback(true);
425 }
426 
TEST_F(UsageStatsDatabaseTest,SetSuspensionsFailure)427 TEST_F(UsageStatsDatabaseTest, SetSuspensionsFailure) {
428   fake_suspension_db()->InitStatusCallback(
429       leveldb_proto::Enums::InitStatus::kOK);
430 
431   base::flat_set<std::string> domains({kFqdn1, kFqdn2});
432 
433   usage_stats_database()->SetSuspensions(
434       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
435                               base::Unretained(this)));
436 
437   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kUnknownError));
438 
439   fake_suspension_db()->UpdateCallback(false);
440 }
441 
TEST_F(UsageStatsDatabaseTest,GetAllSuspensionsSuccess)442 TEST_F(UsageStatsDatabaseTest, GetAllSuspensionsSuccess) {
443   fake_suspension_db()->InitStatusCallback(
444       leveldb_proto::Enums::InitStatus::kOK);
445 
446   std::vector<std::string> expected;
447 
448   usage_stats_database()->GetAllSuspensions(base::BindOnce(
449       &UsageStatsDatabaseTest::OnGetSuspensionsDone, base::Unretained(this)));
450 
451   EXPECT_CALL(*this, OnGetSuspensionsDone(UsageStatsDatabase::Error::kNoError,
452                                           expected));
453 
454   fake_suspension_db()->LoadCallback(true);
455 }
456 
TEST_F(UsageStatsDatabaseTest,GetAllSuspensionsFailure)457 TEST_F(UsageStatsDatabaseTest, GetAllSuspensionsFailure) {
458   fake_suspension_db()->InitStatusCallback(
459       leveldb_proto::Enums::InitStatus::kOK);
460 
461   std::vector<std::string> expected;
462 
463   usage_stats_database()->GetAllSuspensions(base::BindOnce(
464       &UsageStatsDatabaseTest::OnGetSuspensionsDone, base::Unretained(this)));
465 
466   EXPECT_CALL(*this, OnGetSuspensionsDone(
467                          UsageStatsDatabase::Error::kUnknownError, expected));
468 
469   fake_suspension_db()->LoadCallback(false);
470 }
471 
TEST_F(UsageStatsDatabaseTest,SetAndGetSuspension)472 TEST_F(UsageStatsDatabaseTest, SetAndGetSuspension) {
473   fake_suspension_db()->InitStatusCallback(
474       leveldb_proto::Enums::InitStatus::kOK);
475 
476   // Insert 1 suspension.
477   base::flat_set<std::string> domains({kFqdn1});
478 
479   usage_stats_database()->SetSuspensions(
480       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
481                               base::Unretained(this)));
482 
483   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
484 
485   fake_suspension_db()->UpdateCallback(true);
486 
487   // Get 1 suspension.
488   std::vector<std::string> expected({kFqdn1});
489 
490   usage_stats_database()->GetAllSuspensions(base::BindOnce(
491       &UsageStatsDatabaseTest::OnGetSuspensionsDone, base::Unretained(this)));
492 
493   EXPECT_CALL(*this, OnGetSuspensionsDone(UsageStatsDatabase::Error::kNoError,
494                                           expected));
495 
496   fake_suspension_db()->LoadCallback(true);
497 }
498 
TEST_F(UsageStatsDatabaseTest,SetRemoveAndGetSuspension)499 TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetSuspension) {
500   fake_suspension_db()->InitStatusCallback(
501       leveldb_proto::Enums::InitStatus::kOK);
502 
503   // Insert 2 suspensions.
504   base::flat_set<std::string> domains1({kFqdn1, kFqdn2});
505 
506   usage_stats_database()->SetSuspensions(
507       domains1, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
508                                base::Unretained(this)));
509 
510   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
511 
512   fake_suspension_db()->UpdateCallback(true);
513 
514   // Insert 1 suspension, and remove the other.
515   base::flat_set<std::string> domains2({kFqdn1});
516 
517   usage_stats_database()->SetSuspensions(
518       domains2, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
519                                base::Unretained(this)));
520 
521   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
522 
523   fake_suspension_db()->UpdateCallback(true);
524 
525   // Get 1 suspension.
526   std::vector<std::string> expected({kFqdn1});
527 
528   usage_stats_database()->GetAllSuspensions(base::BindOnce(
529       &UsageStatsDatabaseTest::OnGetSuspensionsDone, base::Unretained(this)));
530 
531   EXPECT_CALL(*this, OnGetSuspensionsDone(UsageStatsDatabase::Error::kNoError,
532                                           expected));
533 
534   fake_suspension_db()->LoadCallback(true);
535 }
536 
TEST_F(UsageStatsDatabaseTest,SetAndGetSuspensionDeferred)537 TEST_F(UsageStatsDatabaseTest, SetAndGetSuspensionDeferred) {
538   // Fail to initialize the database.
539   fake_suspension_db()->InitStatusCallback(
540       leveldb_proto::Enums::InitStatus::kError);
541 
542   // Insert 1 suspension.
543   base::flat_set<std::string> domains({kFqdn1});
544 
545   usage_stats_database()->SetSuspensions(
546       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
547                               base::Unretained(this)));
548 
549   // Get 1 suspension.
550   std::vector<std::string> expected({kFqdn1});
551 
552   usage_stats_database()->GetAllSuspensions(base::BindOnce(
553       &UsageStatsDatabaseTest::OnGetSuspensionsDone, base::Unretained(this)));
554 
555   // Now successfully initialize database, and expect previous callbacks to be
556   // run.
557   fake_suspension_db()->InitStatusCallback(
558       leveldb_proto::Enums::InitStatus::kOK);
559 
560   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
561 
562   fake_suspension_db()->UpdateCallback(true);
563 
564   EXPECT_CALL(*this, OnGetSuspensionsDone(UsageStatsDatabase::Error::kNoError,
565                                           expected));
566 
567   fake_suspension_db()->LoadCallback(true);
568 }
569 
570 // Token Mapping Tests
TEST_F(UsageStatsDatabaseTest,SetTokenMappingsSuccess)571 TEST_F(UsageStatsDatabaseTest, SetTokenMappingsSuccess) {
572   fake_token_mapping_db()->InitStatusCallback(
573       leveldb_proto::Enums::InitStatus::kOK);
574 
575   UsageStatsDatabase::TokenMap mappings({{kToken1, kFqdn1}, {kToken2, kFqdn2}});
576 
577   usage_stats_database()->SetTokenMappings(
578       mappings, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
579                                base::Unretained(this)));
580 
581   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
582 
583   fake_token_mapping_db()->UpdateCallback(true);
584 }
585 
TEST_F(UsageStatsDatabaseTest,SetTokenMappingsFailure)586 TEST_F(UsageStatsDatabaseTest, SetTokenMappingsFailure) {
587   fake_token_mapping_db()->InitStatusCallback(
588       leveldb_proto::Enums::InitStatus::kOK);
589 
590   UsageStatsDatabase::TokenMap mappings({{kToken1, kFqdn1}, {kToken2, kFqdn2}});
591 
592   usage_stats_database()->SetTokenMappings(
593       mappings, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
594                                base::Unretained(this)));
595 
596   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kUnknownError));
597 
598   fake_token_mapping_db()->UpdateCallback(false);
599 }
600 
TEST_F(UsageStatsDatabaseTest,GetAllTokenMappingsSuccess)601 TEST_F(UsageStatsDatabaseTest, GetAllTokenMappingsSuccess) {
602   fake_token_mapping_db()->InitStatusCallback(
603       leveldb_proto::Enums::InitStatus::kOK);
604 
605   UsageStatsDatabase::TokenMap expected;
606 
607   usage_stats_database()->GetAllTokenMappings(base::BindOnce(
608       &UsageStatsDatabaseTest::OnGetTokenMappingsDone, base::Unretained(this)));
609 
610   EXPECT_CALL(*this, OnGetTokenMappingsDone(UsageStatsDatabase::Error::kNoError,
611                                             expected));
612 
613   fake_token_mapping_db()->LoadCallback(true);
614 }
615 
TEST_F(UsageStatsDatabaseTest,GetAllTokenMappingsFailure)616 TEST_F(UsageStatsDatabaseTest, GetAllTokenMappingsFailure) {
617   fake_token_mapping_db()->InitStatusCallback(
618       leveldb_proto::Enums::InitStatus::kOK);
619 
620   UsageStatsDatabase::TokenMap expected;
621 
622   usage_stats_database()->GetAllTokenMappings(base::BindOnce(
623       &UsageStatsDatabaseTest::OnGetTokenMappingsDone, base::Unretained(this)));
624 
625   EXPECT_CALL(*this, OnGetTokenMappingsDone(
626                          UsageStatsDatabase::Error::kUnknownError, expected));
627 
628   fake_token_mapping_db()->LoadCallback(false);
629 }
630 
TEST_F(UsageStatsDatabaseTest,SetAndGetTokenMapping)631 TEST_F(UsageStatsDatabaseTest, SetAndGetTokenMapping) {
632   fake_token_mapping_db()->InitStatusCallback(
633       leveldb_proto::Enums::InitStatus::kOK);
634 
635   UsageStatsDatabase::TokenMap mapping({{kToken1, kFqdn1}});
636 
637   // Insert 1 token mapping.
638   usage_stats_database()->SetTokenMappings(
639       mapping, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
640                               base::Unretained(this)));
641 
642   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
643 
644   fake_token_mapping_db()->UpdateCallback(true);
645 
646   // Get 1 token mapping.
647   usage_stats_database()->GetAllTokenMappings(base::BindOnce(
648       &UsageStatsDatabaseTest::OnGetTokenMappingsDone, base::Unretained(this)));
649 
650   EXPECT_CALL(*this, OnGetTokenMappingsDone(UsageStatsDatabase::Error::kNoError,
651                                             mapping));
652 
653   fake_token_mapping_db()->LoadCallback(true);
654 }
655 
TEST_F(UsageStatsDatabaseTest,SetRemoveAndGetTokenMapping)656 TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetTokenMapping) {
657   fake_token_mapping_db()->InitStatusCallback(
658       leveldb_proto::Enums::InitStatus::kOK);
659 
660   // Insert 2 token mappings.
661   UsageStatsDatabase::TokenMap mappings1(
662       {{kToken1, kFqdn1}, {kToken2, kFqdn2}});
663 
664   usage_stats_database()->SetTokenMappings(
665       mappings1, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
666                                 base::Unretained(this)));
667 
668   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
669 
670   fake_token_mapping_db()->UpdateCallback(true);
671 
672   // Re-insert 1 token mapping, and remove the other.apping) {
673   UsageStatsDatabase::TokenMap mappings2({{kToken1, kFqdn1}});
674 
675   usage_stats_database()->SetTokenMappings(
676       mappings2, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
677                                 base::Unretained(this)));
678 
679   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
680 
681   fake_token_mapping_db()->UpdateCallback(true);
682 
683   // Get 1 remaining token mapping.
684   usage_stats_database()->GetAllTokenMappings(base::BindOnce(
685       &UsageStatsDatabaseTest::OnGetTokenMappingsDone, base::Unretained(this)));
686 
687   EXPECT_CALL(*this, OnGetTokenMappingsDone(UsageStatsDatabase::Error::kNoError,
688                                             mappings2));
689 
690   fake_token_mapping_db()->LoadCallback(true);
691 }
692 
TEST_F(UsageStatsDatabaseTest,SetTokenMappingsUninitialized)693 TEST_F(UsageStatsDatabaseTest, SetTokenMappingsUninitialized) {
694   // Fail to initialize database.
695   fake_token_mapping_db()->InitStatusCallback(
696       leveldb_proto::Enums::InitStatus::kError);
697 
698   UsageStatsDatabase::TokenMap mappings({{kToken1, kFqdn1}, {kToken2, kFqdn2}});
699 
700   // Expect callback will not be run.
701   EXPECT_CALL(*this, OnUpdateDone).Times(0);
702 
703   usage_stats_database()->SetTokenMappings(
704       mappings, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
705                                base::Unretained(this)));
706 }
707 
708 }  // namespace usage_stats
709