[iOS] Update BadgeMediator to take a Browser.

This CL updates BadgeMediator to be constructed with a Browser in
preparation for a follow-up CL where the mediator will observe the
OverlayPresenter for the Browser in order to update the
InfobarBadgeTabHelper for banner presentations and dismissals.

This CL also cleans up and updates the mediator:
- Adds readwrite property for the active WebState, allowing easy
  access to the WebState as well as facilitating idempotent updates
  via |-updateForNewWebState:oldWebState:|.
- Adds a readonly property for the active WebState's
  InfobarBadgeTabHelper.
- Since a Browser's BrowserState is constant over the Browser's
  lifetime, we don't need to check whether the newly-activated
  WebState is incognito.  Instead, the incognito badge item is just
  created once upon initialization.
- Updates order of protocol implementations to match declaration
  order.
- Updates test fixture to use params so that both the incognito and
  normal browsing behavior can be tested.

Bug: 1030357
Change-Id: If7a5b4761ea4fd8e558120d1cfc14eb806c75290
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1986860
Commit-Queue: Kurt Horimoto <kkhorimoto@chromium.org>
Reviewed-by: Eugene But <eugenebut@chromium.org>
Reviewed-by: Chris Lu <thegreenfrog@chromium.org>
Cr-Commit-Position: refs/heads/master@{#728740}

ios/chrome/browser/ui/badges/BUILD.gn

@@ -49,8 +49,10 @@ source_set("badges") {
     "resources:wrench_badge",
     "//base:base",
     "//ios/chrome/app/strings:ios_strings_grit",
+    "//ios/chrome/browser/browser_state",
     "//ios/chrome/browser/infobars:badge",
     "//ios/chrome/browser/infobars:public",
+    "//ios/chrome/browser/main:public",
     "//ios/chrome/browser/ui/commands",
     "//ios/chrome/browser/ui/coordinators:chrome_coordinators",
     "//ios/chrome/browser/ui/elements",
@@ -112,6 +114,7 @@ source_set("unit_tests") {
     "//ios/chrome/browser/infobars:badge",
     "//ios/chrome/browser/infobars:badge_public",
     "//ios/chrome/browser/infobars/test",
+    "//ios/chrome/browser/main:test_support",
     "//ios/chrome/browser/ui/infobars:feature_flags",
     "//ios/chrome/browser/ui/infobars:test_support",
     "//ios/chrome/browser/web_state_list",

ios/chrome/browser/ui/badges/badge_mediator.h

@@ -11,16 +11,14 @@
 
 @protocol BadgeConsumer;
 @protocol BadgeItem;
+class Browser;
 @protocol BrowserCoordinatorCommands;
 @protocol InfobarCommands;
-class WebStateList;
 
 // A mediator object that updates the consumer when the state of badges changes.
 @interface BadgeMediator : NSObject <BadgeDelegate>
 
-- (instancetype)initWithConsumer:(id<BadgeConsumer>)consumer
-                    webStateList:(WebStateList*)webStateList
-    NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithBrowser:(Browser*)browser NS_DESIGNATED_INITIALIZER;
 - (instancetype)init NS_UNAVAILABLE;
 
 // Stops observing all objects.
@@ -30,6 +28,10 @@ class WebStateList;
 @property(nonatomic, weak) id<InfobarCommands, BrowserCoordinatorCommands>
     dispatcher;
 
+// The consumer being set up by this mediator.  Setting to a new value updates
+// the new consumer.
+@property(nonatomic, weak) id<BadgeConsumer> consumer;
+
 @end
 
 #endif  // IOS_CHROME_BROWSER_UI_BADGES_BADGE_MEDIATOR_H_

ios/chrome/browser/ui/badges/badge_mediator.mm

@@ -6,10 +6,12 @@
 
 #include "base/mac/foundation_util.h"
 #include "base/metrics/user_metrics.h"
+#include "ios/chrome/browser/browser_state/chrome_browser_state.h"
 #include "ios/chrome/browser/infobars/infobar_badge_tab_helper.h"
 #include "ios/chrome/browser/infobars/infobar_badge_tab_helper_delegate.h"
 #include "ios/chrome/browser/infobars/infobar_metrics_recorder.h"
 #import "ios/chrome/browser/infobars/infobar_type.h"
+#include "ios/chrome/browser/main/browser.h"
 #import "ios/chrome/browser/ui/badges/badge_button.h"
 #import "ios/chrome/browser/ui/badges/badge_consumer.h"
 #import "ios/chrome/browser/ui/badges/badge_item.h"
@@ -20,7 +22,6 @@
 #import "ios/chrome/browser/ui/list_model/list_model.h"
 #import "ios/chrome/browser/web_state_list/web_state_list.h"
 #import "ios/chrome/browser/web_state_list/web_state_list_observer_bridge.h"
-#include "ios/web/public/browser_state.h"
 
 #if !defined(__has_feature) || !__has_feature(objc_arc)
 #error "This file requires ARC support."
@@ -32,38 +33,45 @@ const int kNumberOfFullScrenBadges = 1;
 // The minimum number of non-Fullscreen badges to display the overflow popup
 // menu.
 const int kMinimumNonFullScreenBadgesForOverflow = 2;
-}
+}  // namespace
 
 @interface BadgeMediator () <InfobarBadgeTabHelperDelegate,
                              WebStateListObserving> {
-  std::unique_ptr<WebStateListObserverBridge> _webStateListObserver;
+  std::unique_ptr<WebStateListObserver> _webStateListObserver;
 }
 
 // The WebStateList that this mediator listens for any changes on the active web
 // state.
-@property(nonatomic, assign) WebStateList* webStateList;
+@property(nonatomic, readonly) WebStateList* webStateList;
+
+// The WebStateList's active WebState.
+@property(nonatomic, assign) web::WebState* webState;
+
+// The active WebState's badge tab helper.
+@property(nonatomic, readonly) InfobarBadgeTabHelper* badgeTabHelper;
+
+// The incognito badge, or nil if the Browser is not off-the-record.
+@property(nonatomic, readonly) id<BadgeItem> offTheRecordBadge;
 
 // Array of all available badges.
 @property(nonatomic, strong) NSMutableArray<id<BadgeItem>>* badges;
 
-// The consumer of the mediator.
-@property(nonatomic, weak) id<BadgeConsumer> consumer;
-
 @end
 
 @implementation BadgeMediator
-@synthesize webStateList = _webStateList;
 
-- (instancetype)initWithConsumer:(id<BadgeConsumer>)consumer
-                    webStateList:(WebStateList*)webStateList {
+- (instancetype)initWithBrowser:(Browser*)browser {
   self = [super init];
   if (self) {
-    _consumer = consumer;
-    _webStateList = webStateList;
-    web::WebState* activeWebState = webStateList->GetActiveWebState();
-    if (activeWebState) {
-      [self updateNewWebState:activeWebState withWebStateList:webStateList];
+    DCHECK(browser);
+    // Create the incognito badge if |browser| is off-the-record.
+    if (browser->GetBrowserState()->IsOffTheRecord()) {
+      _offTheRecordBadge = [[BadgeStaticItem alloc]
+          initWithBadgeType:BadgeType::kBadgeTypeIncognito];
     }
+    // Set up the WebStateList and its observer.
+    _webStateList = browser->GetWebStateList();
+    _webState = _webStateList->GetActiveWebState();
     _webStateListObserver = std::make_unique<WebStateListObserverBridge>(self);
     _webStateList->AddObserver(_webStateListObserver.get());
   }
@@ -71,45 +79,82 @@ const int kMinimumNonFullScreenBadgesForOverflow = 2;
 }
 
 - (void)dealloc {
-  [self disconnect];
+  // |-disconnect| must be called before deallocation.
+  DCHECK(!_webStateList);
 }
 
 - (void)disconnect {
   if (_webStateList) {
+    self.webState = nullptr;
     _webStateList->RemoveObserver(_webStateListObserver.get());
-    _webStateListObserver.reset();
+    _webStateListObserver = nullptr;
     _webStateList = nullptr;
   }
 }
 
-#pragma mark - InfobarBadgeTabHelperDelegate
+#pragma mark - Accessors
 
-- (void)addInfobarBadge:(id<BadgeItem>)badgeItem {
-  if (!self.badges) {
-    self.badges = [NSMutableArray array];
-  }
-  [self.badges addObject:badgeItem];
-  [self updateBadgesShown];
+- (void)setConsumer:(id<BadgeConsumer>)consumer {
+  if (_consumer == consumer)
+    return;
+  _consumer = consumer;
+  [self updateConsumer];
 }
 
-- (void)removeInfobarBadge:(id<BadgeItem>)badgeItem {
-  for (id<BadgeItem> item in self.badges) {
-    if (item.badgeType == badgeItem.badgeType) {
-      [self.badges removeObject:item];
-      [self updateBadgesShown];
+- (void)setWebState:(web::WebState*)webState {
+  if (_webState == webState)
     return;
-    }
+  if (_webState)
+    InfobarBadgeTabHelper::FromWebState(_webState)->SetDelegate(nil);
+  _webState = webState;
+  if (_webState)
+    InfobarBadgeTabHelper::FromWebState(_webState)->SetDelegate(self);
+  [self updateBadgesForActiveWebState];
+  [self updateConsumer];
+}
+
+- (InfobarBadgeTabHelper*)badgeTabHelper {
+  return self.webState ? InfobarBadgeTabHelper::FromWebState(self.webState)
+                       : nullptr;
+}
+
+#pragma mark - Accessor helpers
+
+- (void)updateBadgesForActiveWebState {
+  if (self.webState) {
+    self.badges = [self.badgeTabHelper->GetInfobarBadgeItems() mutableCopy];
+    if (self.offTheRecordBadge)
+      [self.badges addObject:self.offTheRecordBadge];
+  } else {
+    self.badges = [NSMutableArray<id<BadgeItem>> array];
   }
 }
 
-- (void)updateInfobarBadge:(id<BadgeItem>)badgeItem {
-  for (id<BadgeItem> item in self.badges) {
-    if (item.badgeType == badgeItem.badgeType) {
-      item.badgeState = badgeItem.badgeState;
-      [self updateBadgesShown];
+// Updates the consumer for the current active WebState.
+- (void)updateConsumer {
+  if (!self.consumer)
     return;
+
+  // Update the badges array if necessary.
+  if (!self.badges)
+    [self updateBadgesForActiveWebState];
+
+  // Show the overflow badge if there are multiple BadgeItems.  Otherwise, use
+  // the last badge.
+  NSUInteger fullscreenBadgeCount = self.offTheRecordBadge ? 1U : 0U;
+  BOOL shouldDisplayOverflowBadge =
+      self.badges.count - fullscreenBadgeCount > 1;
+  id<BadgeItem> displayedBadge = nil;
+  if (shouldDisplayOverflowBadge) {
+    displayedBadge = [[BadgeTappableItem alloc]
+        initWithBadgeType:BadgeType::kBadgeTypeOverflow];
+  } else {
+    id<BadgeItem> firstBadge = [self.badges firstObject];
+    displayedBadge = firstBadge.fullScreen ? nil : firstBadge;
   }
-  }
+  // Update the consumer with the new badge items.
+  [self.consumer setupWithDisplayedBadge:displayedBadge
+                         fullScreenBadge:self.offTheRecordBadge];
 }
 
 #pragma mark - BadgeDelegate
@@ -166,15 +211,42 @@ const int kMinimumNonFullScreenBadgesForOverflow = 2;
   // TODO(crbug.com/976901): Add metric for this action.
 }
 
+#pragma mark - InfobarBadgeTabHelperDelegate
+
+- (void)addInfobarBadge:(id<BadgeItem>)badgeItem {
+  [self.badges addObject:badgeItem];
+  [self updateBadgesShown];
+}
+
+- (void)removeInfobarBadge:(id<BadgeItem>)badgeItem {
+  for (id<BadgeItem> item in self.badges) {
+    if (item.badgeType == badgeItem.badgeType) {
+      [self.badges removeObject:item];
+      [self updateBadgesShown];
+      return;
+    }
+  }
+}
+
+- (void)updateInfobarBadge:(id<BadgeItem>)badgeItem {
+  for (id<BadgeItem> item in self.badges) {
+    if (item.badgeType == badgeItem.badgeType) {
+      item.badgeState = badgeItem.badgeState;
+      [self updateBadgesShown];
+      return;
+    }
+  }
+}
+
 #pragma mark - WebStateListObserver
 
 - (void)webStateList:(WebStateList*)webStateList
     didReplaceWebState:(web::WebState*)oldWebState
           withWebState:(web::WebState*)newWebState
                atIndex:(int)atIndex {
-  if (newWebState && newWebState == webStateList->GetActiveWebState()) {
-    [self updateNewWebState:newWebState withWebStateList:webStateList];
-  }
+  DCHECK_EQ(self.webStateList, webStateList);
+  if (atIndex == webStateList->active_index())
+    self.webState = newWebState;
 }
 
 - (void)webStateList:(WebStateList*)webStateList
@@ -182,11 +254,8 @@ const int kMinimumNonFullScreenBadgesForOverflow = 2;
                 oldWebState:(web::WebState*)oldWebState
                     atIndex:(int)atIndex
                      reason:(int)reason {
-  // Only attempt to retrieve badges if there is a new current web state, since
-  // |newWebState| can be null.
-  if (newWebState) {
-    [self updateNewWebState:newWebState withWebStateList:webStateList];
-  }
+  DCHECK_EQ(self.webStateList, webStateList);
+  self.webState = newWebState;
 }
 
 #pragma mark - Helpers
@@ -279,37 +348,4 @@ const int kMinimumNonFullScreenBadgesForOverflow = 2;
   [self updateConsumerReadStatus];
 }
 
-- (void)updateNewWebState:(web::WebState*)newWebState
-         withWebStateList:(WebStateList*)webStateList {
-  DCHECK_EQ(_webStateList, webStateList);
-  InfobarBadgeTabHelper* infobarBadgeTabHelper =
-      InfobarBadgeTabHelper::FromWebState(newWebState);
-  DCHECK(infobarBadgeTabHelper);
-  infobarBadgeTabHelper->SetDelegate(self);
-  // Whenever the WebState changes ask the corresponding
-  // InfobarBadgeTabHelper for all the badges for that WebState.
-  self.badges = [infobarBadgeTabHelper->GetInfobarBadgeItems() mutableCopy];
-
-  id<BadgeItem> displayedBadge;
-  if ([self.badges count] > 1) {
-    // Show the overflow menu badge when there are multiple badges.
-    displayedBadge = [[BadgeTappableItem alloc]
-        initWithBadgeType:BadgeType::kBadgeTypeOverflow];
-  } else if ([self.badges count] == 1) {
-    displayedBadge = [self.badges lastObject];
-  }
-  id<BadgeItem> fullScreenBadge;
-  if (newWebState->GetBrowserState()->IsOffTheRecord()) {
-    BadgeStaticItem* incognitoItem = [[BadgeStaticItem alloc]
-        initWithBadgeType:BadgeType::kBadgeTypeIncognito];
-    fullScreenBadge = incognitoItem;
-    // Keep track of presence of an incognito badge so that the mediator knows
-    // whether or not there is a fullscreen badge when calling
-    // updateDisplayedBadge:fullScreenBadge:.
-    [self.badges addObject:incognitoItem];
-  }
-  [self.consumer setupWithDisplayedBadge:displayedBadge
-                         fullScreenBadge:fullScreenBadge];
-}
-
 @end

ios/chrome/browser/ui/badges/badge_mediator_unittest.mm

@@ -13,6 +13,7 @@
 #include "ios/chrome/browser/infobars/infobar_ios.h"
 #include "ios/chrome/browser/infobars/infobar_manager_impl.h"
 #import "ios/chrome/browser/infobars/test/fake_infobar_ios.h"
+#import "ios/chrome/browser/main/test_browser.h"
 #import "ios/chrome/browser/ui/badges/badge_consumer.h"
 #import "ios/chrome/browser/ui/badges/badge_item.h"
 #import "ios/chrome/browser/ui/infobars/infobar_feature.h"
@@ -24,7 +25,6 @@
 #import "ios/web/public/test/fakes/test_web_state.h"
 #import "ios/web/public/web_state_user_data.h"
 #include "testing/gtest/include/gtest/gtest.h"
-#include "testing/platform_test.h"
 
 #if !defined(__has_feature) || !__has_feature(objc_arc)
 #error "This file requires ARC support."
@@ -34,24 +34,29 @@ namespace {
 // The two infobar types used in tests.  Both support badges.
 InfobarType kFirstInfobarType = InfobarType::kInfobarTypePasswordSave;
 InfobarType kSecondInfobarType = InfobarType::kInfobarTypePasswordUpdate;
+// Parameters used for BadgeMediator test fixtures.
+enum class TestParam {
+  kNormal,
+  kOffTheRecord,
+};
 }  // namespace
 
 // Fake of BadgeConsumer.
 @interface FakeBadgeConsumer : NSObject <BadgeConsumer>
 @property(nonatomic, strong) id<BadgeItem> displayedBadge;
-@property(nonatomic, assign) BOOL hasIncognitoBadge;
+@property(nonatomic, assign) BOOL hasFullscreenBadge;
 @property(nonatomic, assign) BOOL hasUnreadBadge;
 @end
 
 @implementation FakeBadgeConsumer
 - (void)setupWithDisplayedBadge:(id<BadgeItem>)displayedBadgeItem
                 fullScreenBadge:(id<BadgeItem>)fullscreenBadgeItem {
-  self.hasIncognitoBadge = fullscreenBadgeItem != nil;
+  self.hasFullscreenBadge = fullscreenBadgeItem != nil;
   self.displayedBadge = displayedBadgeItem;
 }
 - (void)updateDisplayedBadge:(id<BadgeItem>)displayedBadgeItem
              fullScreenBadge:(id<BadgeItem>)fullscreenBadgeItem {
-  self.hasIncognitoBadge = fullscreenBadgeItem != nil;
+  self.hasFullscreenBadge = fullscreenBadgeItem != nil;
   self.displayedBadge = displayedBadgeItem;
 }
 - (void)markDisplayedBadgeAsRead:(BOOL)read {
@@ -59,39 +64,31 @@ InfobarType kSecondInfobarType = InfobarType::kInfobarTypePasswordUpdate;
 }
 @end
 
-class BadgeMediatorTest : public PlatformTest {
+class BadgeMediatorTest : public testing::TestWithParam<TestParam> {
  protected:
   BadgeMediatorTest()
-      : browser_state_(TestChromeBrowserState::Builder().Build()),
-        web_state_list_(
-            std::make_unique<WebStateList>(&web_state_list_delegate_)) {
+      : badge_consumer_([[FakeBadgeConsumer alloc] init]),
+        browser_state_(TestChromeBrowserState::Builder().Build()),
+        web_state_list_(&web_state_list_delegate_) {
     feature_list_.InitAndEnableFeature(kInfobarUIReboot);
-    badge_consumer_ = [[FakeBadgeConsumer alloc] init];
-    badge_mediator_ =
-        [[BadgeMediator alloc] initWithConsumer:badge_consumer_
-                                   webStateList:web_state_list_.get()];
+    badge_mediator_ = [[BadgeMediator alloc] initWithBrowser:browser()];
+    badge_mediator_.consumer = badge_consumer_;
   }
 
   ~BadgeMediatorTest() override { [badge_mediator_ disconnect]; }
 
-  // Adds a new WebState to the WebStateList and activates it.
-  void AddAndActivateWebState(int index, bool incognito) {
+  // Inserts a new WebState to the WebStateList at |index| and activates it.
+  void InsertActivatedWebState(int index) {
     std::unique_ptr<web::TestWebState> web_state =
         std::make_unique<web::TestWebState>();
     web_state->SetNavigationManager(
         std::make_unique<web::TestNavigationManager>());
+    web_state->SetBrowserState(browser_state());
     InfoBarManagerImpl::CreateForWebState(web_state.get());
     InfobarBadgeTabHelper::CreateForWebState(web_state.get());
-    if (incognito) {
-      web_state->SetBrowserState(
-          browser_state_->GetOffTheRecordChromeBrowserState());
-    } else {
-      web_state->SetBrowserState(browser_state_.get());
-    }
-    web_state_list_->InsertWebState(index, std::move(web_state),
-                                    WebStateList::INSERT_NO_FLAGS,
+    web_state_list_.InsertWebState(index, std::move(web_state),
+                                   WebStateList::INSERT_ACTIVATE,
                                    WebStateOpener());
-    web_state_list_->ActivateWebStateAt(index);
   }
 
   // Adds an Infobar of |type| to the InfoBarManager and returns the infobar.
@@ -108,8 +105,26 @@ class BadgeMediatorTest : public PlatformTest {
     infobar_manager()->RemoveInfoBar(infobar);
   }
 
+  // Returns whether the test fixture is for an incognito BrowserState.
+  bool is_off_the_record() const {
+    return GetParam() == TestParam::kOffTheRecord;
+  }
+  // Returns the BrowserState to use for the test fixture.
+  ios::ChromeBrowserState* browser_state() {
+    return is_off_the_record()
+               ? browser_state_->GetOffTheRecordChromeBrowserState()
+               : browser_state_.get();
+  }
+  // Returns the Browser to use for the test fixture.
+  Browser* browser() {
+    if (!browser_) {
+      browser_ =
+          std::make_unique<TestBrowser>(browser_state(), &web_state_list_);
+    }
+    return browser_.get();
+  }
   // Returns the active WebState.
-  web::WebState* web_state() { return web_state_list_->GetActiveWebState(); }
+  web::WebState* web_state() { return web_state_list_.GetActiveWebState(); }
   // Returns the active WebState's InfoBarManagerImpl.
   InfoBarManagerImpl* infobar_manager() {
     return InfoBarManagerImpl::FromWebState(web_state());
@@ -123,23 +138,24 @@ class BadgeMediatorTest : public PlatformTest {
   base::test::TaskEnvironment environment_;
   FakeBadgeConsumer* badge_consumer_;
   std::unique_ptr<ios::ChromeBrowserState> browser_state_;
-  BadgeMediator* badge_mediator_;
-  std::unique_ptr<WebStateList> web_state_list_;
   FakeWebStateListDelegate web_state_list_delegate_;
+  WebStateList web_state_list_;
+  std::unique_ptr<Browser> browser_;
+  BadgeMediator* badge_mediator_ = nil;
 };
 
 // Test that the BadgeMediator responds with no displayed and fullscreen badge
 // when there are no Infobars added and the BrowserState is not OffTheRecord.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestNoInfobar) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestNoInfobar) {
+  InsertActivatedWebState(/*index=*/0);
   EXPECT_FALSE(badge_consumer_.displayedBadge);
-  EXPECT_FALSE(badge_consumer_.hasIncognitoBadge);
+  EXPECT_EQ(is_off_the_record(), badge_consumer_.hasFullscreenBadge);
 }
 
 // Test that the BadgeMediator responds with one new badge when an infobar is
 // added
-TEST_F(BadgeMediatorTest, BadgeMediatorTestAddInfobar) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestAddInfobar) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
   ASSERT_TRUE(badge_consumer_.displayedBadge);
   EXPECT_EQ(badge_consumer_.displayedBadge.badgeType,
@@ -148,8 +164,8 @@ TEST_F(BadgeMediatorTest, BadgeMediatorTestAddInfobar) {
 
 // Test that the BadgeMediator handled the removal of the correct badge when two
 // infobars are added and then one is removed.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestRemoveInfobar) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestRemoveInfobar) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
   InfoBarIOS* second_infobar = AddInfobar(kSecondInfobarType);
   EXPECT_EQ(badge_consumer_.displayedBadge.badgeType,
@@ -160,8 +176,8 @@ TEST_F(BadgeMediatorTest, BadgeMediatorTestRemoveInfobar) {
             BadgeType::kBadgeTypePasswordSave);
 }
 
-TEST_F(BadgeMediatorTest, BadgeMediatorTestMarkAsRead) {
-  AddAndActivateWebState(/*index=*/0, /*incognito=*/false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestMarkAsRead) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
   // Since there is only one badge, it should be marked as read.
   EXPECT_FALSE(badge_consumer_.hasUnreadBadge);
@@ -178,19 +194,19 @@ TEST_F(BadgeMediatorTest, BadgeMediatorTestMarkAsRead) {
 
 // Test that the BadgeMediator updates the current badges to none when switching
 // to a second WebState after an infobar is added to the first WebState.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestSwitchWebState) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestSwitchWebState) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
   ASSERT_TRUE(badge_consumer_.displayedBadge);
   EXPECT_EQ(badge_consumer_.displayedBadge.badgeType,
             BadgeType::kBadgeTypePasswordSave);
-  AddAndActivateWebState(1, false);
+  InsertActivatedWebState(/*index=*/1);
   EXPECT_FALSE(badge_consumer_.displayedBadge);
 }
 
 // Test that the BadgeMediator updates the badge when it is accepted.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestAcceptedBadge) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestAcceptedBadge) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
   ASSERT_TRUE(badge_consumer_.displayedBadge);
   EXPECT_FALSE(badge_consumer_.displayedBadge.badgeState &= BadgeStateAccepted);
@@ -199,29 +215,22 @@ TEST_F(BadgeMediatorTest, BadgeMediatorTestAcceptedBadge) {
   EXPECT_TRUE(badge_consumer_.displayedBadge.badgeState &= BadgeStateAccepted);
 }
 
-// Test that the BadgeMediator adds an incognito badge when the webstatelist
-// changes.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestIncognito) {
-  AddAndActivateWebState(0, true);
-  EXPECT_TRUE(badge_consumer_.hasIncognitoBadge);
-}
-
 // Test that the BadgeMediator updates the current badges when the starting
 // active WebState already has a badge. This simulates an app launch after an
 // update when the WebStateList is preserved but the LocationBar (and therefore
 // the BadgeMediator) is restarted from scratch.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestRestartWithInfobar) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestRestartWithInfobar) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
 
   // Simulate reload of app, but preservation of WebStateList.
+  [badge_mediator_ disconnect];
   badge_mediator_ = nil;
   badge_consumer_ = nil;
 
   badge_consumer_ = [[FakeBadgeConsumer alloc] init];
-  badge_mediator_ =
-      [[BadgeMediator alloc] initWithConsumer:badge_consumer_
-                                 webStateList:web_state_list_.get()];
+  badge_mediator_ = [[BadgeMediator alloc] initWithBrowser:browser()];
+  badge_mediator_.consumer = badge_consumer_;
   ASSERT_TRUE(badge_consumer_.displayedBadge);
   EXPECT_EQ(badge_consumer_.displayedBadge.badgeType,
             BadgeType::kBadgeTypePasswordSave);
@@ -230,13 +239,18 @@ TEST_F(BadgeMediatorTest, BadgeMediatorTestRestartWithInfobar) {
 // Test that the BadgeMediator clears its badges when the last WebState is
 // detached and a new WebState is added. This test also makes sure that closing
 // the last WebState doesn't break anything.
-TEST_F(BadgeMediatorTest, BadgeMediatorTestCloseLastTab) {
-  AddAndActivateWebState(0, false);
+TEST_P(BadgeMediatorTest, BadgeMediatorTestCloseLastTab) {
+  InsertActivatedWebState(/*index=*/0);
   AddInfobar(kFirstInfobarType);
   ASSERT_TRUE(badge_consumer_.displayedBadge);
   EXPECT_EQ(badge_consumer_.displayedBadge.badgeType,
             BadgeType::kBadgeTypePasswordSave);
-  web_state_list_->DetachWebStateAt(0);
-  AddAndActivateWebState(0, false);
+  web_state_list_.DetachWebStateAt(0);
+  InsertActivatedWebState(/*index=*/0);
   ASSERT_FALSE(badge_consumer_.displayedBadge);
 }
+
+INSTANTIATE_TEST_SUITE_P(/* No InstantiationName */,
+                         BadgeMediatorTest,
+                         testing::Values(TestParam::kNormal,
+                                         TestParam::kOffTheRecord));

ios/chrome/browser/ui/location_bar/location_bar_coordinator.mm

@@ -183,9 +183,8 @@ const int kLocationAuthorizationStatusCount = 5;
   [self.viewController setBadgeView:self.badgeViewController.view];
   [self.badgeViewController didMoveToParentViewController:self.viewController];
   // Create BadgeMediator and set the viewController as its consumer.
-  self.badgeMediator =
-      [[BadgeMediator alloc] initWithConsumer:self.badgeViewController
-                                 webStateList:self.webStateList];
+  self.badgeMediator = [[BadgeMediator alloc] initWithBrowser:self.browser];
+  self.badgeMediator.consumer = self.badgeViewController;
   self.badgeMediator.dispatcher =
       static_cast<id<InfobarCommands, BrowserCoordinatorCommands>>(
           self.dispatcher);
@@ -218,6 +217,7 @@ const int kLocationAuthorizationStatusCount = 5;
   [self.omniboxPopupCoordinator stop];
   [self.omniboxCoordinator stop];
   [self.badgeMediator disconnect];
+  self.badgeMediator = nil;
   _editController.reset();
 
   self.viewController = nil;