Adding more functionality 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 3d34e06663735018885742b6c7005964a3636db7..f10a2ba9217a73978ed2a924b29be9426fdfa798 100644
--- a/chrome/browser/budget_service/budget_database_unittest.cc
+++ b/chrome/browser/budget_service/budget_database_unittest.cc
@@ -5,6 +5,7 @@
 #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"
@@ -16,13 +17,14 @@
 
 namespace {
 
-const double kDefaultDebt = -0.72;
 const double kDefaultBudget1 = 1.234;
 const double kDefaultBudget2 = 2.345;
-const double kDefaultExpiration = 3600000000;
+const double kDefaultExpirationInHours = 72;
 
 const char kTestOrigin[] = "https://example.com";
 
+}  // namespace
+
 class BudgetDatabaseTest : public ::testing::Test {
  public:
   BudgetDatabaseTest()
@@ -30,49 +32,23 @@ class BudgetDatabaseTest : public ::testing::Test {
         db_(profile_.GetPath().Append(FILE_PATH_LITERAL("BudgetDabase")),
             base::ThreadTaskRunnerHandle::Get()) {}
 
-  void SetBudgetComplete(base::Closure run_loop_closure, bool success) {
-    success_ = success;
-    run_loop_closure.Run();
-  }
-
-  // Set up basic default values.
-  bool SetBudgetWithDefaultValues() {
-    budget_service::Budget budget;
-    budget.set_debt(kDefaultDebt);
-
-    // Create two chunks and give them default values.
-    budget_service::BudgetChunk* budget_chunk = budget.add_budget();
-    budget_chunk->set_amount(kDefaultBudget1);
-    budget_chunk->set_expiration(kDefaultExpiration);
-    budget_chunk = budget.add_budget();
-    budget_chunk->set_amount(kDefaultBudget2);
-    budget_chunk->set_expiration(kDefaultExpiration + 1);
-
-    base::RunLoop run_loop;
-    db_.SetValue(GURL(kTestOrigin), budget,
-                 base::Bind(&BudgetDatabaseTest::SetBudgetComplete,
-                            base::Unretained(this), run_loop.QuitClosure()));
-    run_loop.Run();
-
-    return success_;
-  }
+  // 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 GetBudgetComplete(base::Closure run_loop_closure,
-                         bool success,
-                         std::unique_ptr<budget_service::Budget> budget) {
-    budget_ = std::move(budget);
+  void AddBudgetComplete(base::Closure run_loop_closure, bool success) {
     success_ = success;
     run_loop_closure.Run();
   }
 
-  // Excercise the very basic get method.
-  std::unique_ptr<budget_service::Budget> GetBudget() {
+  // Add budget to the origin.
+  bool AddBudget(const GURL& origin, double amount) {
     base::RunLoop run_loop;
-    db_.GetValue(GURL(kTestOrigin),
-                 base::Bind(&BudgetDatabaseTest::GetBudgetComplete,
-                            base::Unretained(this), run_loop.QuitClosure()));
+    db_.AddBudget(origin, amount,
+                  base::Bind(&BudgetDatabaseTest::AddBudgetComplete,
+                             base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
-    return std::move(budget_);
+    return success_;
   }
 
   void GetBudgetDetailsComplete(
@@ -86,6 +62,7 @@ class BudgetDatabaseTest : public ::testing::Test {
     run_loop_closure.Run();
   }
 
+  // Get the full set of budget expectations for the origin.
   void GetBudgetDetails() {
     base::RunLoop run_loop;
     db_.GetBudgetDetails(
@@ -100,6 +77,13 @@ class BudgetDatabaseTest : public ::testing::Test {
     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;
+  }
+
  protected:
   bool success_;
 
@@ -113,18 +97,19 @@ class BudgetDatabaseTest : public ::testing::Test {
 };
 
 TEST_F(BudgetDatabaseTest, ReadAndWriteTest) {
-  ASSERT_TRUE(SetBudgetWithDefaultValues());
-  std::unique_ptr<budget_service::Budget> b = GetBudget();
-
-  ASSERT_TRUE(success_);
-  EXPECT_EQ(kDefaultDebt, b->debt());
-  EXPECT_EQ(kDefaultBudget1, b->budget(0).amount());
-  EXPECT_EQ(kDefaultBudget2, b->budget(1).amount());
-  EXPECT_EQ(kDefaultExpiration, b->budget(0).expiration());
-}
-
-TEST_F(BudgetDatabaseTest, BudgetDetailsTest) {
-  ASSERT_TRUE(SetBudgetWithDefaultValues());
+  const GURL origin(kTestOrigin);
+  base::SimpleTestClock* clock = SetClockForTesting();
+  base::TimeDelta expiration(
+      base::TimeDelta::FromHours(kDefaultExpirationInHours));
+  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 expectation and validate it.
@@ -138,18 +123,17 @@ TEST_F(BudgetDatabaseTest, BudgetDetailsTest) {
   // First value should be [total_budget, now]
   EXPECT_EQ(kDefaultBudget1 + kDefaultBudget2, iter->first);
   // TODO(harkness): This will be "now" in the final version. For now, it's
-  // just 0.0.
+  // just 0.
   EXPECT_EQ(0, iter->second);
 
   // The next value should be the budget after the first chunk expires.
   iter++;
   EXPECT_EQ(kDefaultBudget2, iter->first);
-  EXPECT_EQ(kDefaultExpiration, iter->second);
+  EXPECT_EQ(expiration_time.ToInternalValue(), iter->second);
 
   // 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->first);
-  EXPECT_EQ(kDefaultExpiration + 1, iter->second);
+  EXPECT_EQ(expiration_time.ToInternalValue(), iter->second);
 }
-
-}  // namespace