Discovery cancellation (#24713)

eleanorjboyd · web-flow · commit 8c54b8aa69e1 · 2025-01-14T23:55:40.000Z
fixes #24602
diff --git a/src/client/testing/testController/common/types.ts b/src/client/testing/testController/common/types.ts
@@ -156,18 +156,19 @@ export interface ITestResultResolver {
 }
 export interface ITestDiscoveryAdapter {
     // ** first line old method signature, second line new method signature
-    discoverTests(uri: Uri): Promise<DiscoveredTestPayload>;
+    discoverTests(uri: Uri): Promise<void>;
     discoverTests(
         uri: Uri,
-        executionFactory: IPythonExecutionFactory,
+        executionFactory?: IPythonExecutionFactory,
+        token?: CancellationToken,
         interpreter?: PythonEnvironment,
-    ): Promise<DiscoveredTestPayload>;
+    ): Promise<void>;
 }
 
 // interface for execution/runner adapter
 export interface ITestExecutionAdapter {
     // ** first line old method signature, second line new method signature
-    runTests(uri: Uri, testIds: string[], profileKind?: boolean | TestRunProfileKind): Promise<ExecutionTestPayload>;
+    runTests(uri: Uri, testIds: string[], profileKind?: boolean | TestRunProfileKind): Promise<void>;
     runTests(
         uri: Uri,
         testIds: string[],
@@ -176,7 +177,7 @@ export interface ITestExecutionAdapter {
         executionFactory?: IPythonExecutionFactory,
         debugLauncher?: ITestDebugLauncher,
         interpreter?: PythonEnvironment,
-    ): Promise<ExecutionTestPayload>;
+    ): Promise<void>;
 }
 
 // Same types as in python_files/unittestadapter/utils.py
diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -1,15 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 import * as path from 'path';
-import { Uri } from 'vscode';
+import { CancellationToken, CancellationTokenSource, Uri } from 'vscode';
 import * as fs from 'fs';
+import { ChildProcess } from 'child_process';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
 import { IConfigurationService, ITestOutputChannel } from '../../../common/types';
-import { Deferred } from '../../../common/utils/async';
+import { createDeferred, Deferred } from '../../../common/utils/async';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
 import { traceError, traceInfo, traceVerbose, traceWarn } from '../../../logging';
 import { DiscoveredTestPayload, ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
@@ -40,24 +41,39 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
     async discoverTests(
         uri: Uri,
         executionFactory?: IPythonExecutionFactory,
+        token?: CancellationToken,
         interpreter?: PythonEnvironment,
-    ): Promise<DiscoveredTestPayload> {
-        const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-            this.resultResolver?.resolveDiscovery(data);
+    ): Promise<void> {
+        const cSource = new CancellationTokenSource();
+        const deferredReturn = createDeferred<void>();
+
+        token?.onCancellationRequested(() => {
+            traceInfo(`Test discovery cancelled.`);
+            cSource.cancel();
+            deferredReturn.resolve();
         });
 
-        await this.runPytestDiscovery(uri, name, executionFactory, interpreter);
+        const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+            // if the token is cancelled, we don't want process the data
+            if (!token?.isCancellationRequested) {
+                this.resultResolver?.resolveDiscovery(data);
+            }
+        }, cSource.token);
+
+        this.runPytestDiscovery(uri, name, cSource, executionFactory, interpreter, token).then(() => {
+            deferredReturn.resolve();
+        });
 
-        // this is only a placeholder to handle function overloading until rewrite is finished
-        const discoveryPayload: DiscoveredTestPayload = { cwd: uri.fsPath, status: 'success' };
-        return discoveryPayload;
+        return deferredReturn.promise;
     }
 
     async runPytestDiscovery(
         uri: Uri,
         discoveryPipeName: string,
+        cSource: CancellationTokenSource,
         executionFactory?: IPythonExecutionFactory,
         interpreter?: PythonEnvironment,
+        token?: CancellationToken,
     ): Promise<void> {
         const relativePathToPytest = 'python_files';
         const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
@@ -111,6 +127,12 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                     args: execArgs,
                     env: (mutableEnv as unknown) as { [key: string]: string },
                 });
+                token?.onCancellationRequested(() => {
+                    traceInfo(`Test discovery cancelled, killing pytest subprocess for workspace ${uri.fsPath}`);
+                    proc.kill();
+                    deferredTillExecClose.resolve();
+                    cSource.cancel();
+                });
                 proc.stdout.on('data', (data) => {
                     const out = fixLogLinesNoTrailing(data.toString());
                     traceInfo(out);
@@ -143,6 +165,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             throwOnStdErr: true,
             outputChannel: this.outputChannel,
             env: mutableEnv,
+            token,
         };
 
         // Create the Python environment in which to execute the command.
@@ -154,7 +177,21 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         const execService = await executionFactory?.createActivatedEnvironment(creationOptions);
 
         const deferredTillExecClose: Deferred<void> = createTestingDeferred();
+
+        let resultProc: ChildProcess | undefined;
+
+        token?.onCancellationRequested(() => {
+            traceInfo(`Test discovery cancelled, killing pytest subprocess for workspace ${uri.fsPath}`);
+            // if the resultProc exists just call kill on it which will handle resolving the ExecClose deferred, otherwise resolve the deferred here.
+            if (resultProc) {
+                resultProc?.kill();
+            } else {
+                deferredTillExecClose.resolve();
+                cSource.cancel();
+            }
+        });
         const result = execService?.execObservable(execArgs, spawnOptions);
+        resultProc = result?.proc;
 
         // Take all output from the subprocess and add it to the test output channel. This will be the pytest output.
         // Displays output to user and ensure the subprocess doesn't run into buffer overflow.
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -38,7 +38,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         executionFactory?: IPythonExecutionFactory,
         debugLauncher?: ITestDebugLauncher,
         interpreter?: PythonEnvironment,
-    ): Promise<ExecutionTestPayload> {
+    ): Promise<void> {
         const deferredTillServerClose: Deferred<void> = utils.createTestingDeferred();
 
         // create callback to handle data received on the named pipe
@@ -59,12 +59,6 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         );
         runInstance?.token.onCancellationRequested(() => {
             traceInfo(`Test run cancelled, resolving 'TillServerClose' deferred for ${uri.fsPath}.`);
-            const executionPayload: ExecutionTestPayload = {
-                cwd: uri.fsPath,
-                status: 'success',
-                error: '',
-            };
-            return executionPayload;
         });
 
         try {
@@ -82,15 +76,6 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         } finally {
             await deferredTillServerClose.promise;
         }
-
-        // placeholder until after the rewrite is adopted
-        // TODO: remove after adoption.
-        const executionPayload: ExecutionTestPayload = {
-            cwd: uri.fsPath,
-            status: 'success',
-            error: '',
-        };
-        return executionPayload;
     }
 
     private async runTestsNew(
@@ -244,7 +229,6 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 });
 
                 const result = execService?.execObservable(runArgs, spawnOptions);
-                resultProc = result?.proc;
 
                 // Take all output from the subprocess and add it to the test output channel. This will be the pytest output.
                 // Displays output to user and ensure the subprocess doesn't run into buffer overflow.
diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -2,7 +2,9 @@
 // Licensed under the MIT License.
 
 import * as path from 'path';
-import { Uri } from 'vscode';
+import { CancellationTokenSource, Uri } from 'vscode';
+import { CancellationToken } from 'vscode-jsonrpc';
+import { ChildProcess } from 'child_process';
 import { IConfigurationService, ITestOutputChannel } from '../../../common/types';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
 import {
@@ -40,15 +42,31 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         private readonly envVarsService?: IEnvironmentVariablesProvider,
     ) {}
 
-    public async discoverTests(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
+    public async discoverTests(
+        uri: Uri,
+        executionFactory?: IPythonExecutionFactory,
+        token?: CancellationToken,
+    ): Promise<void> {
         const settings = this.configSettings.getSettings(uri);
         const { unittestArgs } = settings.testing;
         const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
 
-        const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-            this.resultResolver?.resolveDiscovery(data);
+        const cSource = new CancellationTokenSource();
+        // Create a deferred to return to the caller
+        const deferredReturn = createDeferred<void>();
+
+        token?.onCancellationRequested(() => {
+            traceInfo(`Test discovery cancelled.`);
+            cSource.cancel();
+            deferredReturn.resolve();
         });
 
+        const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+            if (!token?.isCancellationRequested) {
+                this.resultResolver?.resolveDiscovery(data);
+            }
+        }, cSource.token);
+
         // set up env with the pipe name
         let env: EnvironmentVariables | undefined = await this.envVarsService?.getEnvironmentVariables(uri);
         if (env === undefined) {
@@ -62,24 +80,22 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             command,
             cwd,
             outChannel: this.outputChannel,
+            token,
         };
 
-        try {
-            await this.runDiscovery(uri, options, name, cwd, executionFactory);
-        } finally {
-            // none
-        }
-        // placeholder until after the rewrite is adopted
-        // TODO: remove after adoption.
-        const discoveryPayload: DiscoveredTestPayload = { cwd, status: 'success' };
-        return discoveryPayload;
+        this.runDiscovery(uri, options, name, cwd, cSource, executionFactory).then(() => {
+            deferredReturn.resolve();
+        });
+
+        return deferredReturn.promise;
     }
 
     async runDiscovery(
         uri: Uri,
         options: TestCommandOptions,
         testRunPipeName: string,
         cwd: string,
+        cSource: CancellationTokenSource,
         executionFactory?: IPythonExecutionFactory,
     ): Promise<void> {
         // get and edit env vars
@@ -103,6 +119,12 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                     args,
                     env: (mutableEnv as unknown) as { [key: string]: string },
                 });
+                options.token?.onCancellationRequested(() => {
+                    traceInfo(`Test discovery cancelled, killing unittest subprocess for workspace ${uri.fsPath}`);
+                    proc.kill();
+                    deferredTillExecClose.resolve();
+                    cSource.cancel();
+                });
                 proc.stdout.on('data', (data) => {
                     const out = fixLogLinesNoTrailing(data.toString());
                     traceInfo(out);
@@ -148,7 +170,19 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
             const execService = await executionFactory?.createActivatedEnvironment(creationOptions);
 
+            let resultProc: ChildProcess | undefined;
+            options.token?.onCancellationRequested(() => {
+                traceInfo(`Test discovery cancelled, killing unittest subprocess for workspace ${uri.fsPath}`);
+                // if the resultProc exists just call kill on it which will handle resolving the ExecClose deferred, otherwise resolve the deferred here.
+                if (resultProc) {
+                    resultProc?.kill();
+                } else {
+                    deferredTillExecClose.resolve();
+                    cSource.cancel();
+                }
+            });
             const result = execService?.execObservable(args, spawnOptions);
+            resultProc = result?.proc;
 
             // Displays output to user and ensure the subprocess doesn't run into buffer overflow.
             // TODO: after a release, remove discovery output from the "Python Test Log" channel and send it to the "Python" channel instead.
diff --git a/src/client/testing/testController/unittest/testExecutionAdapter.ts b/src/client/testing/testController/unittest/testExecutionAdapter.ts
@@ -47,7 +47,7 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         runInstance?: TestRun,
         executionFactory?: IPythonExecutionFactory,
         debugLauncher?: ITestDebugLauncher,
-    ): Promise<ExecutionTestPayload> {
+    ): Promise<void> {
         // deferredTillServerClose awaits named pipe server close
         const deferredTillServerClose: Deferred<void> = utils.createTestingDeferred();
 
@@ -87,12 +87,6 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         } finally {
             await deferredTillServerClose.promise;
         }
-        const executionPayload: ExecutionTestPayload = {
-            cwd: uri.fsPath,
-            status: 'success',
-            error: '',
-        };
-        return executionPayload;
     }
 
     private async runTestsNew(
diff --git a/src/client/testing/testController/workspaceTestAdapter.ts b/src/client/testing/testController/workspaceTestAdapter.ts
@@ -134,7 +134,7 @@ export class WorkspaceTestAdapter {
         try {
             // ** execution factory only defined for new rewrite way
             if (executionFactory !== undefined) {
-                await this.discoveryAdapter.discoverTests(this.workspaceUri, executionFactory, interpreter);
+                await this.discoveryAdapter.discoverTests(this.workspaceUri, executionFactory, token, interpreter);
             } else {
                 await this.discoveryAdapter.discoverTests(this.workspaceUri);
             }
diff --git a/src/test/testing/testController/pytest/pytestDiscoveryAdapter.unit.test.ts b/src/test/testing/testController/pytest/pytestDiscoveryAdapter.unit.test.ts
diff --git a/src/test/testing/testController/unittest/testDiscoveryAdapter.unit.test.ts b/src/test/testing/testController/unittest/testDiscoveryAdapter.unit.test.ts

Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,7 @@ export class WorkspaceTestAdapter {`
`134`	`134`	`try {`
`135`	`135`	`// ** execution factory only defined for new rewrite way`
`136`	`136`	`if (executionFactory !== undefined) {`
`137`		`- await this.discoveryAdapter.discoverTests(this.workspaceUri, executionFactory, interpreter);`
	`137`	`+ await this.discoveryAdapter.discoverTests(this.workspaceUri, executionFactory, token, interpreter);`
`138`	`138`	`} else {`
`139`	`139`	`await this.discoveryAdapter.discoverTests(this.workspaceUri);`
`140`	`140`	`}`