From fdfbe66513a22a8da1ae5895612de3670d018c06 Mon Sep 17 00:00:00 2001
From: Cesar Gonzalez <cesar.a.gonzalezcs@gmail.com>
Date: Tue, 8 Oct 2024 15:47:20 -0500
Subject: [PATCH] [PM-13284] Implement method to ensure that we can handle
 logic when switching the notification improvements feature flag (#11468)

---
 .../overlay-notifications.background.spec.ts  |  7 ++-
 .../overlay-notifications.background.ts       | 47 +++++++++++++++----
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/apps/browser/src/autofill/background/overlay-notifications.background.spec.ts b/apps/browser/src/autofill/background/overlay-notifications.background.spec.ts
index 1d4ff20b79..8bac8ea691 100644
--- a/apps/browser/src/autofill/background/overlay-notifications.background.spec.ts
+++ b/apps/browser/src/autofill/background/overlay-notifications.background.spec.ts
@@ -1,4 +1,5 @@
 import { mock, MockProxy } from "jest-mock-extended";
+import { BehaviorSubject } from "rxjs";
 
 import { CLEAR_NOTIFICATION_LOGIN_DATA_DURATION } from "@bitwarden/common/autofill/constants";
 import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
@@ -24,6 +25,7 @@ import { OverlayNotificationsBackground } from "./overlay-notifications.backgrou
 
 describe("OverlayNotificationsBackground", () => {
   let logService: MockProxy<LogService>;
+  let getFeatureFlagMock$: BehaviorSubject<boolean>;
   let configService: MockProxy<ConfigService>;
   let notificationBackground: NotificationBackground;
   let getEnableChangedPasswordPromptSpy: jest.SpyInstance;
@@ -33,7 +35,10 @@ describe("OverlayNotificationsBackground", () => {
   beforeEach(async () => {
     jest.useFakeTimers();
     logService = mock<LogService>();
-    configService = mock<ConfigService>();
+    getFeatureFlagMock$ = new BehaviorSubject(true);
+    configService = mock<ConfigService>({
+      getFeatureFlag$: jest.fn().mockReturnValue(getFeatureFlagMock$),
+    });
     notificationBackground = mock<NotificationBackground>();
     getEnableChangedPasswordPromptSpy = jest
       .spyOn(notificationBackground, "getEnableChangedPasswordPrompt")
diff --git a/apps/browser/src/autofill/background/overlay-notifications.background.ts b/apps/browser/src/autofill/background/overlay-notifications.background.ts
index 5ba3e388b7..5ea3e8b8d6 100644
--- a/apps/browser/src/autofill/background/overlay-notifications.background.ts
+++ b/apps/browser/src/autofill/background/overlay-notifications.background.ts
@@ -1,4 +1,5 @@
-import { Subject, switchMap, timer } from "rxjs";
+import { startWith, Subject, Subscription, switchMap, timer } from "rxjs";
+import { pairwise } from "rxjs/operators";
 
 import { CLEAR_NOTIFICATION_LOGIN_DATA_DURATION } from "@bitwarden/common/autofill/constants";
 import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
@@ -23,6 +24,7 @@ export class OverlayNotificationsBackground implements OverlayNotificationsBackg
   private websiteOriginsWithFields: WebsiteOriginsWithFields = new Map();
   private activeFormSubmissionRequests: ActiveFormSubmissionRequests = new Set();
   private modifyLoginCipherFormData: ModifyLoginCipherFormDataForTab = new Map();
+  private featureFlagState$: Subscription;
   private clearLoginCipherFormDataSubject: Subject<void> = new Subject();
   private notificationFallbackTimeout: number | NodeJS.Timeout | null;
   private readonly formSubmissionRequestMethods: Set<string> = new Set(["POST", "PUT", "PATCH"]);
@@ -42,19 +44,35 @@ export class OverlayNotificationsBackground implements OverlayNotificationsBackg
    * Initialize the overlay notifications background service.
    */
   async init() {
-    const featureFlagActive = await this.configService.getFeatureFlag(
-      FeatureFlag.NotificationBarAddLoginImprovements,
-    );
-    if (!featureFlagActive) {
-      return;
-    }
-
-    this.setupExtensionListeners();
+    this.featureFlagState$ = this.configService
+      .getFeatureFlag$(FeatureFlag.NotificationBarAddLoginImprovements)
+      .pipe(startWith(undefined), pairwise())
+      .subscribe(([prev, current]) => this.handleInitFeatureFlagChange(prev, current));
     this.clearLoginCipherFormDataSubject
       .pipe(switchMap(() => timer(CLEAR_NOTIFICATION_LOGIN_DATA_DURATION)))
       .subscribe(() => this.modifyLoginCipherFormData.clear());
   }
 
+  /**
+   * Handles enabling/disabling the extension listeners that trigger the
+   * overlay notifications based on the feature flag state.
+   *
+   * @param previousValue - The previous value of the feature flag
+   * @param currentValue - The current value of the feature flag
+   */
+  private handleInitFeatureFlagChange = (previousValue: boolean, currentValue: boolean) => {
+    if (previousValue === currentValue) {
+      return;
+    }
+
+    if (currentValue) {
+      this.setupExtensionListeners();
+      return;
+    }
+
+    this.removeExtensionListeners();
+  };
+
   /**
    * Handles the response from the content script with the page details. Triggers an initialization
    * of the add login or change password notification if the conditions are met.
@@ -495,11 +513,20 @@ export class OverlayNotificationsBackground implements OverlayNotificationsBackg
    * Sets up the listeners for the extension messages and the tab events.
    */
   private setupExtensionListeners() {
-    BrowserApi.messageListener("overlay-notifications", this.handleExtensionMessage);
+    BrowserApi.addListener(chrome.runtime.onMessage, this.handleExtensionMessage);
     chrome.tabs.onRemoved.addListener(this.handleTabRemoved);
     chrome.tabs.onUpdated.addListener(this.handleTabUpdated);
   }
 
+  /**
+   * Removes the listeners for the extension messages and the tab events.
+   */
+  private removeExtensionListeners() {
+    BrowserApi.removeListener(chrome.runtime.onMessage, this.handleExtensionMessage);
+    chrome.tabs.onRemoved.removeListener(this.handleTabRemoved);
+    chrome.tabs.onUpdated.removeListener(this.handleTabUpdated);
+  }
+
   /**
    * Handles messages that are sent to the extension background.
    *