AddEngagementBudget and SpendBudget added to 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 a8b7930d15f183c1396d3ba93667fe375bba6dba..eb83fff049d8e10499ed55705cac50b2d45313f8 100644
--- a/chrome/browser/budget_service/budget_database_unittest.cc
+++ b/chrome/browser/budget_service/budget_database_unittest.cc
@@ -20,6 +20,7 @@ namespace {
 const double kDefaultBudget1 = 1.234;
 const double kDefaultBudget2 = 2.345;
 const double kDefaultExpirationInHours = 72;
+const double kDefaultEngagement = 30.0;
 
 const char kTestOrigin[] = "https://example.com";
 
@@ -36,7 +37,7 @@ class BudgetDatabaseTest : public ::testing::Test {
   // is written to. Use GetBudgetDetails() to pre-populate the cache.
   void SetUp() override { GetBudgetDetails(); }
 
-  void AddBudgetComplete(base::Closure run_loop_closure, bool success) {
+  void WriteBudgetComplete(base::Closure run_loop_closure, bool success) {
     success_ = success;
     run_loop_closure.Run();
   }
@@ -45,22 +46,49 @@ class BudgetDatabaseTest : public ::testing::Test {
   bool AddBudget(const GURL& origin, double amount) {
     base::RunLoop run_loop;
     db_.AddBudget(origin, amount,
-                  base::Bind(&BudgetDatabaseTest::AddBudgetComplete,
+                  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;
+    db_.SpendBudget(origin, amount,
+                    base::Bind(&BudgetDatabaseTest::WriteBudgetComplete,
+                               base::Unretained(this), run_loop.QuitClosure()));
+    run_loop.Run();
+    return success_;
+  }
+
   void GetBudgetDetailsComplete(
       base::Closure run_loop_closure,
       bool success,
-      const BudgetDatabase::BudgetExpectation& expectation) {
+      const BudgetDatabase::BudgetPrediction& prediction) {
     success_ = success;
-    expectation_ = expectation;
+    // 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();
   }
 
-  // Get the full set of budget expectations for the origin.
+  // Get the full set of budget predictions for the origin.
   void GetBudgetDetails() {
     base::RunLoop run_loop;
     db_.GetBudgetDetails(
@@ -71,9 +99,6 @@ class BudgetDatabaseTest : public ::testing::Test {
   }
 
   Profile* profile() { return &profile_; }
-  const BudgetDatabase::BudgetExpectation& expectation() {
-    return expectation_;
-  }
 
   // Setup a test clock so that the tests can control time.
   base::SimpleTestClock* SetClockForTesting() {
@@ -87,13 +112,13 @@ class BudgetDatabaseTest : public ::testing::Test {
 
  protected:
   bool success_;
+  std::vector<BudgetDatabase::BudgetStatus> prediction_;
 
  private:
   content::TestBrowserThreadBundle thread_bundle_;
   std::unique_ptr<budget_service::Budget> budget_;
   TestingProfile profile_;
   BudgetDatabase db_;
-  BudgetDatabase::BudgetExpectation expectation_;
 };
 
 TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
@@ -113,28 +138,23 @@ TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
   // Get the budget.
   GetBudgetDetails();
 
-  // Get the expectation and validate it.
-  const auto& expected_value = expectation();
+  // Get the prediction and validate it.
   ASSERT_TRUE(success_);
-  ASSERT_EQ(3U, expected_value.size());
+  ASSERT_EQ(3U, prediction_.size());
 
   // Make sure that the correct data is returned.
-  auto iter = expected_value.begin();
-
   // First value should be [total_budget, now]
-  EXPECT_EQ(kDefaultBudget1 + kDefaultBudget2, iter->budget_at);
-  EXPECT_EQ(clock->Now(), iter->time);
+  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.
-  iter++;
-  EXPECT_EQ(kDefaultBudget2, iter->budget_at);
-  EXPECT_EQ(expiration_time, iter->time);
+  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);
-  iter++;
-  EXPECT_EQ(0, iter->budget_at);
-  EXPECT_EQ(expiration_time, iter->time);
+  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 +
@@ -142,11 +162,9 @@ TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
 
   // Get the new budget and check that kDefaultBudget1 has been removed.
   GetBudgetDetails();
-  iter = expectation().begin();
-  ASSERT_EQ(2U, expectation().size());
-  EXPECT_EQ(kDefaultBudget2, iter->budget_at);
-  iter++;
-  EXPECT_EQ(0, iter->budget_at);
+  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 +
@@ -154,11 +172,100 @@ TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
                 base::TimeDelta::FromDays(1));
 
   GetBudgetDetails();
-  iter = expectation().begin();
-  ASSERT_EQ(1U, expectation().size());
-  EXPECT_EQ(0, iter->budget_at);
+  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));
+
+  // The budget should include a full share of the engagement.
+  GetBudgetDetails();
+  ASSERT_TRUE(success_);
+  ASSERT_EQ(2U, prediction_.size());
+  ASSERT_EQ(kDefaultEngagement, prediction_[0].budget_at);
+  ASSERT_EQ(0, prediction_[1].budget_at);
+  ASSERT_EQ(expiration_time, prediction_[1].time);
+
+  // Advance time 1 day and add more engagement budget.
+  clock->Advance(base::TimeDelta::FromDays(1));
+  ASSERT_TRUE(AddEngagementBudget(origin, kDefaultEngagement));
+
+  // 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);
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 1 / 3, prediction_[1].budget_at);
+  ASSERT_EQ(expiration_time, prediction_[1].time);
+  ASSERT_EQ(0, prediction_[2].budget_at);
+  ASSERT_EQ(expiration_time + base::TimeDelta::FromDays(1),
+            prediction_[2].time);
+
+  // 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();
+  ASSERT_TRUE(success_);
+  ASSERT_EQ(3U, prediction_.size());
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 4 / 3, prediction_[0].budget_at);
+}
+
+TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
+  const GURL origin(kTestOrigin);
+  base::SimpleTestClock* clock = SetClockForTesting();
+  base::Time starting_time = clock->Now();
+
+  // Intialize the budget with several chunks.
+  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
+  clock->Advance(base::TimeDelta::FromDays(1));
+  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
+  clock->Advance(base::TimeDelta::FromDays(1));
+  ASSERT_TRUE(AddBudget(origin, kDefaultBudget1));
+
+  // Reset the clock then spend an amount of budget less than kDefaultBudget.
+  clock->SetNow(starting_time);
+  ASSERT_TRUE(SpendBudget(origin, 1));
+  GetBudgetDetails();
+
+  // There should still be three chunks of budget of size kDefaultBudget-1,
+  // kDefaultBudget, and kDefaultBudget.
+  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(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));
+  GetBudgetDetails();
+  ASSERT_EQ(2U, prediction_.size());
+  ASSERT_DOUBLE_EQ(kDefaultBudget1 - 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));
+  GetBudgetDetails();
+  ASSERT_EQ(2U, prediction_.size());
+  ASSERT_DOUBLE_EQ(kDefaultBudget1 - 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.
+  clock->Advance(base::TimeDelta::FromDays(6));
+  EXPECT_FALSE(SpendBudget(origin, 0.01));
+}