AddEngagementBudget and SpendBudget added to BudgetDatabase.

chrome/browser/budget_service/budget_database_unittest.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "chrome/browser/budget_service/budget_database.h"
 
 #include "base/run_loop.h"
 #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/test/base/testing_profile.h"
 #include "components/leveldb_proto/proto_database.h"
 #include "components/leveldb_proto/proto_database_impl.h"
 #include "content/public/browser/browser_thread.h"
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 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";
 
 }  // namespace
 
 class BudgetDatabaseTest : public ::testing::Test {
  public:
   BudgetDatabaseTest()
       : success_(false),
         db_(profile_.GetPath().Append(FILE_PATH_LITERAL("BudgetDabase")),
             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 AddBudgetComplete(base::Closure run_loop_closure, bool success) {
+  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::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) {
     success_ = success;
     expectation_ = expectation;
     run_loop_closure.Run();
   }
 
   // Get the full set of budget expectations for the origin.
   void GetBudgetDetails() {
     base::RunLoop run_loop;
     db_.GetBudgetDetails(
         GURL(kTestOrigin),
         base::Bind(&BudgetDatabaseTest::GetBudgetDetailsComplete,
                    base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
   }
 
   Profile* profile() { return &profile_; }
   const BudgetDatabase::BudgetExpectation& expectation() {

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

Nit: newline before please

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

Done.

     return expectation_;
   }
 
   // Setup a test clock so that the tests can control time.
   base::SimpleTestClock* SetClockForTesting() {
     base::SimpleTestClock* clock = new base::SimpleTestClock();
     db_.SetClockForTesting(base::WrapUnique(clock));
     return clock;
   }
 
@@ skipping 70 matching lines @@
 
   GetBudgetDetails();
   iter = expectation().begin();
   ASSERT_EQ(1U, expectation().size());
   EXPECT_EQ(0, iter->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();
+
+  // 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));
+}