Replace the 1 in 10 grace period with an accumulating budget based on SES.

chrome/browser/push_messaging/background_budget_service_unittest.cc

Index: chrome/browser/push_messaging/background_budget_service_unittest.cc
diff --git a/chrome/browser/push_messaging/background_budget_service_unittest.cc b/chrome/browser/push_messaging/background_budget_service_unittest.cc
index 255dede4589ebc121c20e07507357c9b79027588..501602881cd4def8dce4a089004b1cf294094aac 100644
--- a/chrome/browser/push_messaging/background_budget_service_unittest.cc
+++ b/chrome/browser/push_messaging/background_budget_service_unittest.cc
@@ -5,16 +5,27 @@
 #include <stdint.h>
 #include <string>
 
+#include "base/memory/ptr_util.h"
+#include "base/test/simple_test_clock.h"
+#include "chrome/browser/engagement/site_engagement_service.h"
 #include "chrome/browser/push_messaging/background_budget_service.h"
 #include "chrome/browser/push_messaging/background_budget_service_factory.h"
+#include "chrome/common/pref_names.h"
 #include "chrome/test/base/testing_profile.h"
+#include "components/prefs/pref_service.h"
+#include "components/prefs/scoped_user_pref_update.h"
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace {
 
 const char kTestOrigin[] = "https://example.com";
-const char kTestData[] = "1111101111";
+const double kTestBudget = 10.0;
+const double kTestSES = 24.0;
+const double kLowSES = 1.0;
+const double kMaxSES = 100.0;
+// Mirrors definition in BackgroundBudgetService, this is 10 days of seconds.
+const double kSecondsToAccumulate = 864000.0;
 
 }  // namespace
 
@@ -23,29 +34,181 @@ class BackgroundBudgetServiceTest : public testing::Test {
   BackgroundBudgetServiceTest() {}
   ~BackgroundBudgetServiceTest() override {}
 
-  BackgroundBudgetService* service() {
+  BackgroundBudgetService* GetService() {
     return BackgroundBudgetServiceFactory::GetForProfile(&profile_);
   }
 
+  void SetSiteEngagementScore(const GURL& url, double score) {
+    SiteEngagementService* service = SiteEngagementService::Get(&profile_);
+    service->ResetScoreForURL(url, score);
+  }
+
+  Profile* profile() { return &profile_; }
+
+  base::SimpleTestClock* SetClockForTesting() {
+    base::SimpleTestClock* clock = new base::SimpleTestClock();
+    BackgroundBudgetServiceFactory::GetForProfile(&profile_)
+        ->SetClockForTesting(base::WrapUnique(clock));
+    return clock;
+  }
+
  private:
   content::TestBrowserThreadBundle thread_bundle_;
   TestingProfile profile_;
 };
 
-TEST_F(BackgroundBudgetServiceTest, GetBudgetFailure) {
+TEST_F(BackgroundBudgetServiceTest, GetBudgetNoBudgetOrSES) {
+  const GURL origin(kTestOrigin);
+
+  double budget = GetService()->GetBudget(origin);
+
+  EXPECT_DOUBLE_EQ(budget, 0.0);
+}
+
+TEST_F(BackgroundBudgetServiceTest, GetBudgetNoBudgetSESExists) {
+  const GURL origin(kTestOrigin);
+
+  // Set a starting SES for the url but no stored budget info.
+  SetSiteEngagementScore(origin, kTestSES);
+
+  double budget = GetService()->GetBudget(origin);
+
+  EXPECT_DOUBLE_EQ(budget, kTestSES);
+}
+
+TEST_F(BackgroundBudgetServiceTest, GetBudgetNoElapsedTime) {
+  const GURL origin(kTestOrigin);
+
+  std::unique_ptr<BackgroundBudgetService> service(
+      new BackgroundBudgetService(profile()));
+
+  service->StoreBudget(origin, kTestBudget);
+
+  double budget = service->GetBudget(origin);
+
+  EXPECT_NEAR(budget, kTestBudget, kTestBudget);
+}
+
+TEST_F(BackgroundBudgetServiceTest, GetBudgetElapsedTime) {
+  // Manually construct a BackgroundBudgetServie with a clock that the test
+  // can control so that we can fast forward in time.
+  BackgroundBudgetService* service = GetService();
+  base::SimpleTestClock* clock = SetClockForTesting();
+  base::Time starting_time = clock->Now();
+
+  // Set initial SES and budget values.
   const GURL origin(kTestOrigin);
+  SetSiteEngagementScore(origin, kTestSES);
+  service->StoreBudget(origin, kTestBudget);
+
+  double budget = service->GetBudget(origin);
+  EXPECT_DOUBLE_EQ(budget, kTestBudget);
+
+  // Query for the budget after 1 second has passed.
+  clock->SetNow(starting_time + base::TimeDelta::FromSeconds(1));
+  budget = service->GetBudget(origin);
+  EXPECT_NEAR(budget, kTestBudget, kTestSES * 1.0 / kSecondsToAccumulate);
+  EXPECT_GT(budget, kTestBudget);
+
+  // Query for the budget after 1 hour has passed.
+  clock->SetNow(starting_time + base::TimeDelta::FromHours(1));
+  budget = service->GetBudget(origin);
+  EXPECT_NEAR(budget, kTestBudget, kTestSES * 3600.0 / kSecondsToAccumulate);
+  EXPECT_GT(budget, kTestBudget);
+
+  // Query for the budget after 5 days have passed. The budget should be
+  // increasing, but not up the SES score.
+  clock->SetNow(starting_time + base::TimeDelta::FromDays(5));
+  budget = service->GetBudget(origin);
+  EXPECT_GT(budget, kTestBudget);
+  EXPECT_LT(budget, kTestSES);
+  double moderate_ses_budget = budget;
+
+  // Query for the budget after 10 days have passed. By this point, the budget
+  // should converge to the SES score.
+  clock->SetNow(starting_time + base::TimeDelta::FromDays(10));
+  budget = service->GetBudget(origin);
+  EXPECT_DOUBLE_EQ(budget, kTestSES);
+
+  // Now, change the SES score to the maximum amount and reinitialize budget.
+  SetSiteEngagementScore(origin, kMaxSES);
+  service->StoreBudget(origin, kTestBudget);
+  starting_time = clock->Now();
+
+  // Query for the budget after 1 second has passed.
+  clock->SetNow(starting_time + base::TimeDelta::FromSeconds(1));
+  budget = service->GetBudget(origin);
+  EXPECT_NEAR(budget, kTestBudget, kMaxSES * 1.0 / kSecondsToAccumulate);
+
+  // Query for the budget after 5 days have passed. Again, the budget should be
+  // approaching the SES, but not have reached it.
+  clock->SetNow(starting_time + base::TimeDelta::FromDays(5));
+  budget = service->GetBudget(origin);
+  EXPECT_GT(budget, kTestBudget);
+  EXPECT_LT(budget, kMaxSES);
 
-  std::string budget_string = service()->GetBudget(origin);
+  // The budget after 5 days with max SES should be greater than the budget
+  // after 5 days with moderate SES.
+  EXPECT_GT(budget, moderate_ses_budget);
 
-  EXPECT_EQ(std::string(), budget_string);
+  // Now, change the SES score to a low amount and reinitialize budget.
+  SetSiteEngagementScore(origin, kLowSES);
+  service->StoreBudget(origin, kTestBudget);
+  starting_time = clock->Now();
+
+  // Query for the budget after 5 days have passed. Again, the budget should be
+  // approaching the SES, this time decreasing, but not have reached it.
+  clock->SetNow(starting_time + base::TimeDelta::FromDays(5));
+  budget = service->GetBudget(origin);
+  EXPECT_LT(budget, kTestBudget);
+  EXPECT_GT(budget, kLowSES);
+}
+
+TEST_F(BackgroundBudgetServiceTest, GetBudgetConsumedOverTime) {
+  // Manually construct a BackgroundBudgetService with a clock that the test
+  // can control so that we can fast forward in time.
+  BackgroundBudgetService* service = GetService();
+  base::SimpleTestClock* clock = SetClockForTesting();
+
+  // Set initial SES and budget values.
+  const GURL origin(kTestOrigin);
+  SetSiteEngagementScore(origin, kTestSES);
+  service->StoreBudget(origin, kTestBudget);
+  double budget = 0.0;
+
+  // Measure over 200 hours. In each hour a message is received, and for 1 in
+  // 10, budget is consumed.
+  for (int i = 0; i < 200; i++) {
+    // Query for the budget after 1 hour has passed.
+    clock->Advance(base::TimeDelta::FromHours(1));
+    budget = service->GetBudget(origin);
+
+    if (i % 10 == 0) {
+      service->StoreBudget(origin, budget - 1.0);
+    }
+  }
+
+  // With a SES of 24.0, the origin will get a budget of 2.4 per day, but the
+  // old budget will also decay. At the end, we expect the budget to be lower
+  // than the starting budget.
+  EXPECT_GT(budget, 0.0);
+  EXPECT_LT(budget, kTestBudget);
 }
 
-TEST_F(BackgroundBudgetServiceTest, GetBudgetSuccess) {
+TEST_F(BackgroundBudgetServiceTest, GetBudgetInvalidBudget) {
   const GURL origin(kTestOrigin);
 
-  service()->StoreBudget(origin, kTestData);
+  // Set a starting SES for the url.
+  SetSiteEngagementScore(origin, kTestSES);
+
+  // Set a badly formatted budget in the user preferences.
+  DictionaryPrefUpdate update(profile()->GetPrefs(),
+                              prefs::kBackgroundBudgetMap);
+  base::DictionaryValue* update_map = update.Get();
+  update_map->SetStringWithoutPathExpansion(origin.spec(), "20#2.0");
 
-  std::string budget_string = service()->GetBudget(origin);
+  // Get the budget, expect that it will return SES.
+  double budget = GetService()->GetBudget(origin);
 
-  EXPECT_EQ(kTestData, budget_string);
+  EXPECT_DOUBLE_EQ(budget, kTestSES);
 }