Enable debug pytest (#21228)

fixes #21147

---------

Co-authored-by: Aidos Kanapyanov <65722512+aidoskanapyanov@users.noreply.github.com>
Co-authored-by: Karthik Nadig <kanadig@microsoft.com>

.vscode/launch.json

@@ -22,7 +22,8 @@
                 // Enable this to log telemetry to the output during debugging
                 "XVSC_PYTHON_LOG_TELEMETRY": "1",
                 // Enable this to log debugger output. Directory must exist ahead of time
-                "XDEBUGPY_LOG_DIR": "${workspaceRoot}/tmp/Debug_Output_Ex"
+                "XDEBUGPY_LOG_DIR": "${workspaceRoot}/tmp/Debug_Output_Ex",
+                "ENABLE_PYTHON_TESTING_REWRITE": "1"
             }
         },
         {

pythonFiles/vscode_pytest/__init__.py

@@ -179,6 +179,12 @@ def pytest_sessionfinish(session, exitstatus):
     4: Pytest encountered an internal error or exception during test execution.
     5: Pytest was unable to find any tests to run.
     """
+    print(
+        "pytest session has finished, exit status: ",
+        exitstatus,
+        "in discovery? ",
+        IS_DISCOVERY,
+    )
     cwd = pathlib.Path.cwd()
     if IS_DISCOVERY:
         try:
@@ -209,7 +215,6 @@ def pytest_sessionfinish(session, exitstatus):
                 f"Pytest exited with error status: {exitstatus}, {ERROR_MESSAGE_CONST[exitstatus]}"
             )
             exitstatus_bool = "error"
-
         execution_post(
             os.fsdecode(cwd),
             exitstatus_bool,

src/client/testing/common/debugLauncher.ts

@@ -15,6 +15,7 @@ import { ITestDebugLauncher, LaunchOptions } from './types';
 import { getConfigurationsForWorkspace } from '../../debugger/extension/configuration/launch.json/launchJsonReader';
 import { getWorkspaceFolder, getWorkspaceFolders } from '../../common/vscodeApis/workspaceApis';
 import { showErrorMessage } from '../../common/vscodeApis/windowApis';
+import { createDeferred } from '../../common/utils/async';
 
 @injectable()
 export class DebugLauncher implements ITestDebugLauncher {
@@ -42,16 +43,12 @@ export class DebugLauncher implements ITestDebugLauncher {
         );
         const debugManager = this.serviceContainer.get<IDebugService>(IDebugService);
 
-        return debugManager.startDebugging(workspaceFolder, launchArgs).then(
-            // Wait for debug session to be complete.
-            () =>
-                new Promise<void>((resolve) => {
-                    debugManager.onDidTerminateDebugSession(() => {
-                        resolve();
-                    });
-                }),
-            (ex) => traceError('Failed to start debugging tests', ex),
-        );
+        const deferred = createDeferred<void>();
+        debugManager.onDidTerminateDebugSession(() => {
+            deferred.resolve();
+        });
+        debugManager.startDebugging(workspaceFolder, launchArgs);
+        return deferred.promise;
     }
 
     private static resolveWorkspaceFolder(cwd: string): WorkspaceFolder {
@@ -181,6 +178,12 @@ export class DebugLauncher implements ITestDebugLauncher {
         const args = script(testArgs);
         const [program] = args;
         configArgs.program = program;
+        // if the test provider is pytest, then use the pytest module instead of using a program
+        const rewriteTestingEnabled = process.env.ENABLE_PYTHON_TESTING_REWRITE;
+        if (options.testProvider === 'pytest' && rewriteTestingEnabled) {
+            configArgs.module = 'pytest';
+            configArgs.program = undefined;
+        }
         configArgs.args = args.slice(1);
         // We leave configArgs.request as "test" so it will be sent in telemetry.
 
@@ -201,6 +204,21 @@ export class DebugLauncher implements ITestDebugLauncher {
             throw Error(`Invalid debug config "${debugConfig.name}"`);
         }
         launchArgs.request = 'launch';
+        if (options.testProvider === 'pytest' && rewriteTestingEnabled) {
+            if (options.pytestPort && options.pytestUUID) {
+                launchArgs.env = {
+                    ...launchArgs.env,
+                    TEST_PORT: options.pytestPort,
+                    TEST_UUID: options.pytestUUID,
+                };
+            } else {
+                throw Error(
+                    `Missing value for debug setup, both port and uuid need to be defined. port: "${options.pytestPort}" uuid: "${options.pytestUUID}"`,
+                );
+            }
+            const pluginPath = path.join(EXTENSION_ROOT_DIR, 'pythonFiles');
+            launchArgs.env.PYTHONPATH = pluginPath;
+        }
 
         // Clear out purpose so we can detect if the configuration was used to
         // run via F5 style debugging.
@@ -210,13 +228,19 @@ export class DebugLauncher implements ITestDebugLauncher {
     }
 
     private static getTestLauncherScript(testProvider: TestProvider) {
+        const rewriteTestingEnabled = process.env.ENABLE_PYTHON_TESTING_REWRITE;
         switch (testProvider) {
             case 'unittest': {
+                if (rewriteTestingEnabled) {
+                    return internalScripts.execution_py_testlauncher; // this is the new way to run unittest execution, debugger
+                }
                 return internalScripts.visualstudio_py_testlauncher; // old way unittest execution, debugger
-                // return internalScripts.execution_py_testlauncher; // this is the new way to run unittest execution, debugger
             }
             case 'pytest': {
-                return internalScripts.testlauncher;
+                if (rewriteTestingEnabled) {
+                    return (testArgs: string[]) => testArgs;
+                }
+                return internalScripts.testlauncher; // old way pytest execution, debugger
             }
             default: {
                 throw new Error(`Unknown test provider '${testProvider}'`);

src/client/testing/common/types.ts

@@ -25,6 +25,8 @@ export type LaunchOptions = {
     testProvider: TestProvider;
     token?: CancellationToken;
     outChannel?: OutputChannel;
+    pytestPort?: string;
+    pytestUUID?: string;
 };
 
 export type ParserOptions = TestDiscoveryOptions;

src/client/testing/testController/common/types.ts

@@ -12,7 +12,7 @@ import {
     Uri,
     WorkspaceFolder,
 } from 'vscode';
-import { TestDiscoveryOptions } from '../../common/types';
+import { ITestDebugLauncher, TestDiscoveryOptions } from '../../common/types';
 import { IPythonExecutionFactory } from '../../../common/process/types';
 
 export type TestRunInstanceOptions = TestRunOptions & {
@@ -193,6 +193,7 @@ export interface ITestExecutionAdapter {
         testIds: string[],
         debugBool?: boolean,
         executionFactory?: IPythonExecutionFactory,
+        debugLauncher?: ITestDebugLauncher,
     ): Promise<ExecutionTestPayload>;
 }
 

src/client/testing/testController/controller.ts

@@ -409,6 +409,7 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                                     token,
                                     request.profile?.kind === TestRunProfileKind.Debug,
                                     this.pythonExecFactory,
+                                    this.debugLauncher,
                                 );
                             }
                             return this.pytest.runTests(
@@ -438,6 +439,7 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                                     testItems,
                                     token,
                                     request.profile?.kind === TestRunProfileKind.Debug,
+                                    this.pythonExecFactory,
                                 );
                             }
                             // below is old way of running unittest execution

src/client/testing/testController/pytest/pytestExecutionAdapter.ts

@@ -12,11 +12,13 @@ import {
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
-import { EXTENSION_ROOT_DIR } from '../../../constants';
 import { removePositionalFoldersAndFiles } from './arguments';
+import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
+import { PYTEST_PROVIDER } from '../../common/constants';
+import { EXTENSION_ROOT_DIR } from '../../../common/constants';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-(global as any).EXTENSION_ROOT_DIR = EXTENSION_ROOT_DIR;
+// (global as any).EXTENSION_ROOT_DIR = EXTENSION_ROOT_DIR;
 /**
  * Wrapper Class for pytest test execution. This is where we call `runTestCommand`?
  */
@@ -47,11 +49,12 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         testIds: string[],
         debugBool?: boolean,
         executionFactory?: IPythonExecutionFactory,
+        debugLauncher?: ITestDebugLauncher,
     ): Promise<ExecutionTestPayload> {
         traceVerbose(uri, testIds, debugBool);
         if (executionFactory !== undefined) {
             // ** new version of run tests.
-            return this.runTestsNew(uri, testIds, debugBool, executionFactory);
+            return this.runTestsNew(uri, testIds, debugBool, executionFactory, debugLauncher);
         }
         // if executionFactory is undefined, we are using the old method signature of run tests.
         this.outputChannel.appendLine('Running tests.');
@@ -64,6 +67,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         testIds: string[],
         debugBool?: boolean,
         executionFactory?: IPythonExecutionFactory,
+        debugLauncher?: ITestDebugLauncher,
     ): Promise<ExecutionTestPayload> {
         const deferred = createDeferred<ExecutionTestPayload>();
         const relativePathToPytest = 'pythonFiles';
@@ -106,16 +110,29 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 testArgs.splice(0, 0, '--rootdir', uri.fsPath);
             }
 
+            // why is this needed?
             if (debugBool && !testArgs.some((a) => a.startsWith('--capture') || a === '-s')) {
                 testArgs.push('--capture', 'no');
             }
-
-            console.debug(`Running test with arguments: ${testArgs.join(' ')}\r\n`);
-            console.debug(`Current working directory: ${uri.fsPath}\r\n`);
-
-            const argArray = ['-m', 'pytest', '-p', 'vscode_pytest'].concat(testArgs).concat(testIds);
-            console.debug('argArray', argArray);
-            execService?.exec(argArray, spawnOptions);
+            const pluginArgs = ['-p', 'vscode_pytest', '-v'].concat(testArgs).concat(testIds);
+            if (debugBool) {
+                const pytestPort = this.testServer.getPort().toString();
+                const pytestUUID = uuid.toString();
+                const launchOptions: LaunchOptions = {
+                    cwd: uri.fsPath,
+                    args: pluginArgs,
+                    token: spawnOptions.token,
+                    testProvider: PYTEST_PROVIDER,
+                    pytestPort,
+                    pytestUUID,
+                };
+                console.debug(`Running debug test with arguments: ${pluginArgs.join(' ')}\r\n`);
+                await debugLauncher!.launchDebugger(launchOptions);
+            } else {
+                const runArgs = ['-m', 'pytest'].concat(pluginArgs);
+                console.debug(`Running test with arguments: ${runArgs.join(' ')}\r\n`);
+                execService?.exec(runArgs, spawnOptions);
+            }
         } catch (ex) {
             console.debug(`Error while running tests: ${testIds}\r\n${ex}\r\n\r\n`);
             return Promise.reject(ex);

src/client/testing/testController/workspaceTestAdapter.ts

@@ -38,6 +38,7 @@ import {
 } from './common/types';
 import { fixLogLines } from './common/utils';
 import { IPythonExecutionFactory } from '../../common/process/types';
+import { ITestDebugLauncher } from '../common/types';
 
 /**
  * This class exposes a test-provider-agnostic way of discovering tests.
@@ -77,6 +78,7 @@ export class WorkspaceTestAdapter {
         token?: CancellationToken,
         debugBool?: boolean,
         executionFactory?: IPythonExecutionFactory,
+        debugLauncher?: ITestDebugLauncher,
     ): Promise<void> {
         if (this.executing) {
             return this.executing.promise;
@@ -110,6 +112,7 @@ export class WorkspaceTestAdapter {
                     testCaseIds,
                     debugBool,
                     executionFactory,
+                    debugLauncher,
                 );
                 traceVerbose('executionFactory defined');
             } else {