Update the CLI to v0.231.0 (#1407)

## Changes
Update the CLI to v0.231.0

Also improve stability of unit and integ tests

## Tests
<!-- How is this tested? -->

Author: ilia-db

packages/databricks-vscode/package.json

@@ -907,7 +907,7 @@
         "useYarn": false
     },
     "cli": {
-        "version": "0.230.0"
+        "version": "0.231.0"
     },
     "scripts": {
         "vscode:prepublish": "rm -rf out && yarn run package:compile && yarn run package:wrappers:write && yarn run package:jupyter-init-script:write && yarn run package:copy-webview-toolkit && yarn run generate-telemetry",

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

@@ -3,6 +3,7 @@ import * as fs from "fs/promises";
 import assert from "node:assert";
 import {
     dismissNotifications,
+    executeCommandWhenAvailable,
     openFile,
     waitForLogin,
     waitForWorkflowWebview,
@@ -43,14 +44,11 @@ describe("Run files", async function () {
 
     beforeEach(async () => {
         await openFile("hello.py");
-        // We enable the run command when the active editor is changed.
-        // We wait here to avoid race conditions between the extension logic enabling this command and tests executing it.
-        await sleep(1000);
     });
 
     it("should run a python file on a cluster", async () => {
         const workbench = await driver.getWorkbench();
-        await workbench.executeQuickPick("Databricks: Upload and Run File");
+        await executeCommandWhenAvailable("Databricks: Upload and Run File");
 
         const debugOutput = await workbench
             .getBottomBar()
@@ -67,7 +65,7 @@ describe("Run files", async function () {
 
     it("should cancel a run during deployment", async () => {
         const workbench = await driver.getWorkbench();
-        await workbench.executeQuickPick("Databricks: Upload and Run File");
+        await executeCommandWhenAvailable("Databricks: Upload and Run File");
         await browser.waitUntil(async () => {
             const notifications = await workbench.getNotifications();
             for (const notification of notifications) {

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

@@ -1,9 +1,9 @@
 import path from "node:path";
 import * as fs from "fs/promises";
 import assert from "node:assert";
-import {sleep} from "wdio-vscode-service";
 import {
     dismissNotifications,
+    executeCommandWhenAvailable,
     openFile,
     waitForLogin,
     waitForWorkflowWebview,
@@ -74,21 +74,13 @@ describe("Run notebooks", async function () {
 
     it("should run a notebook.py file as a workflow", async () => {
         await openFile("notebook.py");
-        // We enable the run command when the active editor is changed.
-        // We wait here to avoid race conditions between the extension logic enabling this command and tests executing it.
-        await sleep(1000);
-        const workbench = await driver.getWorkbench();
-        await workbench.executeQuickPick("Databricks: Run File as Workflow");
+        await executeCommandWhenAvailable("Databricks: Run File as Workflow");
         await waitForWorkflowWebview("a/b c");
     });
 
     it("should run a notebook.ipynb file as a workflow", async () => {
         await openFile("notebook.ipynb");
-        // We enable the run command when the active editor is changed.
-        // We wait here to avoid race conditions between the extension logic enabling this command and tests executing it.
-        await sleep(1000);
-        const workbench = await driver.getWorkbench();
-        await workbench.executeQuickPick("Databricks: Run File as Workflow");
+        await executeCommandWhenAvailable("Databricks: Run File as Workflow");
         await waitForWorkflowWebview("hello world");
     });
 });

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

@@ -303,3 +303,15 @@ export async function openFile(fileName: string) {
         return title.includes(fileName);
     });
 }
+
+export async function executeCommandWhenAvailable(command: string) {
+    const workbench = await driver.getWorkbench();
+    return browser.waitUntil(async () => {
+        try {
+            await workbench.executeQuickPick(command);
+            return true;
+        } catch (e) {
+            return false;
+        }
+    });
+}

packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.test.ts

@@ -236,7 +236,7 @@ describe(__filename, async () => {
                     "dbutils.library.restartPython()",
                     "print('hello2')",
                 ];
-                writeFile(
+                await writeFile(
                     localFilePath.path,
                     comment.concat(data).join("\n"),
                     "utf-8"
@@ -443,7 +443,7 @@ describe(__filename, async () => {
                         },
                     ],
                 };
-                writeFile(
+                await writeFile(
                     localFilePath.path,
                     JSON.stringify(originalData),
                     "utf-8"