Start plumbing connections from the BudgetManager to the BudgetDatabase

chrome/browser/budget_service/budget_database_unittest.cc

Index: chrome/browser/budget_service/budget_database_unittest.cc
diff --git a/chrome/browser/budget_service/budget_database_unittest.cc b/chrome/browser/budget_service/budget_database_unittest.cc
index eb83fff049d8e10499ed55705cac50b2d45313f8..1794610bdd055d2837cd205b7f059b7133db5f3b 100644
--- a/chrome/browser/budget_service/budget_database_unittest.cc
+++ b/chrome/browser/budget_service/budget_database_unittest.cc
@@ -8,6 +8,7 @@
 #include "base/test/simple_test_clock.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "chrome/browser/budget_service/budget.pb.h"
+#include "chrome/browser/engagement/site_engagement_service.h"
 #include "chrome/test/base/testing_profile.h"
 #include "components/leveldb_proto/proto_database.h"
 #include "components/leveldb_proto/proto_database_impl.h"
@@ -17,8 +18,6 @@
 
 namespace {
 
-const double kDefaultBudget1 = 1.234;
-const double kDefaultBudget2 = 2.345;
 const double kDefaultExpirationInHours = 72;
 const double kDefaultEngagement = 30.0;
 
@@ -30,39 +29,15 @@ class BudgetDatabaseTest : public ::testing::Test {
  public:
   BudgetDatabaseTest()
       : success_(false),
-        db_(profile_.GetPath().Append(FILE_PATH_LITERAL("BudgetDabase")),
+        db_(&profile_,
+            profile_.GetPath().Append(FILE_PATH_LITERAL("BudgetDatabase")),
             base::ThreadTaskRunnerHandle::Get()) {}
 
-  // The BudgetDatabase assumes that a budget will always be queried before it
-  // is written to. Use GetBudgetDetails() to pre-populate the cache.
-  void SetUp() override { GetBudgetDetails(); }
-
   void WriteBudgetComplete(base::Closure run_loop_closure, bool success) {
     success_ = success;
     run_loop_closure.Run();
   }
 
-  // Add budget to the origin.
-  bool AddBudget(const GURL& origin, double amount) {
-    base::RunLoop run_loop;
-    db_.AddBudget(origin, amount,
-                  base::Bind(&BudgetDatabaseTest::WriteBudgetComplete,
-                             base::Unretained(this), run_loop.QuitClosure()));
-    run_loop.Run();
-    return success_;
-  }
-
-  // Add engagement based budget to the origin.
-  bool AddEngagementBudget(const GURL& origin, double sesScore) {
-    base::RunLoop run_loop;
-    db_.AddEngagementBudget(
-        origin, sesScore,
-        base::Bind(&BudgetDatabaseTest::WriteBudgetComplete,
-                   base::Unretained(this), run_loop.QuitClosure()));
-    run_loop.Run();
-    return success_;
-  }
-
   // Spend budget for the origin.
   bool SpendBudget(const GURL& origin, double amount) {
     base::RunLoop run_loop;
@@ -79,11 +54,6 @@ class BudgetDatabaseTest : public ::testing::Test {
       const BudgetDatabase::BudgetPrediction& prediction) {
     success_ = success;
     // Convert BudgetPrediction to a vector for random access to check values.
-    /*
-    prediction_.clear();
-    for (auto& chunk : prediction)
-      prediction_.push_back(chunk);
-      */
     prediction_.assign(prediction.begin(), prediction.end());
     run_loop_closure.Run();
   }
@@ -107,8 +77,10 @@ class BudgetDatabaseTest : public ::testing::Test {
     return clock;
   }
 
-  // Query the database to check if the origin is in the cache.
-  bool IsCached(const GURL& origin) { return db_.IsCached(origin); }
+  void SetSiteEngagementScore(const GURL& url, double score) {
+    SiteEngagementService* service = SiteEngagementService::Get(&profile_);
+    service->ResetScoreForURL(url, score);
+  }
 
  protected:
   bool success_;
@@ -121,74 +93,14 @@ class BudgetDatabaseTest : public ::testing::Test {
   BudgetDatabase db_;
 };
 
-TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
-  const GURL origin(kTestOrigin);
-  base::SimpleTestClock* clock = SetClockForTesting();
-  base::TimeDelta expiration(
-      base::TimeDelta::FromHours(kDefaultExpirationInHours));
-  base::Time starting_time = clock->Now();
-  base::Time expiration_time = clock->Now() + expiration;
-
-  // Add two budget chunks with different expirations (default expiration and
-  // default expiration + 1 day).
-  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
-  clock->Advance(base::TimeDelta::FromDays(1));
-  ASSERT_TRUE(AddBudget(origin, kDefaultBudget2));
-
-  // Get the budget.
-  GetBudgetDetails();
-
-  // Get the prediction and validate it.
-  ASSERT_TRUE(success_);
-  ASSERT_EQ(3U, prediction_.size());
-
-  // Make sure that the correct data is returned.
-  // First value should be [total_budget, now]
-  EXPECT_EQ(kDefaultBudget1 + kDefaultBudget2, prediction_[0].budget_at);
-  EXPECT_EQ(clock->Now(), prediction_[0].time);
-
-  // The next value should be the budget after the first chunk expires.
-  EXPECT_EQ(kDefaultBudget2, prediction_[1].budget_at);
-  EXPECT_EQ(expiration_time, prediction_[1].time);
-
-  // The final value gives the budget of 0.0 after the second chunk expires.
-  expiration_time += base::TimeDelta::FromDays(1);
-  EXPECT_EQ(0, prediction_[2].budget_at);
-  EXPECT_EQ(expiration_time, prediction_[2].time);
-
-  // Advance the time until the first chunk of budget should be expired.
-  clock->SetNow(starting_time +
-                base::TimeDelta::FromHours(kDefaultExpirationInHours));
-
-  // Get the new budget and check that kDefaultBudget1 has been removed.
-  GetBudgetDetails();
-  ASSERT_EQ(2U, prediction_.size());
-  EXPECT_EQ(kDefaultBudget2, prediction_[0].budget_at);
-  EXPECT_EQ(0, prediction_[1].budget_at);
-
-  // Advace the time until both chunks of budget should be expired.
-  clock->SetNow(starting_time +
-                base::TimeDelta::FromHours(kDefaultExpirationInHours) +
-                base::TimeDelta::FromDays(1));
-
-  GetBudgetDetails();
-  ASSERT_EQ(1U, prediction_.size());
-  EXPECT_EQ(0, prediction_[0].budget_at);
-
-  // Now that the entire budget has expired, check that the entry in the map
-  // has been removed.
-  EXPECT_FALSE(IsCached(origin));
-}
-
 TEST_F(BudgetDatabaseTest, AddEngagementBudgetTest) {
   const GURL origin(kTestOrigin);
   base::SimpleTestClock* clock = SetClockForTesting();
   base::Time expiration_time =
       clock->Now() + base::TimeDelta::FromHours(kDefaultExpirationInHours);
 
-  // Add a chunk of budget to a non-existant origin. This should add the full
-  // amount of engagement.
-  ASSERT_TRUE(AddEngagementBudget(origin, kDefaultEngagement));
+  // Set the default site engagement.
+  SetSiteEngagementScore(origin, kDefaultEngagement);
 
   // The budget should include a full share of the engagement.
   GetBudgetDetails();
@@ -200,10 +112,9 @@ TEST_F(BudgetDatabaseTest, AddEngagementBudgetTest) {
 
   // Advance time 1 day and add more engagement budget.
   clock->Advance(base::TimeDelta::FromDays(1));
-  ASSERT_TRUE(AddEngagementBudget(origin, kDefaultEngagement));
+  GetBudgetDetails();
 
   // The budget should now have 1 full share plus 1/3 share.
-  GetBudgetDetails();
   ASSERT_TRUE(success_);
   ASSERT_EQ(3U, prediction_.size());
   ASSERT_DOUBLE_EQ(kDefaultEngagement * 4 / 3, prediction_[0].budget_at);
@@ -216,7 +127,6 @@ TEST_F(BudgetDatabaseTest, AddEngagementBudgetTest) {
   // Advance time by 59 minutes and check that no engagement budget is added
   // since budget should only be added for > 1 hour increments.
   clock->Advance(base::TimeDelta::FromMinutes(59));
-  ASSERT_TRUE(AddEngagementBudget(origin, kDefaultEngagement));
 
   // The budget should be the same as before the attempted add.
   GetBudgetDetails();
@@ -230,42 +140,48 @@ TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
   base::SimpleTestClock* clock = SetClockForTesting();
   base::Time starting_time = clock->Now();
 
+  // Set the default site engagement.
+  SetSiteEngagementScore(origin, kDefaultEngagement);
+
   // Intialize the budget with several chunks.
-  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
+  GetBudgetDetails();
   clock->Advance(base::TimeDelta::FromDays(1));
-  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
+  GetBudgetDetails();
   clock->Advance(base::TimeDelta::FromDays(1));
-  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
+  GetBudgetDetails();
 
-  // Reset the clock then spend an amount of budget less than kDefaultBudget.
+  // Reset the clock then spend an amount of budget less than
+  // kDefaultEngagement.
   clock->SetNow(starting_time);

[Peter Beverloo, 2016/08/25 15:31:19]

This is a bit of a weird paradigm— by adding values in the future and then
resetting the clock "back to the past", this test is now in a state which is
unlikely for production code to reach. (Unless the user messages with their
clock.)

What is significant about resetting the clock?

[harkness, 2016/08/25 16:36:13]

That's actually an artifact from when chunks partially decayed over time. The
only times we could be sure of values for testing was at the start. Now that we
have discrete chunks that expire entirely at a fixed time, I can remove this.

   ASSERT_TRUE(SpendBudget(origin, 1));
   GetBudgetDetails();
 
-  // There should still be three chunks of budget of size kDefaultBudget-1,
-  // kDefaultBudget, and kDefaultBudget.
+  // There should still be three chunks of budget of size kDefaultEngagement-1,
+  // kDefaultEngagement, and kDefaultEngagement.
   ASSERT_EQ(4U, prediction_.size());
-  ASSERT_DOUBLE_EQ(kDefaultBudget1 * 3 - 1, prediction_[0].budget_at);
-  ASSERT_DOUBLE_EQ(kDefaultBudget1 * 2, prediction_[1].budget_at);
-  ASSERT_DOUBLE_EQ(kDefaultBudget1, prediction_[2].budget_at);
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 5 / 3 - 1, prediction_[0].budget_at);
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 2 / 3, prediction_[1].budget_at);
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 1 / 3, prediction_[2].budget_at);
   ASSERT_DOUBLE_EQ(0, prediction_[3].budget_at);
 
   // Now spend enough that it will use up the rest of the first chunk and all of
   // the second chunk, but not all of the third chunk.
-  ASSERT_TRUE(SpendBudget(origin, kDefaultBudget1 * 2));
+  ASSERT_TRUE(SpendBudget(origin, kDefaultEngagement * 4 / 3));
   GetBudgetDetails();
   ASSERT_EQ(2U, prediction_.size());
-  ASSERT_DOUBLE_EQ(kDefaultBudget1 - 1, prediction_.begin()->budget_at);
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 1 / 3 - 1,
+                   prediction_.begin()->budget_at);
 
   // Validate that the code returns false if SpendBudget tries to spend more
   // budget than the origin has.
-  EXPECT_FALSE(SpendBudget(origin, kDefaultBudget1));
+  EXPECT_FALSE(SpendBudget(origin, kDefaultEngagement));
   GetBudgetDetails();
   ASSERT_EQ(2U, prediction_.size());
-  ASSERT_DOUBLE_EQ(kDefaultBudget1 - 1, prediction_.begin()->budget_at);
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 1 / 3 - 1,
+                   prediction_.begin()->budget_at);
 
   // Advance time until the last remaining chunk should be expired, then query
-  // for what would be a valid amount of budget if the chunks weren't expired.
+  // for the full engagement worth of budget.
   clock->Advance(base::TimeDelta::FromDays(6));
-  EXPECT_FALSE(SpendBudget(origin, 0.01));
+  EXPECT_TRUE(SpendBudget(origin, kDefaultEngagement));
 }