From ebe6f65705300c669a58e4ebef49e4f265569886 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 10:42:20 +1000
Subject: [PATCH 01/28] add dvc python path into user settings

---
 extension/src/python/index.ts       | 6 +-----
 extension/src/python/path.ts        | 7 +++++++
 extension/src/test/cli/util.ts      | 2 +-
 extension/src/test/e2e/wdio.conf.ts | 8 ++++++++
 4 files changed, 17 insertions(+), 6 deletions(-)
 create mode 100644 extension/src/python/path.ts

diff --git a/extension/src/python/index.ts b/extension/src/python/index.ts
index 8795813950..d987921b7c 100644
--- a/extension/src/python/index.ts
+++ b/extension/src/python/index.ts
@@ -1,4 +1,5 @@
 import { join } from 'path'
+import { getVenvBinPath } from './path'
 import { getProcessPlatform } from '../env'
 import { exists } from '../fileSystem'
 import { createProcessWithOutput } from '../processExecution'
@@ -11,11 +12,6 @@ const installPackages = (cwd: string, pythonBin: string, ...args: string[]) => {
   })
 }
 
-export const getVenvBinPath = (cwd: string, envDir: string, name: string) =>
-  getProcessPlatform() === 'win32'
-    ? join(cwd, envDir, 'Scripts', `${name}.exe`)
-    : join(cwd, envDir, 'bin', name)
-
 export const setupVenv = async (
   cwd: string,
   envDir: string,
diff --git a/extension/src/python/path.ts b/extension/src/python/path.ts
new file mode 100644
index 0000000000..70468ba1bf
--- /dev/null
+++ b/extension/src/python/path.ts
@@ -0,0 +1,7 @@
+import { join } from 'path'
+import { getProcessPlatform } from '../env'
+
+export const getVenvBinPath = (cwd: string, envDir: string, name: string) =>
+  getProcessPlatform() === 'win32'
+    ? join(cwd, envDir, 'Scripts', `${name}.exe`)
+    : join(cwd, envDir, 'bin', name)
diff --git a/extension/src/test/cli/util.ts b/extension/src/test/cli/util.ts
index 0aa9634942..07fb063f22 100644
--- a/extension/src/test/cli/util.ts
+++ b/extension/src/test/cli/util.ts
@@ -4,7 +4,7 @@ import { CliExecutor } from '../../cli/executor'
 import { CliReader } from '../../cli/reader'
 import { Config } from '../../config'
 import { exists } from '../../fileSystem'
-import { getVenvBinPath } from '../../python'
+import { getVenvBinPath } from '../../python/path'
 import { dvcDemoPath } from '../util'
 
 const config = {
diff --git a/extension/src/test/e2e/wdio.conf.ts b/extension/src/test/e2e/wdio.conf.ts
index 82b9f2dce8..4d5d955832 100644
--- a/extension/src/test/e2e/wdio.conf.ts
+++ b/extension/src/test/e2e/wdio.conf.ts
@@ -1,5 +1,6 @@
 import { join, resolve } from 'path'
 import { Options } from '@wdio/types'
+import { getVenvBinPath } from '../../python/path'
 
 export const config: Options.Testrunner = {
   after: async function () {
@@ -17,6 +18,13 @@ export const config: Options.Testrunner = {
       // @ts-expect-error these caps are not typed in WebdriverIO
       'wdio:vscodeOptions': {
         extensionPath: resolve(__dirname, '..', '..', '..'),
+        userSettings: {
+          'dvc.pythonPath': getVenvBinPath(
+            resolve(__dirname, '..', '..', '..', '..', 'demo'),
+            '.env',
+            'python'
+          )
+        },
         verboseLogging: false,
         workspacePath: resolve(__dirname, '..', '..', '..', '..', 'demo')
       }

From 5d0d9e9eb264c6685e3f115b922f928345bcff1e Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 11:02:15 +1000
Subject: [PATCH 02/28] use tree for the table role (update to React types)

---
 extension/src/test/e2e/pageObjects/locators.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/extension/src/test/e2e/pageObjects/locators.ts b/extension/src/test/e2e/pageObjects/locators.ts
index 88701000fb..9c3cec4553 100644
--- a/extension/src/test/e2e/pageObjects/locators.ts
+++ b/extension/src/test/e2e/pageObjects/locators.ts
@@ -7,7 +7,7 @@ export const experiments = {
   ...webview,
   expandRowButton: 'button[title="Expand Row"]',
   row: '[role=row]',
-  table: '[role=table]'
+  table: '[role=tree]'
 }
 
 export const plots = {

From 18e8d4b57bbb8c9c2637968a35ace3ad5219fb66 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 11:19:28 +1000
Subject: [PATCH 03/28] add end to end test workflow

---
 .github/workflows/end-to-end-test.yml | 36 +++++++++++++++++++++++++++
 extension/src/test/e2e/wdio.conf.ts   | 15 +++++++++++
 2 files changed, 51 insertions(+)
 create mode 100644 .github/workflows/end-to-end-test.yml

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
new file mode 100644
index 0000000000..68446c9197
--- /dev/null
+++ b/.github/workflows/end-to-end-test.yml
@@ -0,0 +1,36 @@
+name: End-To-End Test
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-node@v3
+        with:
+          node-version: 16
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+      - run: yarn run install-frozen-lockfile
+
+      - name: Test E2E
+        uses: GabrielBB/xvfb-action@v1
+        with:
+          run: yarn test:e2e
+          options: '-screen 0 1600x1200x24'
+      - uses: actions/upload-artifact@v3
+        if: failure()
+        with:
+          name: screenshots
+          path: |
+            test/screenshots
+            test/logs
diff --git a/extension/src/test/e2e/wdio.conf.ts b/extension/src/test/e2e/wdio.conf.ts
index 4d5d955832..8881d77458 100644
--- a/extension/src/test/e2e/wdio.conf.ts
+++ b/extension/src/test/e2e/wdio.conf.ts
@@ -1,12 +1,27 @@
 import { join, resolve } from 'path'
+import { mkdirp } from 'fs-extra'
 import { Options } from '@wdio/types'
 import { getVenvBinPath } from '../../python/path'
+import { Logger } from '../../common/logger'
 
 export const config: Options.Testrunner = {
   after: async function () {
     await browser.switchToFrame(null)
     await browser.switchToFrame(null)
   },
+  afterTest: async (test, __, { passed }) => {
+    if (passed) {
+      return
+    }
+
+    Logger.log('Capturing screenshot for debugging')
+
+    const screenshotDir = join(__dirname, 'screenshots')
+    mkdirp(screenshotDir)
+    await browser.saveScreenshot(
+      join(screenshotDir, `${test.parent} - ${test.title}.png`)
+    )
+  },
   baseUrl: 'http://localhost',
   before: async function () {
     await browser.setWindowSize(1600, 1200)

From 28cee39581f648829d1b3fe2e4322befdf9c7719 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 11:27:57 +1000
Subject: [PATCH 04/28] try mac for screenshot

---
 .github/workflows/end-to-end-test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index 68446c9197..32afec5fe2 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -7,10 +7,10 @@ on:
   workflow_dispatch:
 
 jobs:
-  build:
+  test:
     strategy:
       matrix:
-        os: [ubuntu-latest]
+        os: [macOS-latest]
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v3

From 3aa068c89492ef4181313c3706a241ee03741981 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 13:21:34 +1000
Subject: [PATCH 05/28] increase timeout and take screenshot

---
 extension/src/test/e2e/extension.test.ts | 11 +++++++++++
 extension/src/test/e2e/wdio.conf.ts      |  7 ++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 1985703be1..d592bd7187 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -1,3 +1,4 @@
+import { join } from 'path'
 import { suite, before, describe, it } from 'mocha'
 import {
   closeAllEditors,
@@ -11,7 +12,17 @@ import { delay } from '../../util/time'
 
 suite('DVC Extension For Visual Studio Code', () => {
   before('should finish loading the extension', async () => {
+    const screenshotDir = join(__dirname, 'screenshots')
+    await browser.saveScreenshot(
+      join(screenshotDir, 'startup - before container.png')
+    )
+
     await waitForViewContainerToLoad()
+
+    await browser.saveScreenshot(
+      join(screenshotDir, 'startup - after container.png')
+    )
+
     return dismissAllNotifications()
   })
 
diff --git a/extension/src/test/e2e/wdio.conf.ts b/extension/src/test/e2e/wdio.conf.ts
index 8881d77458..25d2d669b7 100644
--- a/extension/src/test/e2e/wdio.conf.ts
+++ b/extension/src/test/e2e/wdio.conf.ts
@@ -4,6 +4,8 @@ import { Options } from '@wdio/types'
 import { getVenvBinPath } from '../../python/path'
 import { Logger } from '../../common/logger'
 
+const screenshotDir = join(__dirname, 'screenshots')
+
 export const config: Options.Testrunner = {
   after: async function () {
     await browser.switchToFrame(null)
@@ -16,14 +18,13 @@ export const config: Options.Testrunner = {
 
     Logger.log('Capturing screenshot for debugging')
 
-    const screenshotDir = join(__dirname, 'screenshots')
-    mkdirp(screenshotDir)
     await browser.saveScreenshot(
       join(screenshotDir, `${test.parent} - ${test.title}.png`)
     )
   },
   baseUrl: 'http://localhost',
   before: async function () {
+    mkdirp(screenshotDir)
     await browser.setWindowSize(1600, 1200)
   },
   capabilities: [
@@ -60,5 +61,5 @@ export const config: Options.Testrunner = {
   reporters: ['spec'],
   services: ['vscode'],
   specs: ['./src/test/e2e/*.test.ts'],
-  waitforTimeout: 10000
+  waitforTimeout: 60000
 }

From f1856ebf832b4130c899a1662cc180b05e39e76e Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 13:24:25 +1000
Subject: [PATCH 06/28] add screenshots directory to gitignore

---
 extension/.gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/extension/.gitignore b/extension/.gitignore
index f1af6ec301..43200f9254 100644
--- a/extension/.gitignore
+++ b/extension/.gitignore
@@ -1,6 +1,7 @@
 *.vsix
 .vscode-test
 .wdio-vscode-service
+screenshots
 README.md
 LICENSE
 CHANGELOG.md

From 2cc9dfccba0ca235b9b33afdcf5bbb5e76ff68a0 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 14:02:16 +1000
Subject: [PATCH 07/28] pull data once on startup

---
 extension/src/test/e2e/extension.test.ts | 11 -----------
 extension/src/test/e2e/util.ts           | 18 ++++++++++++------
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index d592bd7187..1985703be1 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -1,4 +1,3 @@
-import { join } from 'path'
 import { suite, before, describe, it } from 'mocha'
 import {
   closeAllEditors,
@@ -12,17 +11,7 @@ import { delay } from '../../util/time'
 
 suite('DVC Extension For Visual Studio Code', () => {
   before('should finish loading the extension', async () => {
-    const screenshotDir = join(__dirname, 'screenshots')
-    await browser.saveScreenshot(
-      join(screenshotDir, 'startup - before container.png')
-    )
-
     await waitForViewContainerToLoad()
-
-    await browser.saveScreenshot(
-      join(screenshotDir, 'startup - after container.png')
-    )
-
     return dismissAllNotifications()
   })
 
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index 6fe063586b..9231697fc8 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -1,4 +1,5 @@
 import { ViewControl } from 'wdio-vscode-service'
+import { ElementArray } from 'webdriverio'
 
 export const dismissAllNotifications = () =>
   browser.waitUntil(async () => {
@@ -35,17 +36,22 @@ export const waitForViewContainerToLoad = async () => {
     return !!view
   })
 
-  return browser.waitUntil(async () => {
+  let currentProgressBars: ElementArray
+
+  await browser.waitUntil(async () => {
     const numberOfProgressBarsInContainer = 7
-    const currentProgressBars = await $$('.monaco-progress-container')
+    currentProgressBars = await $$('.monaco-progress-container')
 
-    if (
+    return !(
       currentProgressBars.length <
       initialProgressBars.length + numberOfProgressBarsInContainer
-    ) {
-      return false
-    }
+    )
+  })
+
+  const workbench = await browser.getWorkbench()
+  await workbench.executeCommand('DVC: Pull')
 
+  return browser.waitUntil(async () => {
     for (const progress of currentProgressBars) {
       if ((await progress.getAttribute('aria-hidden')) !== 'true') {
         return false

From 764640bcd078a5ab3540c2bec723d4ca181114c5 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 14:10:26 +1000
Subject: [PATCH 08/28] explicitly set timeout in before and wait step

---
 extension/src/test/e2e/extension.test.ts |  3 ++-
 extension/src/test/e2e/util.ts           | 17 ++++++++++-------
 extension/src/test/e2e/wdio.conf.ts      |  2 +-
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 1985703be1..fbe5a6bf3d 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -10,7 +10,8 @@ import { PlotsWebview } from './pageObjects/plotsWebview'
 import { delay } from '../../util/time'
 
 suite('DVC Extension For Visual Studio Code', () => {
-  before('should finish loading the extension', async () => {
+  before('should finish loading the extension', async function () {
+    this.timeout(180000)
     await waitForViewContainerToLoad()
     return dismissAllNotifications()
   })
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index 9231697fc8..7982accfc2 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -51,15 +51,18 @@ export const waitForViewContainerToLoad = async () => {
   const workbench = await browser.getWorkbench()
   await workbench.executeCommand('DVC: Pull')
 
-  return browser.waitUntil(async () => {
-    for (const progress of currentProgressBars) {
-      if ((await progress.getAttribute('aria-hidden')) !== 'true') {
-        return false
+  return browser.waitUntil(
+    async () => {
+      for (const progress of currentProgressBars) {
+        if ((await progress.getAttribute('aria-hidden')) !== 'true') {
+          return false
+        }
       }
-    }
 
-    return true
-  })
+      return true
+    },
+    { timeout: 180000 }
+  )
 }
 
 export const closeAllEditors = async () => {
diff --git a/extension/src/test/e2e/wdio.conf.ts b/extension/src/test/e2e/wdio.conf.ts
index 25d2d669b7..550f5dbc24 100644
--- a/extension/src/test/e2e/wdio.conf.ts
+++ b/extension/src/test/e2e/wdio.conf.ts
@@ -61,5 +61,5 @@ export const config: Options.Testrunner = {
   reporters: ['spec'],
   services: ['vscode'],
   specs: ['./src/test/e2e/*.test.ts'],
-  waitforTimeout: 60000
+  waitforTimeout: 10000
 }

From a9bc755385cd7a9333cd3cb3a64cf2dda259f367 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 14:25:38 +1000
Subject: [PATCH 09/28] increase timeout of experiment

---
 extension/src/test/e2e/extension.test.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index fbe5a6bf3d..3a1df10c24 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -67,14 +67,14 @@ suite('DVC Extension For Visual Studio Code', () => {
           const currentRows = await webview.row$$
           return currentRows.length >= initialRows.length + epochs
         },
-        { timeout: 120000 }
+        { timeout: 600000 }
       )
 
       const finalRows = await webview.row$$
 
       expect(finalRows.length).toStrictEqual(initialRows.length + epochs)
       await webview.close()
-    }).timeout(180000)
+    }).timeout(600000)
   })
 
   describe('Plots Webview', () => {

From 2dcfe058dbd7a4fc65aae842392d24d7107de751 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 15:35:38 +1000
Subject: [PATCH 10/28] fix up pulling all data before proceeding

---
 .github/workflows/end-to-end-test.yml    |  4 +--
 extension/src/test/e2e/extension.test.ts |  8 +++---
 extension/src/test/e2e/util.ts           | 31 +++++++++++++++++++-----
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index 32afec5fe2..afabd1ea37 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -32,5 +32,5 @@ jobs:
         with:
           name: screenshots
           path: |
-            test/screenshots
-            test/logs
+            extension/src/test/e2e/screenshots
+            extension/src/test/e2e/logs
diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 3a1df10c24..e6d9922966 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -2,12 +2,12 @@ import { suite, before, describe, it } from 'mocha'
 import {
   closeAllEditors,
   dismissAllNotifications,
+  dvcIsWorking,
   getDVCActivityBarIcon,
   waitForViewContainerToLoad
 } from './util'
 import { ExperimentsWebview } from './pageObjects/experimentsWebview'
 import { PlotsWebview } from './pageObjects/plotsWebview'
-import { delay } from '../../util/time'
 
 suite('DVC Extension For Visual Studio Code', () => {
   before('should finish loading the extension', async function () {
@@ -17,7 +17,7 @@ suite('DVC Extension For Visual Studio Code', () => {
   })
 
   // avoid killing any background process after experiments have finished run
-  after(() => delay(30000))
+  after(() => browser.waitUntil(() => !dvcIsWorking()))
 
   afterEach(() => browser.switchToFrame(null))
 
@@ -67,14 +67,14 @@ suite('DVC Extension For Visual Studio Code', () => {
           const currentRows = await webview.row$$
           return currentRows.length >= initialRows.length + epochs
         },
-        { timeout: 600000 }
+        { timeout: 180000 }
       )
 
       const finalRows = await webview.row$$
 
       expect(finalRows.length).toStrictEqual(initialRows.length + epochs)
       await webview.close()
-    }).timeout(600000)
+    }).timeout(180000)
   })
 
   describe('Plots Webview', () => {
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index 7982accfc2..6fdaff9566 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -1,5 +1,6 @@
 import { ViewControl } from 'wdio-vscode-service'
-import { ElementArray } from 'webdriverio'
+
+const findProgressBars = () => $$('.monaco-progress-container')
 
 export const dismissAllNotifications = () =>
   browser.waitUntil(async () => {
@@ -12,6 +13,16 @@ export const dismissAllNotifications = () =>
     return openNotifications.length === 0
   })
 
+export const dvcIsWorking = async () => {
+  const workbench = await browser.getWorkbench()
+  const statusBar = workbench.getStatusBar()
+  const statusBarItems = await statusBar.getItems()
+  return statusBarItems.some(
+    statusBarItem =>
+      statusBarItem.includes('loading~spin') && statusBarItem.includes('DVC')
+  )
+}
+
 export const getDVCActivityBarIcon = async (): Promise<ViewControl> => {
   const workbench = await browser.getWorkbench()
 
@@ -24,7 +35,7 @@ export const getDVCActivityBarIcon = async (): Promise<ViewControl> => {
 }
 
 export const waitForViewContainerToLoad = async () => {
-  const initialProgressBars = await $$('.monaco-progress-container')
+  const initialProgressBars = await findProgressBars()
   await browser.waitUntil(async () => {
     const dvcIcon = await getDVCActivityBarIcon()
     if (!dvcIcon) {
@@ -36,11 +47,13 @@ export const waitForViewContainerToLoad = async () => {
     return !!view
   })
 
-  let currentProgressBars: ElementArray
-
   await browser.waitUntil(async () => {
+    if (await dvcIsWorking()) {
+      return false
+    }
+
     const numberOfProgressBarsInContainer = 7
-    currentProgressBars = await $$('.monaco-progress-container')
+    const currentProgressBars = await findProgressBars()
 
     return !(
       currentProgressBars.length <
@@ -53,8 +66,14 @@ export const waitForViewContainerToLoad = async () => {
 
   return browser.waitUntil(
     async () => {
+      if (await dvcIsWorking()) {
+        return false
+      }
+
+      const currentProgressBars = await findProgressBars()
+
       for (const progress of currentProgressBars) {
-        if ((await progress.getAttribute('aria-hidden')) !== 'true') {
+        if (await progress.isDisplayed()) {
           return false
         }
       }

From b917f84ab0a5e4a80bb6361091b6df255be71d28 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 15:43:23 +1000
Subject: [PATCH 11/28] fix timeouts of startup

---
 extension/src/test/e2e/extension.test.ts |  2 +-
 extension/src/test/e2e/util.ts           | 25 +++++++++++++-----------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index e6d9922966..41a2c036c1 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -11,7 +11,7 @@ import { PlotsWebview } from './pageObjects/plotsWebview'
 
 suite('DVC Extension For Visual Studio Code', () => {
   before('should finish loading the extension', async function () {
-    this.timeout(180000)
+    this.timeout(240000)
     await waitForViewContainerToLoad()
     return dismissAllNotifications()
   })
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index 6fdaff9566..80b6b48dc1 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -47,19 +47,22 @@ export const waitForViewContainerToLoad = async () => {
     return !!view
   })
 
-  await browser.waitUntil(async () => {
-    if (await dvcIsWorking()) {
-      return false
-    }
+  await browser.waitUntil(
+    async () => {
+      if (await dvcIsWorking()) {
+        return false
+      }
 
-    const numberOfProgressBarsInContainer = 7
-    const currentProgressBars = await findProgressBars()
+      const numberOfProgressBarsInContainer = 7
+      const currentProgressBars = await findProgressBars()
 
-    return !(
-      currentProgressBars.length <
-      initialProgressBars.length + numberOfProgressBarsInContainer
-    )
-  })
+      return !(
+        currentProgressBars.length <
+        initialProgressBars.length + numberOfProgressBarsInContainer
+      )
+    },
+    { timeout: 60000 }
+  )
 
   const workbench = await browser.getWorkbench()
   await workbench.executeCommand('DVC: Pull')

From 86bbe5805efe15b8d4de2edaf0c24f5238e131a3 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 16:01:32 +1000
Subject: [PATCH 12/28] split wait for blocks

---
 extension/src/test/e2e/extension.test.ts |  2 +-
 extension/src/test/e2e/util.ts           | 25 +++++++++++-------------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 41a2c036c1..f63fb56d8f 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -17,7 +17,7 @@ suite('DVC Extension For Visual Studio Code', () => {
   })
 
   // avoid killing any background process after experiments have finished run
-  after(() => browser.waitUntil(() => !dvcIsWorking()))
+  after(() => browser.waitUntil(async () => !(await dvcIsWorking())))
 
   afterEach(() => browser.switchToFrame(null))
 
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index 80b6b48dc1..db5308c3fc 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -47,22 +47,19 @@ export const waitForViewContainerToLoad = async () => {
     return !!view
   })
 
-  await browser.waitUntil(
-    async () => {
-      if (await dvcIsWorking()) {
-        return false
-      }
+  await browser.waitUntil(async () => {
+    const numberOfProgressBarsInContainer = 7
+    const currentProgressBars = await findProgressBars()
 
-      const numberOfProgressBarsInContainer = 7
-      const currentProgressBars = await findProgressBars()
+    return !(
+      currentProgressBars.length <
+      initialProgressBars.length + numberOfProgressBarsInContainer
+    )
+  })
 
-      return !(
-        currentProgressBars.length <
-        initialProgressBars.length + numberOfProgressBarsInContainer
-      )
-    },
-    { timeout: 60000 }
-  )
+  await browser.waitUntil(async () => !(await dvcIsWorking()), {
+    timeout: 60000
+  })
 
   const workbench = await browser.getWorkbench()
   await workbench.executeCommand('DVC: Pull')

From f7a84a4c1bb752abae6767df1111b394ce4a06ae Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 16:07:49 +1000
Subject: [PATCH 13/28] refactor waiting for the cli functions

---
 extension/src/test/e2e/extension.test.ts | 12 ++++++++----
 extension/src/test/e2e/util.ts           | 11 +++++++----
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index f63fb56d8f..0886ae4531 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -2,8 +2,8 @@ import { suite, before, describe, it } from 'mocha'
 import {
   closeAllEditors,
   dismissAllNotifications,
-  dvcIsWorking,
   getDVCActivityBarIcon,
+  waitForDvcToFinish,
   waitForViewContainerToLoad
 } from './util'
 import { ExperimentsWebview } from './pageObjects/experimentsWebview'
@@ -16,10 +16,14 @@ suite('DVC Extension For Visual Studio Code', () => {
     return dismissAllNotifications()
   })
 
-  // avoid killing any background process after experiments have finished run
-  after(() => browser.waitUntil(async () => !(await dvcIsWorking())))
+  after(function () {
+    this.timeout(60000)
+    return waitForDvcToFinish()
+  })
 
-  afterEach(() => browser.switchToFrame(null))
+  afterEach(function () {
+    return browser.switchToFrame(null)
+  })
 
   describe('Activity Bar', () => {
     it('should show the DVC Icon', async () => {
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index db5308c3fc..38e1d8898d 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -13,7 +13,7 @@ export const dismissAllNotifications = () =>
     return openNotifications.length === 0
   })
 
-export const dvcIsWorking = async () => {
+const dvcIsWorking = async () => {
   const workbench = await browser.getWorkbench()
   const statusBar = workbench.getStatusBar()
   const statusBarItems = await statusBar.getItems()
@@ -23,6 +23,11 @@ export const dvcIsWorking = async () => {
   )
 }
 
+export const waitForDvcToFinish = () =>
+  browser.waitUntil(async () => !(await dvcIsWorking()), {
+    timeout: 60000
+  })
+
 export const getDVCActivityBarIcon = async (): Promise<ViewControl> => {
   const workbench = await browser.getWorkbench()
 
@@ -57,9 +62,7 @@ export const waitForViewContainerToLoad = async () => {
     )
   })
 
-  await browser.waitUntil(async () => !(await dvcIsWorking()), {
-    timeout: 60000
-  })
+  await waitForDvcToFinish()
 
   const workbench = await browser.getWorkbench()
   await workbench.executeCommand('DVC: Pull')

From 6a24e454c4d42a139ab393efed847d8750c68826 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 16:25:50 +1000
Subject: [PATCH 14/28] try ubuntu

---
 .github/workflows/end-to-end-test.yml |  2 +-
 demo/dvc.lock                         | 14 +++++++-------
 demo/training_metrics.json            |  8 ++++----
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index afabd1ea37..1c5cde88fc 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -10,7 +10,7 @@ jobs:
   test:
     strategy:
       matrix:
-        os: [macOS-latest]
+        os: [ubuntu-latest]
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v3
diff --git a/demo/dvc.lock b/demo/dvc.lock
index 3bab9240d7..d9f8eb14f3 100644
--- a/demo/dvc.lock
+++ b/demo/dvc.lock
@@ -17,18 +17,18 @@ stages:
         weight_decay: 0
     outs:
     - path: misclassified.jpg
-      md5: 5edd7dd13144d95ef72e270570d9d338
-      size: 32829
+      md5: 97b176d8fcd0590c82f0f7499850c65b
+      size: 32720
     - path: model.pt
-      md5: 38126781764ca9fb04496ca2c2173056
+      md5: c7a5760efd52d3759d8e546ab867f4a6
       size: 439383
     - path: predictions.json
-      md5: 0aa0e279e904572c28c1203ebfa41b29
+      md5: 77955a85635a5273fd80ea4cfeca5822
       size: 310000
     - path: training_metrics
-      md5: 5abc8d18d6b57bd513acf95e55e2755a.dir
-      size: 947
+      md5: 181cb4d5e6171a7f16592d1301c87358.dir
+      size: 950
       nfiles: 2
     - path: training_metrics.json
-      md5: 2a592caac70651d359b2331ff674beec
+      md5: 1d702331f65c5bc23c7552306d2d3861
       size: 69
diff --git a/demo/training_metrics.json b/demo/training_metrics.json
index fefa4a9f5b..80b65680ac 100644
--- a/demo/training_metrics.json
+++ b/demo/training_metrics.json
@@ -1,5 +1,5 @@
 {
-    "step": 14,
-    "loss": 0.9596208930015564,
-    "acc": 0.7735
-}
\ No newline at end of file
+  "step": 14,
+  "loss": 0.9756642580032349,
+  "acc": 0.7659
+}

From 97eb9f2f51c83600eb36b08bf52a3490b8d8d785 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 16:41:02 +1000
Subject: [PATCH 15/28] reduce number of epochs

---
 demo/params.yaml                         | 2 +-
 extension/src/test/e2e/extension.test.ts | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/demo/params.yaml b/demo/params.yaml
index a90973e184..76c95d589b 100644
--- a/demo/params.yaml
+++ b/demo/params.yaml
@@ -1,3 +1,3 @@
 lr: 0.003
 weight_decay: 0
-epochs: 15
+epochs: 5
diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 0886ae4531..0ba4d588e9 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -54,7 +54,7 @@ suite('DVC Extension For Visual Studio Code', () => {
 
     it('should update with a new row for each checkpoint when an experiment is running', async () => {
       const workbench = await browser.getWorkbench()
-      const epochs = 15
+      const epochs = 5
       await workbench.executeCommand('DVC: Reset and Run Experiment')
 
       await webview.open()

From 2a58c6594520c1b400859dcfaa26c5334cfddec7 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 19:24:41 +1000
Subject: [PATCH 16/28] try and debug if rows are being found or not

---
 extension/src/test/e2e/extension.test.ts |  5 +++++
 extension/src/test/e2e/util.ts           | 22 +++++++++++++---------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 0ba4d588e9..cd602827c9 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -69,6 +69,8 @@ suite('DVC Extension For Visual Studio Code', () => {
         async () => {
           await webview.expandAllRows()
           const currentRows = await webview.row$$
+          // eslint-disable-next-line no-console
+          console.error(JSON.stringify(currentRows.length))
           return currentRows.length >= initialRows.length + epochs
         },
         { timeout: 180000 }
@@ -78,6 +80,7 @@ suite('DVC Extension For Visual Studio Code', () => {
 
       expect(finalRows.length).toStrictEqual(initialRows.length + epochs)
       await webview.close()
+      await waitForDvcToFinish()
     }).timeout(180000)
   })
 
@@ -92,6 +95,8 @@ suite('DVC Extension For Visual Studio Code', () => {
       const workbench = await browser.getWorkbench()
       await workbench.executeCommand('DVC: Show Plots')
 
+      await waitForDvcToFinish()
+
       await webview.open()
 
       await browser.waitUntil(async () => {
diff --git a/extension/src/test/e2e/util.ts b/extension/src/test/e2e/util.ts
index 38e1d8898d..d6c66b3ed6 100644
--- a/extension/src/test/e2e/util.ts
+++ b/extension/src/test/e2e/util.ts
@@ -1,9 +1,11 @@
+import { ChainablePromiseArray, ElementArray } from 'webdriverio'
 import { ViewControl } from 'wdio-vscode-service'
 
-const findProgressBars = () => $$('.monaco-progress-container')
+const findProgressBars = (): ChainablePromiseArray<ElementArray> =>
+  $$('.monaco-progress-container')
 
-export const dismissAllNotifications = () =>
-  browser.waitUntil(async () => {
+export const dismissAllNotifications = async (): Promise<void> => {
+  await browser.waitUntil(async () => {
     const workbench = await browser.getWorkbench()
     const notifications = await workbench.getNotifications()
     for (const n of notifications) {
@@ -12,8 +14,9 @@ export const dismissAllNotifications = () =>
     const openNotifications = await workbench.getNotifications()
     return openNotifications.length === 0
   })
+}
 
-const dvcIsWorking = async () => {
+const dvcIsWorking = async (): Promise<boolean> => {
   const workbench = await browser.getWorkbench()
   const statusBar = workbench.getStatusBar()
   const statusBarItems = await statusBar.getItems()
@@ -23,10 +26,11 @@ const dvcIsWorking = async () => {
   )
 }
 
-export const waitForDvcToFinish = () =>
-  browser.waitUntil(async () => !(await dvcIsWorking()), {
+export const waitForDvcToFinish = async (): Promise<void> => {
+  await browser.waitUntil(async () => !(await dvcIsWorking()), {
     timeout: 60000
   })
+}
 
 export const getDVCActivityBarIcon = async (): Promise<ViewControl> => {
   const workbench = await browser.getWorkbench()
@@ -39,7 +43,7 @@ export const getDVCActivityBarIcon = async (): Promise<ViewControl> => {
   return activityBar.getViewControl('DVC') as Promise<ViewControl>
 }
 
-export const waitForViewContainerToLoad = async () => {
+export const waitForViewContainerToLoad = async (): Promise<void> => {
   const initialProgressBars = await findProgressBars()
   await browser.waitUntil(async () => {
     const dvcIcon = await getDVCActivityBarIcon()
@@ -67,7 +71,7 @@ export const waitForViewContainerToLoad = async () => {
   const workbench = await browser.getWorkbench()
   await workbench.executeCommand('DVC: Pull')
 
-  return browser.waitUntil(
+  await browser.waitUntil(
     async () => {
       if (await dvcIsWorking()) {
         return false
@@ -87,7 +91,7 @@ export const waitForViewContainerToLoad = async () => {
   )
 }
 
-export const closeAllEditors = async () => {
+export const closeAllEditors = async (): Promise<void> => {
   const workbench = await browser.getWorkbench()
   const editorView = workbench.getEditorView()
   return editorView.closeAllEditors()

From adf3176c7dd927b06d6c36977bff969f4e8bf71f Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 19:33:35 +1000
Subject: [PATCH 17/28] check on expand buttons

---
 extension/src/test/e2e/pageObjects/experimentsWebview.ts | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/extension/src/test/e2e/pageObjects/experimentsWebview.ts b/extension/src/test/e2e/pageObjects/experimentsWebview.ts
index c8000f329f..a5986cb852 100644
--- a/extension/src/test/e2e/pageObjects/experimentsWebview.ts
+++ b/extension/src/test/e2e/pageObjects/experimentsWebview.ts
@@ -10,8 +10,10 @@ export interface ExperimentsWebview
 export class ExperimentsWebview extends BaseWebview {
   public async expandAllRows() {
     const expandRowButtons = await this.expandRowButton$$
+    // eslint-disable-next-line no-console
+    console.error('expand buttons', JSON.stringify(expandRowButtons.length))
     for (const button of expandRowButtons) {
-      button.click()
+      await button.click()
     }
     return expandRowButtons.length === 0
   }

From 6010bf0a70b9f08cd1c35d7934b5e865202bed05 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 19:45:10 +1000
Subject: [PATCH 18/28] check on contract buttons

---
 extension/src/test/e2e/pageObjects/experimentsWebview.ts | 6 ++++++
 extension/src/test/e2e/pageObjects/locators.ts           | 1 +
 2 files changed, 7 insertions(+)

diff --git a/extension/src/test/e2e/pageObjects/experimentsWebview.ts b/extension/src/test/e2e/pageObjects/experimentsWebview.ts
index a5986cb852..6e9302e909 100644
--- a/extension/src/test/e2e/pageObjects/experimentsWebview.ts
+++ b/extension/src/test/e2e/pageObjects/experimentsWebview.ts
@@ -12,6 +12,12 @@ export class ExperimentsWebview extends BaseWebview {
     const expandRowButtons = await this.expandRowButton$$
     // eslint-disable-next-line no-console
     console.error('expand buttons', JSON.stringify(expandRowButtons.length))
+    // eslint-disable-next-line no-console
+    console.error(
+      'contract buttons',
+      // eslint-disable-next-line unicorn/no-await-expression-member
+      JSON.stringify((await this.contractRowButton$$).length)
+    )
     for (const button of expandRowButtons) {
       await button.click()
     }
diff --git a/extension/src/test/e2e/pageObjects/locators.ts b/extension/src/test/e2e/pageObjects/locators.ts
index 9c3cec4553..f9db448393 100644
--- a/extension/src/test/e2e/pageObjects/locators.ts
+++ b/extension/src/test/e2e/pageObjects/locators.ts
@@ -5,6 +5,7 @@ export const webview = {
 
 export const experiments = {
   ...webview,
+  contractRowButton: 'button[title="Contract Row"]',
   expandRowButton: 'button[title="Expand Row"]',
   row: '[role=row]',
   table: '[role=tree]'

From c508884ff76e52f89b19f0f33a224c7f27f41144 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 20:12:10 +1000
Subject: [PATCH 19/28] take screenshots

---
 extension/src/test/e2e/extension.test.ts | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index cd602827c9..1df3c8372a 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -1,3 +1,4 @@
+import { join } from 'path'
 import { suite, before, describe, it } from 'mocha'
 import {
   closeAllEditors,
@@ -65,15 +66,22 @@ suite('DVC Extension For Visual Studio Code', () => {
 
       expect(initialRows.length).toBeGreaterThanOrEqual(4)
 
+      const screenshotDir = join(__dirname, 'screenshots')
+
       await browser.waitUntil(
         async () => {
           await webview.expandAllRows()
           const currentRows = await webview.row$$
           // eslint-disable-next-line no-console
           console.error(JSON.stringify(currentRows.length))
+
+          await browser.saveScreenshot(
+            join(screenshotDir, `stuck - ${Date.now()}.png`)
+          )
+
           return currentRows.length >= initialRows.length + epochs
         },
-        { timeout: 180000 }
+        { interval: 30000, timeout: 180000 }
       )
 
       const finalRows = await webview.row$$

From ecc979498d76059281b8371c2fe6be439119b06d Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 20:26:27 +1000
Subject: [PATCH 20/28] checkout the entire git history

---
 .github/workflows/end-to-end-test.yml | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index 1c5cde88fc..b9c38d94a0 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -12,14 +12,24 @@ jobs:
       matrix:
         os: [ubuntu-latest]
     runs-on: ${{ matrix.os }}
+
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-node@v3
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - name: Setup Node.js environment
+        uses: actions/setup-node@v3
         with:
-          node-version: 16
+          node-version: '16'
+          cache: yarn
+
+      - name: Setup Python environment
       - uses: actions/setup-python@v4
         with:
           python-version: '3.9'
+
       - run: yarn run install-frozen-lockfile
 
       - name: Test E2E

From d6046626ca6115d97381b0fbddeccf8499b8180b Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 20:32:05 +1000
Subject: [PATCH 21/28] fix action

---
 .github/workflows/end-to-end-test.yml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index b9c38d94a0..f766d8ea94 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -24,14 +24,13 @@ jobs:
         with:
           node-version: '16'
           cache: yarn
+      - run: yarn run install-frozen-lockfile
 
       - name: Setup Python environment
-      - uses: actions/setup-python@v4
+        uses: actions/setup-python@v4
         with:
           python-version: '3.9'
 
-      - run: yarn run install-frozen-lockfile
-
       - name: Test E2E
         uses: GabrielBB/xvfb-action@v1
         with:

From b244499fbd1faf24045a59ea924eab6bfc31caec Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 20:39:21 +1000
Subject: [PATCH 22/28] remove all debugging code

---
 demo/params.yaml                                     |  2 +-
 extension/src/test/e2e/extension.test.ts             | 12 +-----------
 .../src/test/e2e/pageObjects/experimentsWebview.ts   |  8 --------
 extension/src/test/e2e/pageObjects/locators.ts       |  1 -
 4 files changed, 2 insertions(+), 21 deletions(-)

diff --git a/demo/params.yaml b/demo/params.yaml
index 76c95d589b..a90973e184 100644
--- a/demo/params.yaml
+++ b/demo/params.yaml
@@ -1,3 +1,3 @@
 lr: 0.003
 weight_decay: 0
-epochs: 5
+epochs: 15
diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 1df3c8372a..6c11689471 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -1,4 +1,3 @@
-import { join } from 'path'
 import { suite, before, describe, it } from 'mocha'
 import {
   closeAllEditors,
@@ -55,7 +54,7 @@ suite('DVC Extension For Visual Studio Code', () => {
 
     it('should update with a new row for each checkpoint when an experiment is running', async () => {
       const workbench = await browser.getWorkbench()
-      const epochs = 5
+      const epochs = 15
       await workbench.executeCommand('DVC: Reset and Run Experiment')
 
       await webview.open()
@@ -66,19 +65,10 @@ suite('DVC Extension For Visual Studio Code', () => {
 
       expect(initialRows.length).toBeGreaterThanOrEqual(4)
 
-      const screenshotDir = join(__dirname, 'screenshots')
-
       await browser.waitUntil(
         async () => {
           await webview.expandAllRows()
           const currentRows = await webview.row$$
-          // eslint-disable-next-line no-console
-          console.error(JSON.stringify(currentRows.length))
-
-          await browser.saveScreenshot(
-            join(screenshotDir, `stuck - ${Date.now()}.png`)
-          )
-
           return currentRows.length >= initialRows.length + epochs
         },
         { interval: 30000, timeout: 180000 }
diff --git a/extension/src/test/e2e/pageObjects/experimentsWebview.ts b/extension/src/test/e2e/pageObjects/experimentsWebview.ts
index 6e9302e909..c2e7fad50c 100644
--- a/extension/src/test/e2e/pageObjects/experimentsWebview.ts
+++ b/extension/src/test/e2e/pageObjects/experimentsWebview.ts
@@ -10,14 +10,6 @@ export interface ExperimentsWebview
 export class ExperimentsWebview extends BaseWebview {
   public async expandAllRows() {
     const expandRowButtons = await this.expandRowButton$$
-    // eslint-disable-next-line no-console
-    console.error('expand buttons', JSON.stringify(expandRowButtons.length))
-    // eslint-disable-next-line no-console
-    console.error(
-      'contract buttons',
-      // eslint-disable-next-line unicorn/no-await-expression-member
-      JSON.stringify((await this.contractRowButton$$).length)
-    )
     for (const button of expandRowButtons) {
       await button.click()
     }
diff --git a/extension/src/test/e2e/pageObjects/locators.ts b/extension/src/test/e2e/pageObjects/locators.ts
index f9db448393..9c3cec4553 100644
--- a/extension/src/test/e2e/pageObjects/locators.ts
+++ b/extension/src/test/e2e/pageObjects/locators.ts
@@ -5,7 +5,6 @@ export const webview = {
 
 export const experiments = {
   ...webview,
-  contractRowButton: 'button[title="Contract Row"]',
   expandRowButton: 'button[title="Expand Row"]',
   row: '[role=row]',
   table: '[role=tree]'

From 7e5cf62e8d42d5a01552325691bf6befe8b27d01 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 8 Jul 2022 20:51:15 +1000
Subject: [PATCH 23/28] revert dvc files

---
 demo/dvc.lock              | 14 +++++++-------
 demo/training_metrics.json |  4 ++--
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/demo/dvc.lock b/demo/dvc.lock
index d9f8eb14f3..3bab9240d7 100644
--- a/demo/dvc.lock
+++ b/demo/dvc.lock
@@ -17,18 +17,18 @@ stages:
         weight_decay: 0
     outs:
     - path: misclassified.jpg
-      md5: 97b176d8fcd0590c82f0f7499850c65b
-      size: 32720
+      md5: 5edd7dd13144d95ef72e270570d9d338
+      size: 32829
     - path: model.pt
-      md5: c7a5760efd52d3759d8e546ab867f4a6
+      md5: 38126781764ca9fb04496ca2c2173056
       size: 439383
     - path: predictions.json
-      md5: 77955a85635a5273fd80ea4cfeca5822
+      md5: 0aa0e279e904572c28c1203ebfa41b29
       size: 310000
     - path: training_metrics
-      md5: 181cb4d5e6171a7f16592d1301c87358.dir
-      size: 950
+      md5: 5abc8d18d6b57bd513acf95e55e2755a.dir
+      size: 947
       nfiles: 2
     - path: training_metrics.json
-      md5: 1d702331f65c5bc23c7552306d2d3861
+      md5: 2a592caac70651d359b2331ff674beec
       size: 69
diff --git a/demo/training_metrics.json b/demo/training_metrics.json
index 80b65680ac..486cd24a07 100644
--- a/demo/training_metrics.json
+++ b/demo/training_metrics.json
@@ -1,5 +1,5 @@
 {
   "step": 14,
-  "loss": 0.9756642580032349,
-  "acc": 0.7659
+  "loss": 0.9596208930015564,
+  "acc": 0.7735
 }

From b77a11ac655ec285f2fb9369a4bb87d3e18a065e Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 12 Jul 2022 08:37:19 +1000
Subject: [PATCH 24/28] add end to end into continuous integration

---
 .github/workflows/continuous-integration.yml | 18 ++++++++++++++++++
 .github/workflows/end-to-end-test.yml        |  9 ++++-----
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml
index 8d210180ff..67ef7efa3d 100644
--- a/.github/workflows/continuous-integration.yml
+++ b/.github/workflows/continuous-integration.yml
@@ -19,6 +19,11 @@ jobs:
         with:
           fetch-depth: 0
 
+      - name: Setup Python environment
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+
       - name: Setup Node.js environment
         uses: actions/setup-node@v3
         with:
@@ -55,3 +60,16 @@ jobs:
           storybookBuildDir: 'webview/storybook-static'
           exitOnceUploaded: true
           autoAcceptChanges: 'main'
+
+      - name: Test E2E
+        uses: GabrielBB/xvfb-action@v1
+        with:
+          run: yarn test:e2e
+          options: '-screen 0 1600x1200x24'
+      - uses: actions/upload-artifact@v3
+        if: failure()
+        with:
+          name: screenshots
+          path: |
+            extension/src/test/e2e/screenshots
+            extension/src/test/e2e/logs
diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index f766d8ea94..2758b82cbf 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -1,16 +1,15 @@
 name: End-To-End Test
 on:
-  push:
-    branches:
-      - main
-  pull_request:
+  push: # remove before merge
   workflow_dispatch:
+  schedule:
+    - cron: '0 0 * * *'
 
 jobs:
   test:
     strategy:
       matrix:
-        os: [ubuntu-latest]
+        os: [ubuntu-latest, macos-latest, windows-latest]
     runs-on: ${{ matrix.os }}
 
     steps:

From 1d2303baf07fc578c0af0eb6fcc27829eb3957a1 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 12 Jul 2022 08:43:21 +1000
Subject: [PATCH 25/28] use yarn install for windows

---
 .github/workflows/end-to-end-test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
index 2758b82cbf..e36936f1ef 100644
--- a/.github/workflows/end-to-end-test.yml
+++ b/.github/workflows/end-to-end-test.yml
@@ -23,7 +23,7 @@ jobs:
         with:
           node-version: '16'
           cache: yarn
-      - run: yarn run install-frozen-lockfile
+      - run: yarn install
 
       - name: Setup Python environment
         uses: actions/setup-python@v4

From 5e965cf8bd6264ceeb81241fff08803fcc38eafc Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 12 Jul 2022 08:47:33 +1000
Subject: [PATCH 26/28] remove macos from cross-platform tests

---
 .github/workflows/cross-platform-test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cross-platform-test.yml b/.github/workflows/cross-platform-test.yml
index 23227f0fa0..c305cf3a19 100644
--- a/.github/workflows/cross-platform-test.yml
+++ b/.github/workflows/cross-platform-test.yml
@@ -15,7 +15,7 @@ jobs:
     timeout-minutes: 20
     strategy:
       matrix:
-        os: [macOS-latest, windows-latest]
+        os: [windows-latest]
     steps:
       - name: Checkout
         uses: actions/checkout@v3

From edd489eb7d05865b116bcc90b284cf3c74fa1d12 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 12 Jul 2022 09:23:08 +1000
Subject: [PATCH 27/28] improve test reliability

---
 extension/src/test/e2e/extension.test.ts      | 19 ++++++++++++-------
 .../src/test/e2e/pageObjects/baseWebview.ts   |  4 ++--
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/extension/src/test/e2e/extension.test.ts b/extension/src/test/e2e/extension.test.ts
index 6c11689471..585f225b02 100644
--- a/extension/src/test/e2e/extension.test.ts
+++ b/extension/src/test/e2e/extension.test.ts
@@ -40,7 +40,7 @@ suite('DVC Extension For Visual Studio Code', () => {
 
       await workbench.executeCommand('DVC: Show Experiments')
 
-      await webview.open()
+      await webview.focus()
 
       await browser.waitUntil(async () => {
         const table = await webview.table$
@@ -49,7 +49,7 @@ suite('DVC Extension For Visual Studio Code', () => {
 
       expect(await webview.table$$).toHaveLength(1)
 
-      await webview.close()
+      await webview.unfocus()
     })
 
     it('should update with a new row for each checkpoint when an experiment is running', async () => {
@@ -57,7 +57,7 @@ suite('DVC Extension For Visual Studio Code', () => {
       const epochs = 15
       await workbench.executeCommand('DVC: Reset and Run Experiment')
 
-      await webview.open()
+      await webview.focus()
 
       await browser.waitUntil(() => webview.expandAllRows())
 
@@ -71,14 +71,19 @@ suite('DVC Extension For Visual Studio Code', () => {
           const currentRows = await webview.row$$
           return currentRows.length >= initialRows.length + epochs
         },
-        { interval: 30000, timeout: 180000 }
+        { interval: 5000, timeout: 180000 }
       )
 
+      await webview.unfocus()
+      await waitForDvcToFinish()
+      await webview.focus()
+
       const finalRows = await webview.row$$
 
       expect(finalRows.length).toStrictEqual(initialRows.length + epochs)
-      await webview.close()
+      await webview.unfocus()
       await waitForDvcToFinish()
+      await workbench.executeCommand('Terminal: Kill All Terminals')
     }).timeout(180000)
   })
 
@@ -95,7 +100,7 @@ suite('DVC Extension For Visual Studio Code', () => {
 
       await waitForDvcToFinish()
 
-      await webview.open()
+      await webview.focus()
 
       await browser.waitUntil(async () => {
         return (await webview.vegaVisualization$$.length) === 6
@@ -109,7 +114,7 @@ suite('DVC Extension For Visual Studio Code', () => {
         expect(plotNotEmpty).toBe(true)
       }
 
-      await webview.close()
+      await webview.unfocus()
     })
   })
 })
diff --git a/extension/src/test/e2e/pageObjects/baseWebview.ts b/extension/src/test/e2e/pageObjects/baseWebview.ts
index 6e8343e3f0..92b1b7b348 100644
--- a/extension/src/test/e2e/pageObjects/baseWebview.ts
+++ b/extension/src/test/e2e/pageObjects/baseWebview.ts
@@ -17,7 +17,7 @@ export class BaseWebview extends BasePage<
     this.locatorKey = locatorKey
   }
 
-  public async open() {
+  public async focus() {
     const webviewContainer = await this.outerFrame$
 
     await this.outerFrame$.waitForDisplayed()
@@ -31,7 +31,7 @@ export class BaseWebview extends BasePage<
     return browser.switchToFrame(webviewInner)
   }
 
-  public async close() {
+  public async unfocus() {
     await browser.switchToFrame(null)
     await browser.switchToFrame(null)
   }

From cf52e17c9d8a271cf65c488c6865d3a50d66c646 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 12 Jul 2022 10:00:21 +1000
Subject: [PATCH 28/28] remove end to end workflow

---
 .github/workflows/end-to-end-test.yml | 44 ---------------------------
 1 file changed, 44 deletions(-)
 delete mode 100644 .github/workflows/end-to-end-test.yml

diff --git a/.github/workflows/end-to-end-test.yml b/.github/workflows/end-to-end-test.yml
deleted file mode 100644
index e36936f1ef..0000000000
--- a/.github/workflows/end-to-end-test.yml
+++ /dev/null
@@ -1,44 +0,0 @@
-name: End-To-End Test
-on:
-  push: # remove before merge
-  workflow_dispatch:
-  schedule:
-    - cron: '0 0 * * *'
-
-jobs:
-  test:
-    strategy:
-      matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
-    runs-on: ${{ matrix.os }}
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
-
-      - name: Setup Node.js environment
-        uses: actions/setup-node@v3
-        with:
-          node-version: '16'
-          cache: yarn
-      - run: yarn install
-
-      - name: Setup Python environment
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.9'
-
-      - name: Test E2E
-        uses: GabrielBB/xvfb-action@v1
-        with:
-          run: yarn test:e2e
-          options: '-screen 0 1600x1200x24'
-      - uses: actions/upload-artifact@v3
-        if: failure()
-        with:
-          name: screenshots
-          path: |
-            extension/src/test/e2e/screenshots
-            extension/src/test/e2e/logs