From 6f55671f8f6a622727e549bd6a6eda5096c1fa29 Mon Sep 17 00:00:00 2001
From: Francis McKenzie <francis.mckenzie@gmail.com>
Date: Mon, 30 Dec 2019 22:16:34 +0100
Subject: [PATCH] Container locking - simplified test code

---
 src/js/popup.js | 20 ++++++++++++--------
 test/helper.js  | 18 +++---------------
 2 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/src/js/popup.js b/src/js/popup.js
index 5b4a537..4e90df9 100644
--- a/src/js/popup.js
+++ b/src/js/popup.js
@@ -370,6 +370,16 @@ const Logic = {
       value
     });
   },
+  
+  lockOrUnlockContainer(userContextId, isLocked) {
+    return browser.runtime.sendMessage({
+      method: "lockOrUnlockContainer",
+      message: {
+        userContextId: userContextId,
+        isLocked: !!isLocked
+      }
+    });
+  },
 
   generateIdentityName() {
     const defaultName = "Container #";
@@ -1041,16 +1051,10 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
       
       Logic.addEnterHandler(lockElement, async () => {
         try {
-          await browser.runtime.sendMessage({
-            method: "lockOrUnlockContainer",
-            message: {
-              userContextId: Logic.currentUserContextId(),
-              isLocked: !isLocked
-            }
-          });
+          await Logic.lockOrUnlockContainer(Logic.currentUserContextId(), !isLocked);
           this.showAssignedContainers(assignments, !isLocked);
         } catch (e) {
-          throw new Error("Failed to lock/unlock. ", e.message);
+          throw new Error("Failed to lock/unlock container: " + e.message);
         }
       });
       /* Container locking end */
diff --git a/test/helper.js b/test/helper.js
index 039ffe8..21ca95c 100644
--- a/test/helper.js
+++ b/test/helper.js
@@ -53,21 +53,9 @@ module.exports = {
     
     // https://github.com/mozilla/multi-account-containers/issues/847
     async setContainerIsLocked(cookieStoreId, isLocked) {
-      const identityStateKey = this.getIdentityStateContainerStoreKey(cookieStoreId);
-      const identityState = await background.browser.storage.local.get([identityStateKey]) || {};
-      if (isLocked) {
-        identityState.isLocked = "locked";
-      } else {
-        delete identityState.isLocked;
-      }
-      // Must have valid 'hiddenTabs', otherwise backgroundLogic.showTabs() throws error
-      if (!identityState.hiddenTabs) { identityState.hiddenTabs = []; }
-      await background.browser.storage.local.set({[identityStateKey]: identityState});
-    },
-    
-    getIdentityStateContainerStoreKey(cookieStoreId) {
-      const storagePrefix = "identitiesState@@_";
-      return `${storagePrefix}${cookieStoreId}`;      
+      const Logic = popup.dom.window.eval("Logic");
+      const userContextId = Logic.userContextId(cookieStoreId);
+      await Logic.lockOrUnlockContainer(userContextId, isLocked);
     }
   }
 };