Increase timeouts (#311)

Co-authored-by: Fabian Jakobs <fabian.jakobs@databricks.com>

Author: kartikgupta-db

packages/databricks-vscode/src/test/e2e/configure.e2e.ts

@@ -3,7 +3,7 @@ import path from "node:path";
 import * as fs from "fs/promises";
 import {
     getViewSection,
-    waitForPythonExtension,
+    waitForPythonExtensionWithRetry,
     waitForTreeItems,
 } from "./utils";
 import {
@@ -21,7 +21,7 @@ describe("Configure Databricks Extension", async function () {
     let host: string;
     let workbench: Workbench;
 
-    this.timeout(3 * 60 * 1000);
+    this.timeout(5 * 60 * 1000);
 
     before(async function () {
         assert(
@@ -41,7 +41,7 @@ describe("Configure Databricks Extension", async function () {
         host = process.env.DATABRICKS_HOST;
 
         workbench = await browser.getWorkbench();
-        await waitForPythonExtension();
+        await waitForPythonExtensionWithRetry(1.5 * 60 * 1000, 2);
     });
 
     it("should open VSCode", async function () {

packages/databricks-vscode/src/test/e2e/run_job_on_cluster.e2e.ts

@@ -3,15 +3,15 @@ import * as fs from "fs/promises";
 import assert from "node:assert";
 import {
     getViewSection,
-    getViewSubSection,
-    waitForPythonExtension,
-    waitForTreeItems,
+    startSyncIfStopped,
+    waitForPythonExtensionWithRetry,
+    waitForSyncComplete,
 } from "./utils";
-import {sleep, TreeItem} from "wdio-vscode-service";
+import {sleep} from "wdio-vscode-service";
 
-describe("Run python on cluster", async function () {
+describe("Run as workflow on cluster", async function () {
     let projectDir: string;
-    this.timeout(2 * 60 * 1000);
+    this.timeout(5 * 60 * 1000);
 
     before(async () => {
         assert(process.env.TEST_DEFAULT_CLUSTER_ID);
@@ -42,59 +42,15 @@ describe("Run python on cluster", async function () {
             ].join("\n")
         );
 
-        const section = await getViewSection("CONFIGURATION");
-        assert(section);
-        await waitForTreeItems(section);
-        await waitForPythonExtension();
+        await waitForPythonExtensionWithRetry(1.5 * 60 * 1000, 2);
+
+        await (await getViewSection("CLUSTERS"))?.collapse();
     });
 
     beforeEach(async () => {
-        const section = await getViewSection("CLUSTERS");
-        await section?.collapse();
-
-        const repoConfigItem = await getViewSubSection("CONFIGURATION", "Repo");
-        assert(repoConfigItem);
+        await startSyncIfStopped();
 
-        let status: TreeItem | undefined = undefined;
-        for (const i of await repoConfigItem.getChildren()) {
-            if ((await i.getLabel()).includes("State:")) {
-                status = i;
-                break;
-            }
-        }
-        assert(status);
-        if ((await status.getDescription())?.includes("STOPPED")) {
-            const buttons = await repoConfigItem.getActionButtons();
-            await buttons[0].elem.click();
-        }
-
-        await browser.waitUntil(
-            async () => {
-                const repoConfigItem = await getViewSubSection(
-                    "CONFIGURATION",
-                    "Repo"
-                );
-                assert(repoConfigItem);
-
-                status = undefined;
-                for (const i of await repoConfigItem.getChildren()) {
-                    if ((await i.getLabel()).includes("State:")) {
-                        status = i;
-                        break;
-                    }
-                }
-                assert(status);
-                const description = await status?.getDescription();
-                return (
-                    description !== undefined &&
-                    description.includes("WATCHING_FOR_CHANGES")
-                );
-            },
-            {
-                timeout: 20000,
-                timeoutMsg: "Couldn't finish sync in 20s",
-            }
-        );
+        await waitForSyncComplete();
     });
 
     it("should run a python notebook as a job on a cluster", async () => {
@@ -170,7 +126,7 @@ describe("Run python on cluster", async function () {
         await sleep(200);
         await input.setText("file.py");
         await input.confirm();
-        await sleep(500);
+        await sleep(2000);
 
         // run file
         await workbench.executeQuickPick(

packages/databricks-vscode/src/test/e2e/run_on_cluster.e2e.ts

@@ -3,15 +3,16 @@ import * as fs from "fs/promises";
 import assert from "node:assert";
 import {
     getViewSection,
-    getViewSubSection,
-    waitForPythonExtension,
+    startSyncIfStopped,
+    waitForPythonExtensionWithRetry,
+    waitForSyncComplete,
     waitForTreeItems,
 } from "./utils";
 import {sleep} from "wdio-vscode-service";
 
 describe("Run python on cluster", async function () {
     let projectDir: string;
-    this.timeout(3 * 60 * 1000);
+    this.timeout(5 * 60 * 1000);
 
     before(async () => {
         assert(process.env.DATABRICKS_HOST);
@@ -36,7 +37,7 @@ describe("Run python on cluster", async function () {
             path.join(projectDir, "hello.py"),
             `spark.sql('SELECT "hello world"').show()`
         );
-        await waitForPythonExtension();
+        await waitForPythonExtensionWithRetry(1.5 * 60 * 1000, 2);
     });
 
     it("should connect to Databricks", async () => {
@@ -49,22 +50,10 @@ describe("Run python on cluster", async function () {
         const section = await getViewSection("CLUSTERS");
         await section?.collapse();
 
-        const repoConfigItem = await getViewSubSection("CONFIGURATION", "Repo");
-        assert(repoConfigItem);
-        const buttons = await repoConfigItem.getActionButtons();
-        await buttons[0].elem.click();
+        await startSyncIfStopped();
 
         // wait for sync to finish
-        const workbench = await driver.getWorkbench();
-        const terminalView = await workbench.getBottomBar().openTerminalView();
-
-        while (true) {
-            await sleep(500);
-            const text = await terminalView.getText();
-            if (text.includes("Sync Complete")) {
-                break;
-            }
-        }
+        await waitForSyncComplete();
     });
 
     it("should run a python file on a cluster", async () => {

packages/databricks-vscode/src/test/e2e/utils.ts

@@ -82,23 +82,34 @@ export async function waitForTreeItems(
     }
 }
 
-export async function waitForPythonExtension() {
+export async function waitForPythonExtension(timeoutMs: number) {
     const section = await getViewSection("CONFIGURATION");
     assert(section);
-    const welcome = await section.findWelcomeContent();
-    assert(welcome);
-    sleep(1000);
     const workbench = await browser.getWorkbench();
-    const notifs = await workbench.getNotifications();
-    for (const n of notifs) {
-        if (
-            (await n.getActions()).find(
-                (btn) => btn.getTitle() === "Install and Reload"
-            ) !== undefined
-        ) {
-            await n.takeAction("Install and Reload");
+
+    await browser.waitUntil(
+        async () => {
+            const notifs = await workbench.getNotifications();
+            let found = false;
+            for (const n of notifs) {
+                if (
+                    (await n.getActions()).find(
+                        (btn) => btn.getTitle() === "Install and Reload"
+                    ) !== undefined
+                ) {
+                    await n.takeAction("Install and Reload");
+                    found = true;
+                }
+            }
+            return found;
+        },
+        {
+            timeout: 5 * 1000,
+            interval: 500,
+            timeoutMsg:
+                "Can't find notification to install ms-python extension",
         }
-    }
+    );
 
     await browser.waitUntil(
         async () =>
@@ -108,16 +119,115 @@ export async function waitForPythonExtension() {
                 )?.getTitle()
             )?.includes("README.quickstart.md") === true,
         {
-            timeout: 120000,
+            timeout: timeoutMs,
             timeoutMsg:
                 "Timeout when installing python extension and reloading",
         }
     );
 
-    sleep(500);
+    await sleep(500);
     try {
+        const notifs = await workbench.getNotifications();
         for (const n of notifs) {
             await n.dismiss();
         }
     } catch {}
 }
+
+export async function waitForPythonExtensionWithRetry(
+    timeoutMs: number,
+    retryCount: number
+) {
+    for (let i = 0; i < retryCount; i++) {
+        try {
+            const section = await getViewSection("CONFIGURATION");
+            assert(section);
+            await waitForPythonExtension(timeoutMs);
+        } catch (e) {
+            // eslint-disable-next-line no-console
+            console.error(e);
+            const wb = await browser.getWorkbench();
+            await wb.executeCommand("Developer: Reload Window");
+            if (i === retryCount - 1) {
+                throw e;
+            }
+            continue;
+        }
+        break;
+    }
+}
+
+export async function waitForSyncComplete() {
+    await browser.waitUntil(
+        async () => {
+            const repoConfigItem = await getViewSubSection(
+                "CONFIGURATION",
+                "Repo"
+            );
+            if (repoConfigItem === undefined) {
+                return false;
+            }
+            await repoConfigItem.expand();
+
+            let status: TreeItem | undefined = undefined;
+            for (const i of await repoConfigItem.getChildren()) {
+                if ((await i.getLabel()).includes("State:")) {
+                    status = i;
+                    break;
+                }
+            }
+            if (status === undefined) {
+                return false;
+            }
+
+            const description = await status?.getDescription();
+            return (
+                description !== undefined &&
+                description.includes("WATCHING_FOR_CHANGES")
+            );
+        },
+        {
+            timeout: 60000,
+            interval: 20000,
+            timeoutMsg: "Couldn't finish sync in 1m",
+        }
+    );
+}
+
+export async function startSyncIfStopped() {
+    browser.waitUntil(
+        async () => {
+            const repoConfigItem = await getViewSubSection(
+                "CONFIGURATION",
+                "Repo"
+            );
+            if (repoConfigItem === undefined) {
+                return false;
+            }
+            await repoConfigItem.expand();
+
+            let status: TreeItem | undefined = undefined;
+            for (const i of await repoConfigItem.getChildren()) {
+                if ((await i.getLabel()).includes("State:")) {
+                    status = i;
+                    break;
+                }
+            }
+            if (status === undefined) {
+                return false;
+            }
+
+            if ((await status.getDescription())?.includes("STOPPED")) {
+                const buttons = await repoConfigItem.getActionButtons();
+                if (buttons.length === 0) {
+                    return false;
+                }
+                await buttons[0].elem.click();
+            }
+            return true;
+        },
+        {
+            timeout: 20000,
+        }
+    );
+}