Conditionally close info page in beta e2e tests.

2 files changed, 14 insertions, 4 deletions

diff --git a/test/e2e/beta/from-import-beta-ui.spec.js b/test/e2e/beta/from-import-beta-ui.spec.js
index 7fec5fa09..f24a8e1d2 100644
--- a/test/e2e/beta/from-import-beta-ui.spec.js
+++ b/test/e2e/beta/from-import-beta-ui.spec.js
@@ -89,18 +89,23 @@ describe('Using MetaMask with an existing account', function () {
     })
 
     it('selects the new UI option', async () => {
-      const button = await driver.findElement(By.xpath("//p[contains(text(), 'Try Beta Version')]"))
+      const button = await findElement(driver, By.xpath("//p[contains(text(), 'Try Beta Version')]"))
       await button.click()
       await delay(regularDelayMs)
 
       // Close all other tabs
-      const [oldUi, newUi] = await driver.getAllWindowHandles()
+      let [oldUi, infoPage, newUi] = await driver.getAllWindowHandles()
+      newUi = newUi || infoPage
       await driver.switchTo().window(oldUi)
       await driver.close()
+      if (infoPage !== newUi) {
+        await driver.switchTo().window(infoPage)
+        await driver.close()
+      }
       await driver.switchTo().window(newUi)
       await delay(regularDelayMs)
 
-      const [continueBtn] = await driver.findElements(By.css('.welcome-screen__button'))
+      const continueBtn = await findElement(driver, By.css('.welcome-screen__button'))
       await continueBtn.click()
       await delay(regularDelayMs)
     })
diff --git a/test/e2e/beta/metamask-beta-ui.spec.js b/test/e2e/beta/metamask-beta-ui.spec.js
index c5f163611..f66f4a618 100644
--- a/test/e2e/beta/metamask-beta-ui.spec.js
+++ b/test/e2e/beta/metamask-beta-ui.spec.js
@@ -93,9 +93,14 @@ describe('MetaMask', function () {
       await delay(regularDelayMs)
 
       // Close all other tabs
-      const [oldUi, newUi] = await driver.getAllWindowHandles()
+      let [oldUi, infoPage, newUi] = await driver.getAllWindowHandles()
+      newUi = newUi || infoPage
       await driver.switchTo().window(oldUi)
       await driver.close()
+      if (infoPage !== newUi) {
+        await driver.switchTo().window(infoPage)
+        await driver.close()
+      }
       await driver.switchTo().window(newUi)
       await delay(regularDelayMs)