From 065561e0d73504b633d21121dfa05bc670d21be5 Mon Sep 17 00:00:00 2001
From: Aishat Akinyemi <aishatadenike01@gmail.com>
Date: Thu, 24 Oct 2019 13:37:40 +0100
Subject: [PATCH] refactoring messageHandler.js to use
 badge.displayBrowserActionBadge

---
 src/js/background/badge.js          | 31 ++++++++++++++++++++---------
 src/js/background/messageHandler.js | 13 +-----------
 2 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/src/js/background/badge.js b/src/js/background/badge.js
index 78cd9f1..92e5b71 100644
--- a/src/js/background/badge.js
+++ b/src/js/background/badge.js
@@ -1,8 +1,8 @@
 const MAJOR_VERSIONS = ["2.3.0", "2.4.0"];
 const badge = {
   async init() {
-    const currentWindow = await browser.windows.getCurrent();
-    this.displayBrowserActionBadge(currentWindow.incognito);
+    const showVersionIndicator = await browser.windows.getCurrent();
+    this.displayBrowserActionBadge(showVersionIndicator);
   },
 
   disableAddon(tabId) {
@@ -10,16 +10,29 @@ const badge = {
     browser.browserAction.setTitle({ tabId, title: "Containers disabled in Private Browsing Mode" });
   },
 
-  async displayBrowserActionBadge() {
+  async displayBrowserActionBadge(action) {
     const extensionInfo = await backgroundLogic.getExtensionInfo();
-    const storage = await browser.storage.local.get({browserActionBadgesClicked: []});
-
-    if (MAJOR_VERSIONS.indexOf(extensionInfo.version) > -1 &&
-        storage.browserActionBadgesClicked.indexOf(extensionInfo.version) < 0) {
-      browser.browserAction.setBadgeBackgroundColor({color: "rgba(0,217,0,255)"});
-      browser.browserAction.setBadgeText({text: "NEW"});
+    function changeBadgeColorText(color, text){
+      browser.browserAction.setBadgeBackgroundColor({color: color});
+      browser.browserAction.setBadgeText({text: text});
+    }
+    if(action==="showVersionIndicator") {    
+      const ActionBadgesClickedStorage = await browser.storage.local.get({browserActionBadgesClicked: []});
+      if (MAJOR_VERSIONS.indexOf(extensionInfo.version) > -1 &&
+          ActionBadgesClickedStorage.browserActionBadgesClicked.indexOf(extensionInfo.version) < 0) {
+        changeBadgeColorText("rgba(0,217,0,255)", "NEW");
+      }
+    }
+    else if (action==="showAchievement") {
+      const achievementsStorage = await browser.storage.local.get({achievements: []});
+      achievementsStorage.achievements.push({"name": "manyContainersOpened", "done": false});
+      // use set and spread to create a unique array
+      const achievements = [...new Set(achievementsStorage.achievements)];
+      browser.storage.local.set({achievements});
+      changeBadgeColorText("rgba(0,217,0,255)", "NEW");
     }
   }
+
 };
 
 badge.init();
diff --git a/src/js/background/messageHandler.js b/src/js/background/messageHandler.js
index 9fbe88e..8775e6e 100644
--- a/src/js/background/messageHandler.js
+++ b/src/js/background/messageHandler.js
@@ -171,23 +171,12 @@ const messageHandler = {
 
     // When the user opens their _ tab, give them the achievement
     if (countOfContainerTabsOpened === 100) {
-      const storage = await browser.storage.local.get({achievements: []});
-      storage.achievements.push({"name": "manyContainersOpened", "done": false});
-      // use set and spread to create a unique array
-      const achievements = [...new Set(storage.achievements)];
-      browser.storage.local.set({achievements});
-      browser.browserAction.setBadgeBackgroundColor({color: "rgba(0,217,0,255)"});
-      browser.browserAction.setBadgeText({text: "NEW"});
+      badge.displayBrowserActionBadge("showAchievement");      
     }
   },
 
   async onFocusChangedCallback(windowId) {
     assignManager.removeContextMenu();
-    // browserAction loses background color in new windows ...
-    // https://bugzil.la/1314674
-    // https://github.com/mozilla/testpilot-containers/issues/608
-    // ... so re-call displayBrowserActionBadge on window changes
-    badge.displayBrowserActionBadge();
     browser.tabs.query({active: true, windowId}).then((tabs) => {
       if (tabs && tabs[0]) {
         assignManager.calculateContextMenu(tabs[0]);