Inject env vars into python unit test debugging sessions (#1120)

## Changes  ## Tests
databricks · Mar 7, 2024 · 264cb54 · 264cb54
1 parent 8d60383
commit 264cb54
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 11 deletions.
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
@@ -579,7 +579,10 @@ export async function activate(
 
     // Run/debug group
     const databricksDebugConfigurationProvider =
-        new DatabricksDebugConfigurationProvider(context);
+        new DatabricksDebugConfigurationProvider(
+            context,
+            databricksEnvFileManager
+        );
 
     const runCommands = new RunCommands(
         connectionManager,
@@ -606,6 +609,10 @@ export async function activate(
             "python",
             databricksDebugConfigurationProvider
         ),
+        debug.registerDebugConfigurationProvider(
+            "debugpy",
+            databricksDebugConfigurationProvider
+        ),
         telemetry.registerCommand(
             "databricks.run.dbconnect.debug",
             runCommands.debugFileUsingDbconnect,

diff --git a/packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts b/packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts
@@ -5,6 +5,7 @@ import {
     ExtensionContext,
 } from "vscode";
 import path from "node:path";
+import {DatabricksEnvFileManager} from "../file-managers/DatabricksEnvFileManager";
 
 export interface DatabricksPythonDebugConfiguration extends DebugConfiguration {
     databricks?: boolean;
@@ -14,27 +15,50 @@ export interface DatabricksPythonDebugConfiguration extends DebugConfiguration {
     console?: "integratedTerminal" | "externalTerminal" | "internalConsole";
 }
 
+function isTest(debugConfiguration: DebugConfiguration) {
+    return (
+        debugConfiguration.env?.RUN_TEST_IDS_PORT !== undefined ||
+        debugConfiguration.request === "test" ||
+        debugConfiguration.name === "Debug Unit Test"
+    );
+}
 export class DatabricksDebugConfigurationProvider
     implements DebugConfigurationProvider
 {
-    constructor(private readonly context: ExtensionContext) {}
+    constructor(
+        private readonly context: ExtensionContext,
+        private readonly databricksEnvFileManager: DatabricksEnvFileManager
+    ) {}
+
     async resolveDebugConfigurationWithSubstitutedVariables(
         folder: WorkspaceFolder | undefined,
         debugConfiguration: DebugConfiguration
     ) {
-        if (debugConfiguration.databricks !== true) {
+        if (
+            debugConfiguration.databricks !== true &&
+            !isTest(debugConfiguration)
+        ) {
             return debugConfiguration;
         }
 
-        const userProgram = debugConfiguration.program;
-        debugConfiguration.program = this.context.asAbsolutePath(
-            path.join("resources", "python", "dbconnect-bootstrap.py")
-        );
+        // Only add the bootstrap script if we are running explicit databricks debug
+        // configs. Other sources of configs such as tests, should not have the bootstrap
+        if (debugConfiguration.databricks) {
+            const userProgram = debugConfiguration.program;
+            debugConfiguration.program = this.context.asAbsolutePath(
+                path.join("resources", "python", "dbconnect-bootstrap.py")
+            );
+            debugConfiguration.args = [
+                userProgram,
+                ...(debugConfiguration.args ?? []),
+            ];
+        }
 
-        debugConfiguration.args = [
-            userProgram,
-            ...(debugConfiguration.args ?? []),
-        ];
+        // Explicitly set our env vars even though bootstrap loads them.
+        debugConfiguration.env = {
+            ...(await this.databricksEnvFileManager.getEnv()),
+            ...(debugConfiguration.env ?? {}),
+        };
 
         return debugConfiguration;
     }