From 0bf47ad8069f4b064fd6c415ed3a005ef2f7aa8c Mon Sep 17 00:00:00 2001
From: Francis McKenzie <francis.mckenzie@gmail.com>
Date: Sun, 29 Dec 2019 17:25:27 +0100
Subject: [PATCH] Unlocked tab should open adjacent to locked tab, not at end
 of tabs. Unit tests updated accordingly.

---
 src/js/background/assignManager.js | 31 +++++++++++++++++++++++++++---
 test/features/lock.test.js         | 23 +++++++++++-----------
 test/helper.js                     | 20 +++++++++----------
 3 files changed, 48 insertions(+), 26 deletions(-)

diff --git a/src/js/background/assignManager.js b/src/js/background/assignManager.js
index b80b937..575eab2 100644
--- a/src/js/background/assignManager.js
+++ b/src/js/background/assignManager.js
@@ -211,10 +211,13 @@ const assignManager = {
     }
 
     if (mustReloadPageInDefaultContainerDueToLocking) {
-      // Open new tab in default container
-      browser.tabs.create({url: options.url});
+      this.reloadPageInDefaultContainer(
+        options.url,
+        tab.index + 1,
+        tab.active,
+        openTabId
+      );
     } else {
-      // Open new tab in assigned container
       this.reloadPageInContainer(
         options.url,
         userContextId,
@@ -545,6 +548,28 @@ const assignManager = {
       return `%${charCode}`;
     });
   },
+  
+  reloadPageInDefaultContainer(url, index, active, openerTabId) {
+    // To create a new tab in the default container, it is easiest just to omit the
+    // cookieStoreId entirely.
+    // 
+    // Unfortunately, if you create a new tab WITHOUT a cookieStoreId but WITH an openerTabId,
+    // then the new tab automatically inherits the opener tab's cookieStoreId.
+    // I.e. it opens in the wrong container!
+    // 
+    // So we have to explicitly pass in a cookieStoreId when creating the tab, since we
+    // are specifying the openerTabId. There doesn't seem to be any way
+    // to look up the default container's cookieStoreId programatically, so sadly
+    // we have to hardcode it here as "firefox-default". This is potentially
+    // not cross-browser compatible.
+    // 
+    // Note that we could have just omitted BOTH cookieStoreId and openerTabId. But the
+    // drawback then is that if the user later closes the newly-created tab, the browser
+    // does not automatically return to the original opener tab. To get this desired behaviour,
+    // we MUST specify the openerTabId when creating the new tab.
+    const cookieStoreId = "firefox-default";
+    browser.tabs.create({url, cookieStoreId, index, active, openerTabId});
+  },
 
   reloadPageInContainer(url, currentUserContextId, userContextId, index, active, neverAsk = false, openerTabId = null) {
     const cookieStoreId = backgroundLogic.cookieStoreId(userContextId);
diff --git a/test/features/lock.test.js b/test/features/lock.test.js
index 6b97d88..62bb313 100644
--- a/test/features/lock.test.js
+++ b/test/features/lock.test.js
@@ -4,7 +4,8 @@ describe("Lock Feature", () => {
     id: 1,
     cookieStoreId: "firefox-container-1",
     url: "http://example.com",
-    index: 0
+    index: 0,
+    active: true
   };
   beforeEach(async () => {
     await helper.browser.initializeWithTab(activeTab);
@@ -19,10 +20,7 @@ describe("Lock Feature", () => {
     describe("open different URL in same tab", () => {
       const differentURL = "http://example2.com";
       beforeEach(async () => {
-        await helper.browser.updateTab(activeTab, {
-          url: differentURL,
-          resetHistory: true
-        });
+        await helper.browser.browseToURL(activeTab.id, differentURL);
       });
 
       it("should not open a new tab", () => {
@@ -36,16 +34,17 @@ describe("Lock Feature", () => {
     
         describe("open different URL in same tab", () => {
           beforeEach(async () => {
-            await helper.browser.updateTab(activeTab, {
-              url: differentURL,
-              resetHistory: true
-            });
+            await helper.browser.browseToURL(activeTab.id, differentURL);
           });
 
           it("should open a new tab in the default container", () => {
-            background.browser.tabs.create.should.have.been.calledWith({
-              url: differentURL
-            });
+            background.browser.tabs.create.should.have.been.calledWith(sinon.match({
+              url: differentURL,
+              cookieStoreId: "firefox-default",
+              openerTabId: activeTab.id,
+              index: activeTab.index + 1,
+              active: activeTab.active
+            }));
           });
         });
       });
diff --git a/test/helper.js b/test/helper.js
index 772f7ea..039ffe8 100644
--- a/test/helper.js
+++ b/test/helper.js
@@ -30,18 +30,16 @@ module.exports = {
     async openNewTab(tab, options = {}) {
       return background.browser.tabs._create(tab, options);
     },
-
+    
     // https://github.com/mozilla/multi-account-containers/issues/847
-    async updateTab(tab, options = {}) {
-      const updatedTab = {};
-      for (const key in tab) {
-        updatedTab[key] = tab[key];
-      }
-      for (const key in options) {
-        updatedTab[key] = options[key];
-      }
-      return this.openNewTab(updatedTab);
-    },
+    async browseToURL(tabId, url) {
+      const [promise] = background.browser.webRequest.onBeforeRequest.addListener.yield({
+        frameId: 0,
+        tabId: tabId,
+        url: url
+      });
+      return promise;
+    }
   },
 
   popup: {