fix promise returns and add tests for discovery cancellation

Author: eleanorjboyd

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

@@ -10,7 +10,7 @@ import {
     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';
@@ -44,27 +44,27 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         token?: CancellationToken,
         interpreter?: PythonEnvironment,
     ): Promise<DiscoveredTestPayload> {
+        const cSource = new CancellationTokenSource();
+        const deferredReturn = createDeferred<DiscoveredTestPayload>();
+
         token?.onCancellationRequested(() => {
             traceInfo(`Test discovery cancelled.`);
-            const discoveryPayload: DiscoveredTestPayload = { cwd: uri.fsPath, status: 'success' };
-            return discoveryPayload;
+            cSource.cancel();
+            deferredReturn.resolve({ cwd: uri.fsPath, status: 'success' });
         });
 
-        const cSource = new CancellationTokenSource();
-        token?.onCancellationRequested(() => cSource.cancel());
-
         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);
 
-        await this.runPytestDiscovery(uri, name, cSource, executionFactory, interpreter, token);
+        this.runPytestDiscovery(uri, name, cSource, executionFactory, interpreter, token).then(() => {
+            deferredReturn.resolve({ cwd: uri.fsPath, status: 'success' });
+        });
 
-        // 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(

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

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

@@ -51,27 +51,21 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         const { unittestArgs } = settings.testing;
         const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
 
+        const cSource = new CancellationTokenSource();
+        const deferredReturn = createDeferred<DiscoveredTestPayload>();
+
         token?.onCancellationRequested(() => {
             traceInfo(`Test discovery cancelled.`);
-            const discoveryPayload: DiscoveredTestPayload = { cwd: uri.fsPath, status: 'success' };
-            return discoveryPayload;
+            cSource.cancel();
+            deferredReturn.resolve({ cwd: uri.fsPath, status: 'success' });
         });
 
-        const cSource = new CancellationTokenSource();
-        token?.onCancellationRequested(() => cSource.cancel());
-
         const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
             if (!token?.isCancellationRequested) {
                 this.resultResolver?.resolveDiscovery(data);
             }
         }, cSource.token);
 
-        token?.onCancellationRequested(() => {
-            traceInfo(`Test discovery cancelled.`);
-            const discoveryPayload: DiscoveredTestPayload = { cwd: uri.fsPath, status: 'success' };
-            return discoveryPayload;
-        });
-
         // set up env with the pipe name
         let env: EnvironmentVariables | undefined = await this.envVarsService?.getEnvironmentVariables(uri);
         if (env === undefined) {
@@ -88,15 +82,11 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             token,
         };
 
-        try {
-            await this.runDiscovery(uri, options, name, cwd, cSource, 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(

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.');
+    });
 });