Refactor budget computation to be more tuneable.

chrome/browser/budget_service/budget_manager_browsertest.cc

Index: chrome/browser/budget_service/budget_manager_browsertest.cc
diff --git a/chrome/browser/budget_service/budget_manager_browsertest.cc b/chrome/browser/budget_service/budget_manager_browsertest.cc
index aef3cd4c4fd680ee19be200cca65ab3b373c4f34..19b8abc3d1e17454b055c414ad6791be6a91742c 100644
--- a/chrome/browser/budget_service/budget_manager_browsertest.cc
+++ b/chrome/browser/budget_service/budget_manager_browsertest.cc
@@ -94,20 +94,22 @@ class BudgetManagerBrowserTest : public InProcessBrowserTest {
 IN_PROC_BROWSER_TEST_F(BudgetManagerBrowserTest, BudgetInDocument) {
   std::string script_result;
 
-  // The page will have been loaded once, which gives a budget of 3.
+  LoadTestPage();  // Reload to build site engagement.
+
+  // The page will have been loaded twice, which gives a budget of 2.
   ASSERT_TRUE(RunScript("documentGetBudget()", &script_result));
-  ASSERT_EQ("ok - budget returned value of 3", script_result);
+  EXPECT_EQ("ok - budget returned value of 2", script_result);
 
   ASSERT_TRUE(RunScript("documentReserveBudget()", &script_result));
-  ASSERT_EQ("ok - reserved budget", script_result);
+  EXPECT_EQ("ok - reserved budget", script_result);
 
-  // After reserving budget, the new budget should be at 1.
+  // After reserving budget, the new budget should be at 0.
   ASSERT_TRUE(RunScript("documentGetBudget()", &script_result));
-  ASSERT_EQ("ok - budget returned value of 1", script_result);
+  EXPECT_EQ("ok - budget returned value of 0", script_result);
 
   // A second reserve should fail because there is not enough budget.
   ASSERT_TRUE(RunScript("documentReserveBudget()", &script_result));
-  ASSERT_EQ("failed - not able to reserve budget", script_result);
+  EXPECT_EQ("failed - not able to reserve budget", script_result);
 
   // Consume should succeed because there is an existing reservation.
   ConsumeReservation();
@@ -125,29 +127,32 @@ IN_PROC_BROWSER_TEST_F(BudgetManagerBrowserTest, BudgetInWorker) {
   ASSERT_EQ("ok - service worker registered", script_result);
 
   LoadTestPage();  // Reload to become controlled.
+  LoadTestPage();  // Reload to build site engagement.
+  LoadTestPage();  // Reload to build site engagement.
+  LoadTestPage();  // Reload to build site engagement.
+  LoadTestPage();  // Reload to build site engagement.
 
   ASSERT_TRUE(RunScript("isControlled()", &script_result));
   ASSERT_EQ("true - is controlled", script_result);
 
-  // The page will have been loaded twice and a service worker was registered,
-  // which gives a budget of 4.5.
+  // The page will have been loaded five times, giving a budget of 5.
   ASSERT_TRUE(RunScript("workerGetBudget()", &script_result));
-  ASSERT_EQ("ok - budget returned value of 4.5", script_result);
+  EXPECT_EQ("ok - budget returned value of 5", script_result);
 
-  // With a budget of 4.5, two reservations should succeed.
+  // With a budget of 5, two reservations should succeed.
   ASSERT_TRUE(RunScript("workerReserveBudget()", &script_result));
-  ASSERT_EQ("ok - reserved budget", script_result);
+  EXPECT_EQ("ok - reserved budget", script_result);
 
   ASSERT_TRUE(RunScript("workerReserveBudget()", &script_result));
-  ASSERT_EQ("ok - reserved budget", script_result);
+  EXPECT_EQ("ok - reserved budget", script_result);
 
-  // After reserving budget, the new budget should be at 0.5.
+  // After reserving budget, the new budget should be at 1.
   ASSERT_TRUE(RunScript("workerGetBudget()", &script_result));
-  ASSERT_EQ("ok - budget returned value of 0.5", script_result);
+  EXPECT_EQ("ok - budget returned value of 1", script_result);
 
   // A second reserve should fail because there is not enough budget.
   ASSERT_TRUE(RunScript("workerReserveBudget()", &script_result));
-  ASSERT_EQ("failed - not able to reserve budget", script_result);
+  EXPECT_EQ("failed - not able to reserve budget", script_result);
 
   // Two consumes should succeed because there are existing reservations.
   ConsumeReservation();