Discovery cancellation

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

@@ -159,7 +159,8 @@ export interface ITestDiscoveryAdapter {
     discoverTests(uri: Uri): Promise<DiscoveredTestPayload>;
     discoverTests(
         uri: Uri,
-        executionFactory: IPythonExecutionFactory,
+        executionFactory?: IPythonExecutionFactory,
+        token?: CancellationToken,
         interpreter?: PythonEnvironment,
     ): Promise<DiscoveredTestPayload>;
 }

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);
+        const cSource = new CancellationTokenSource();
+        const deferredReturn = createDeferred<DiscoveredTestPayload>();
+
+        token?.onCancellationRequested(() => {
+            traceInfo(`Test discovery cancelled.`);
+            cSource.cancel();
+            deferredReturn.resolve({ cwd: uri.fsPath, status: 'success' });
         });
 
-        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 is only a placeholder to handle function overloading until rewrite is finished
-        const discoveryPayload: DiscoveredTestPayload = { cwd: uri.fsPath, status: 'success' };
-        return discoveryPayload;
+        this.runPytestDiscovery(uri, name, cSource, executionFactory, interpreter, token).then(() => {
+            deferredReturn.resolve({ cwd: uri.fsPath, status: 'success' });
+        });
+
+        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.

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

@@ -244,7 +244,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.

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,30 @@ 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<DiscoveredTestPayload> {
         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();
+        const deferredReturn = createDeferred<DiscoveredTestPayload>();
+
+        token?.onCancellationRequested(() => {
+            traceInfo(`Test discovery cancelled.`);
+            cSource.cancel();
+            deferredReturn.resolve({ cwd: uri.fsPath, status: 'success' });
         });
 
+        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 +79,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({ cwd: uri.fsPath, status: 'success' });
+        });
+
+        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 +118,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 +169,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.

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);
             }

src/test/testing/testController/pytest/pytestDiscoveryAdapter.unit.test.ts

@@ -2,7 +2,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 import * as assert from 'assert';
-import { Uri } from 'vscode';
+import { Uri, CancellationTokenSource } from 'vscode';
 import * as typeMoq from 'typemoq';
 import * as path from 'path';
 import { Observable } from 'rxjs/Observable';
@@ -13,6 +13,7 @@ import { PytestTestDiscoveryAdapter } from '../../../../client/testing/testContr
 import {
     IPythonExecutionFactory,
     IPythonExecutionService,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
     SpawnOptions,
     Output,
 } from '../../../../client/common/process/types';
@@ -31,11 +32,13 @@ suite('pytest test discovery adapter', () => {
     let outputChannel: typeMoq.IMock<ITestOutputChannel>;
     let expectedPath: string;
     let uri: Uri;
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
     let expectedExtraVariables: Record<string, string>;
     let mockProc: MockChildProcess;
     let deferred2: Deferred<void>;
     let utilsStartDiscoveryNamedPipeStub: sinon.SinonStub;
     let useEnvExtensionStub: sinon.SinonStub;
+    let cancellationTokenSource: CancellationTokenSource;
 
     setup(() => {
         useEnvExtensionStub = sinon.stub(extapi, 'useEnvExtension');
@@ -86,9 +89,12 @@ suite('pytest test discovery adapter', () => {
                     },
                 };
             });
+
+        cancellationTokenSource = new CancellationTokenSource();
     });
     teardown(() => {
         sinon.restore();
+        cancellationTokenSource.dispose();
     });
     test('Discovery should call exec with correct basic args', async () => {
         // set up exec mock
@@ -333,4 +339,77 @@ suite('pytest test discovery adapter', () => {
             typeMoq.Times.once(),
         );
     });
+    test('Test discovery canceled before exec observable call finishes', async () => {
+        // set up exec mock
+        execFactory = typeMoq.Mock.ofType<IPythonExecutionFactory>();
+        execFactory
+            .setup((x) => x.createActivatedEnvironment(typeMoq.It.isAny()))
+            .returns(() => Promise.resolve(execService.object));
+
+        sinon.stub(fs.promises, 'lstat').callsFake(
+            async () =>
+                ({
+                    isFile: () => true,
+                    isSymbolicLink: () => false,
+                } as fs.Stats),
+        );
+        sinon.stub(fs.promises, 'realpath').callsFake(async (pathEntered) => pathEntered.toString());
+
+        adapter = new PytestTestDiscoveryAdapter(configService, outputChannel.object);
+        const discoveryPromise = adapter.discoverTests(uri, execFactory.object, cancellationTokenSource.token);
+
+        // Trigger cancellation before exec observable call finishes
+        cancellationTokenSource.cancel();
+
+        await discoveryPromise;
+
+        assert.ok(
+            true,
+            'Test resolves correctly when triggering a cancellation token immediately after starting discovery.',
+        );
+    });
+
+    test('Test discovery cancelled while exec observable is running and proc is closed', async () => {
+        //
+        const execService2 = typeMoq.Mock.ofType<IPythonExecutionService>();
+        execService2.setup((p) => ((p as unknown) as any).then).returns(() => undefined);
+        execService2
+            .setup((x) => x.execObservable(typeMoq.It.isAny(), typeMoq.It.isAny()))
+            .returns(() => {
+                // Trigger cancellation while exec observable is running
+                cancellationTokenSource.cancel();
+                return {
+                    proc: mockProc as any,
+                    out: new Observable<Output<string>>(),
+                    dispose: () => {
+                        /* no-body */
+                    },
+                };
+            });
+        // set up exec mock
+        deferred = createDeferred();
+        execFactory = typeMoq.Mock.ofType<IPythonExecutionFactory>();
+        execFactory
+            .setup((x) => x.createActivatedEnvironment(typeMoq.It.isAny()))
+            .returns(() => {
+                deferred.resolve();
+                return Promise.resolve(execService2.object);
+            });
+
+        sinon.stub(fs.promises, 'lstat').callsFake(
+            async () =>
+                ({
+                    isFile: () => true,
+                    isSymbolicLink: () => false,
+                } as fs.Stats),
+        );
+        sinon.stub(fs.promises, 'realpath').callsFake(async (pathEntered) => pathEntered.toString());
+
+        adapter = new PytestTestDiscoveryAdapter(configService, outputChannel.object);
+        const discoveryPromise = adapter.discoverTests(uri, execFactory.object, cancellationTokenSource.token);
+
+        // add in await and trigger
+        await discoveryPromise;
+        assert.ok(true, 'Test resolves correctly when triggering a cancellation token in exec observable.');
+    });
 });