1 files changed, 41 insertions, 20 deletions

diff --git a/test/e2e/beta/from-import-beta-ui.spec.js b/test/e2e/beta/from-import-beta-ui.spec.js
index 32f57b157..c271461d3 100644
--- a/test/e2e/beta/from-import-beta-ui.spec.js
+++ b/test/e2e/beta/from-import-beta-ui.spec.js
@@ -12,9 +12,11 @@ const {
 } = require('../func')
 const {
   checkBrowserForConsoleErrors,
+  closeAllWindowHandlesExcept,
   verboseReportOnFailure,
   findElement,
   findElements,
+  loadExtension,
 } = require('./helpers')
 
 
@@ -25,6 +27,7 @@ describe('Using MetaMask with an existing account', function () {
   const testSeedPhrase = 'phrase upgrade clock rough situate wedding elder clever doctor stamp excess tent'
   const testAddress = '0xE18035BF8712672935FDB4e5e431b1a0183d2DFC'
   const testPrivateKey2 = '14abe6f4aab7f9f626fe981c864d0adeb5685f289ac9270c27b8fd790b4235d6'
+  const tinyDelayMs = 500
   const regularDelayMs = 1000
   const largeDelayMs = regularDelayMs * 2
 
@@ -74,37 +77,55 @@ describe('Using MetaMask with an existing account', function () {
 
   describe('New UI setup', async function () {
     it('switches to first tab', async function () {
+      await delay(tinyDelayMs)
       const [firstTab] = await driver.getAllWindowHandles()
       await driver.switchTo().window(firstTab)
       await delay(regularDelayMs)
     })
 
-    it('use the local network', async function () {
-      const networkSelector = await findElement(driver, By.css('#network_component'))
-      await networkSelector.click()
-      await delay(regularDelayMs)
-
-      const [localhost] = await findElements(driver, By.xpath(`//li[contains(text(), 'Localhost')]`))
-      await localhost.click()
-      await delay(regularDelayMs)
-    })
-
     it('selects the new UI option', async () => {
-      const button = await findElement(driver, By.xpath("//p[contains(text(), 'Try Beta Version')]"))
+      try {
+        const overlay = await findElement(driver, By.css('.full-flex-height'))
+        await driver.wait(until.stalenessOf(overlay))
+      } catch (e) {}
+
+      const button = await findElement(driver, By.xpath("//button[contains(text(), 'Try it now')]"))
       await button.click()
       await delay(regularDelayMs)
 
       // Close all other tabs
-      const [oldUi, infoPage, newUi] = await driver.getAllWindowHandles()
-
-      const newUiOrInfoPage = newUi || infoPage
-      await driver.switchTo().window(oldUi)
-      await driver.close()
-      if (infoPage !== newUiOrInfoPage) {
-        await driver.switchTo().window(infoPage)
-        await driver.close()
+      const [tab0, tab1, tab2] = await driver.getAllWindowHandles()
+      await driver.switchTo().window(tab0)
+      await delay(tinyDelayMs)
+
+      let selectedUrl = await driver.getCurrentUrl()
+      await delay(tinyDelayMs)
+      if (tab0 && selectedUrl.match(/popup.html/)) {
+        await closeAllWindowHandlesExcept(driver, tab0)
+      } else if (tab1) {
+        await driver.switchTo().window(tab1)
+        selectedUrl = await driver.getCurrentUrl()
+        await delay(tinyDelayMs)
+        if (selectedUrl.match(/popup.html/)) {
+          await closeAllWindowHandlesExcept(driver, tab1)
+        } else if (tab2) {
+          await driver.switchTo().window(tab2)
+          selectedUrl = await driver.getCurrentUrl()
+          selectedUrl.match(/popup.html/) && await closeAllWindowHandlesExcept(driver, tab2)
+        }
+      } else {
+        throw new Error('popup.html not found')
       }
-      await driver.switchTo().window(newUiOrInfoPage)
+      await delay(regularDelayMs)
+      const [appTab] = await driver.getAllWindowHandles()
+      await driver.switchTo().window(appTab)
+      await delay(tinyDelayMs)
+
+      const okayBtn = await findElement(driver, By.xpath("//button[contains(text(), 'Ok')]"))
+      await okayBtn.click()
+      await delay(regularDelayMs)
+
+      await loadExtension(driver, extensionId)
       await delay(regularDelayMs)
 
       const continueBtn = await findElement(driver, By.css('.welcome-screen__button'))