--- a/browser/components/extensions/test/browser/browser_ext_connect_and_move_tabs.js
+++ b/browser/components/extensions/test/browser/browser_ext_connect_and_move_tabs.js
@@ -32,20 +32,24 @@ function loadExtension() {
             // When the port is unexpectedly disconnected, postMessage will throw an error.
             port.postMessage("ping");
           } catch (e) {
             browser.test.fail(`Error: ${e} :: ${e.stack}`);
             browser.test.sendMessage("port_ping_ponged_before_disconnect");
           }
         });
 
-        browser.runtime.onMessage.addListener((msg, sender) => {
-          browser.test.assertEq("disconnect-me", msg, "expected message");
-          port.disconnect();
-          // Now port.onDisconnect should fire in the content script.
+        browser.runtime.onMessage.addListener(async (msg, sender) => {
+          if (msg == "disconnect-me") {
+            port.disconnect();
+            // Now port.onDisconnect should fire in the content script.
+          } else if (msg == "close-tab") {
+            await browser.tabs.remove(sender.tab.id);
+            browser.test.sendMessage("closed_tab");
+          }
         });
       });
 
       browser.test.onMessage.addListener(msg => {
         browser.test.assertEq("open_extension_tab", msg, "expected message");
         browser.tabs.create({url: "tab.html"});
       });
     },
@@ -57,39 +61,37 @@ function loadExtension() {
       `,
       "script.js": function() {
         let port = browser.runtime.connect();
         port.onMessage.addListener(msg => {
           browser.test.assertEq("ping", msg, "expected message");
           browser.test.sendMessage("port_ping_ponged_before_disconnect");
           port.onDisconnect.addListener(() => {
             browser.test.sendMessage("port_disconnected");
+            browser.runtime.sendMessage("close-tab");
           });
           browser.runtime.sendMessage("disconnect-me");
         });
         port.postMessage("connect_from_script");
       },
     },
   });
 }
 
 add_task(async function contentscript_connect_and_move_tabs() {
   let extension = loadExtension();
   await extension.startup();
   await BrowserTestUtils.openNewForegroundTab(gBrowser, "http://mochi.test:8888/?discoTest");
   await extension.awaitMessage("port_ping_ponged_before_disconnect");
   await extension.awaitMessage("port_disconnected");
-  // Must use gBrowser.selectedTab instead of the return value of
-  // BrowserTestUtils.openNewForegroundTab because the latter does not refer to
-  // the tab because the tab is moved between windows during the test.
-  await BrowserTestUtils.removeTab(gBrowser.selectedTab);
+  await extension.awaitMessage("closed_tab");
   await extension.unload();
 });
 
 add_task(async function extension_tab_connect_and_move_tabs() {
   let extension = loadExtension();
   await extension.startup();
   extension.sendMessage("open_extension_tab");
   await extension.awaitMessage("port_ping_ponged_before_disconnect");
   await extension.awaitMessage("port_disconnected");
-  // Upon unloading the extension, the extension tab is automatically removed.
+  await extension.awaitMessage("closed_tab");
   await extension.unload();
 });