Add e2e tests for running .py notebooks with dbconnect (#1626)

## Changes

Tests + fix init-scripts cleanup login on windows (seems like it was
always broken there).

Partially addresses this bug #1624

## Tests

Tests

Author: ilia-db

packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts

@@ -172,7 +172,10 @@ export class NotebookInitScriptManager implements Disposable {
 
         const sourceFiles = await this.sourceFiles;
         for (const file of await glob(
-            path.join(startupDir, "00-databricks-init-*.py")
+            path.join(startupDir, "00-databricks-init-*.py"),
+            // path.join will use "\" separator on windows, wbut glob always expects "/", unless we specify windowsPathsNoEscape.
+            // This means we can't use "\" to escape special chracters in the glob pattern, but we don't use any.
+            {windowsPathsNoEscape: true}
         )) {
             if (!sourceFiles.includes(path.basename(file))) {
                 await rm(file);

packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts

@@ -15,6 +15,22 @@ import {
     writeRootBundleConfig,
 } from "./utils/dabsFixtures.ts";
 
+async function checkOutputFile(path: string, expectedContent: string) {
+    await browser.waitUntil(
+        async () => {
+            const fileContent = await fs.readFile(path, "utf-8");
+            console.log(`"${path}" contents: `, fileContent);
+            return fileContent.includes(expectedContent);
+        },
+        {
+            timeout: 60_000,
+            interval: 2000,
+            timeoutMsg: `Output file "${path}" did not contain "${expectedContent}"`,
+        }
+    );
+    await fs.rm(path);
+}
+
 describe("Run files on serverless compute", async function () {
     let projectDir: string;
     this.timeout(3 * 60 * 1000);
@@ -62,6 +78,13 @@ describe("Run files on serverless compute", async function () {
                             `df.to_json(os.path.join(os.getcwd(), "notebook-output.json"))`,
                         ],
                     },
+                    {
+                        cell_type: "code",
+                        execution_count: null,
+                        metadata: {},
+                        outputs: [],
+                        source: [`%run "./hello.py"`],
+                    },
                 ],
                 metadata: {
                     kernelspec: {
@@ -77,6 +100,24 @@ describe("Run files on serverless compute", async function () {
             })
         );
 
+        await fs.writeFile(
+            path.join(nestedDir, "databricks-notebook.py"),
+            [
+                "# Databricks notebook source",
+                `spark.sql('SELECT "hello world"').show()`,
+                "# COMMAND ----------",
+                "# DBTITLE 1,My cell title",
+                "# MAGIC %sql",
+                "# MAGIC select 1 + 1;",
+                "# MAGIC select 'hello; world'",
+                "# COMMAND ----------",
+                `df = _sqldf.toPandas()`,
+                `df.to_json(os.path.join(os.getcwd(), "databricks-notebook-output.json"))`,
+                "# COMMAND ----------",
+                "# MAGIC %run './hello.py'",
+            ].join("\n")
+        );
+
         await writeRootBundleConfig(
             getBasicBundleConfig({}, false),
             projectDir
@@ -202,21 +243,8 @@ describe("Run files on serverless compute", async function () {
         await executeCommandWhenAvailable(
             "Databricks: Run current file with Databricks Connect"
         );
-        await browser.waitUntil(
-            async () => {
-                const fileOutput = await fs.readFile(
-                    path.join(projectDir, "file-output.json"),
-                    "utf-8"
-                );
-                console.log("File output: ", fileOutput);
-                return fileOutput.includes("hello world");
-            },
-            {
-                timeout: 60_000,
-                interval: 2000,
-                timeoutMsg: "Terminal output did not contain 'hello world'",
-            }
-        );
+        const output = path.join(projectDir, "file-output.json");
+        await checkOutputFile(output, "hello world");
     });
 
     it("should run a notebook with dbconnect", async () => {
@@ -225,28 +253,43 @@ describe("Run files on serverless compute", async function () {
 
         const kernelInput = await waitForInput();
         await kernelInput.selectQuickPick("Python Environments...");
-        console.log(
-            "Selected 'Python Environments...' option for kernel selection"
-        );
+        console.log("Selected 'Python Environments...' option");
 
         const envInput = await waitForInput();
         await envInput.selectQuickPick(".venv");
         console.log("Selected .venv environment");
 
-        await browser.waitUntil(
-            async () => {
-                const notebookOutput = await fs.readFile(
-                    path.join(projectDir, "nested", "notebook-output.json"),
-                    "utf-8"
-                );
-                console.log("Notebook output: ", notebookOutput);
-                return notebookOutput.includes("hello world");
-            },
-            {
-                timeout: 60_000,
-                interval: 2000,
-                timeoutMsg: "Notebook execution did not complete successfully",
-            }
+        const firstCellOutput = path.join(
+            projectDir,
+            "nested",
+            "notebook-output.json"
+        );
+        await checkOutputFile(firstCellOutput, "hello world");
+
+        const secondCellOutput = path.join(
+            projectDir,
+            "nested",
+            "file-output.json"
+        );
+        await checkOutputFile(secondCellOutput, "hello world");
+    });
+
+    it("should run a databricks notebook with dbconnect and handle magic comments", async () => {
+        await openFile("databricks-notebook.py");
+        await executeCommandWhenAvailable("Jupyter: Run All Cells");
+
+        const sqlOutputFile = path.join(
+            projectDir,
+            "nested",
+            "databricks-notebook-output.json"
+        );
+        await checkOutputFile(sqlOutputFile, "hello; world");
+
+        const runOutputFile = path.join(
+            projectDir,
+            "nested",
+            "file-output.json"
         );
+        await checkOutputFile(runOutputFile, "hello world");
     });
 });