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..465f5ab02a720593a9c0cbb673ea6bbec4e46db9 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,12 +46,33 @@ 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,
@@ -162,3 +184,72 @@ TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
   // has been removed.
   EXPECT_FALSE(IsCached(origin));
 }
+
+TEST_F(BudgetDatabaseTest, AddEngagementBudgetTest) {
+  const GURL origin(kTestOrigin);
+  base::SimpleTestClock* clock = SetClockForTesting();
+
+  // Add a chunk of budget to a non-existant origin. This should add the full
+  // amount of engagement.
+  ASSERT_TRUE(AddEngagementBudget(origin, kDefaultEngagement));
+
+  // Get the budget.
+  GetBudgetDetails();
+

[johnme, 2016/08/19 18:48:36]

Nit: remove newline, since the ASSERTs below are tied to this call to
GetBudgetDetails.

[harkness, 2016/08/22 13:46:26]

Done.

+  // The budget should include a full share of the engagement. Just check the
+  // value of the budget state at the start of the list for the total budget.
+  const auto& expected_value = expectation();

[johnme, 2016/08/19 18:48:36]

Nit: Don't bother caching this reference; the compiler can just inline the
expectation() getter. Or just make expectation_ protected and access it
directly, like success_?

[harkness, 2016/08/22 13:46:26]

Made prediction_ protected, it just makes the code much simpler.

+  ASSERT_TRUE(success_);
+  ASSERT_EQ(2U, expected_value.size());
+  ASSERT_EQ(kDefaultEngagement, expectation().begin()->budget_at);

[johnme, 2016/08/19 18:48:36]

s/expectation()/expected_value/ since you've cached a reference to it

[harkness, 2016/08/22 13:46:26]

Directly accesses prediction_ now.

+
+  // Advance time 1 day and add more engagement budget.
+  clock->Advance(base::TimeDelta::FromDays(1));
+  ASSERT_TRUE(AddEngagementBudget(origin, kDefaultEngagement));
+  GetBudgetDetails();

[johnme, 2016/08/19 18:48:36]

Add newline before getBudgetDetails()

[harkness, 2016/08/22 13:46:26]

Done.

+

[johnme, 2016/08/19 18:48:36]

Ditto remove newline

[harkness, 2016/08/22 13:46:26]

Done.

+  // The budget should now have 1 full share plus 1/3 share.
+  ASSERT_TRUE(success_);
+  ASSERT_EQ(3U, expectation().size());
+  ASSERT_DOUBLE_EQ(kDefaultEngagement * 4 / 3,

[johnme, 2016/08/19 18:48:36]

Might be nice to re-test the other 2 values too, to make sure they haven't
changed. Ditto above, with testing the 2nd value. To make this easier, consider
making expectation_ be a std::vector (initialised as a copy of the std::list in
GetBudgetDetailsComplete).

[harkness, 2016/08/22 13:46:26]

Good idea on converting the member to be a vector. I've done that and cleaned up
the code to use it.

I've added some checks on the remaining values in this case, although it ends up
duplicating the logic checking that was done in the AddBudget case.

+                   expectation().begin()->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, expectation().size());
+  ASSERT_DOUBLE_EQ(kDefaultBudget1 * 3 - 1, expectation().begin()->budget_at);

[johnme, 2016/08/19 18:48:36]

Ditto with checking the other values in expectation()?

[harkness, 2016/08/22 13:46:26]

Done.

+
+  // 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, expectation().size());
+  ASSERT_DOUBLE_EQ(kDefaultBudget1 - 1, expectation().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));

[johnme, 2016/08/19 18:48:36]

After this, check that the budget hasn't changed?

[harkness, 2016/08/22 13:46:26]

Done.

+
+  // 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));
+}