new inactive discovery logic (#20566)

eleanorjboyd · web-flow · commit 3fe705755482 · 2023-01-31T14:46:01.000-08:00
New pytest code in a currently inactive state.
diff --git a/src/client/testing/testController/common/server.ts b/src/client/testing/testController/common/server.ts
@@ -71,6 +71,12 @@ export class PythonTestServer implements ITestServer, Disposable {
         return (this.server.address() as net.AddressInfo).port;
     }
 
+    public createUUID(cwd: string): string {
+        const uuid = crypto.randomUUID();
+        this.uuids.set(uuid, cwd);
+        return uuid;
+    }
+
     public dispose(): void {
         this.server.close();
         this._onDataReceived.dispose();
@@ -81,15 +87,13 @@ export class PythonTestServer implements ITestServer, Disposable {
     }
 
     async sendCommand(options: TestCommandOptions): Promise<void> {
-        const uuid = crypto.randomUUID();
+        const uuid = this.createUUID(options.cwd);
         const spawnOptions: SpawnOptions = {
             token: options.token,
             cwd: options.cwd,
             throwOnStdErr: true,
         };
 
-        this.uuids.set(uuid, options.cwd);
-
         // Create the Python environment in which to execute the command.
         const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
             allowEnvironmentFetchExceptions: false,
diff --git a/src/client/testing/testController/common/types.ts b/src/client/testing/testController/common/types.ts
@@ -151,6 +151,17 @@ export type TestCommandOptions = {
     testIds?: string[];
 };
 
+export type TestCommandOptionsPytest = {
+    workspaceFolder: Uri;
+    cwd: string;
+    commandStr: string;
+    token?: CancellationToken;
+    outChannel?: OutputChannel;
+    debugBool?: boolean;
+    testIds?: string[];
+    env: { [key: string]: string | undefined };
+};
+
 /**
  * Interface describing the server that will send test commands to the Python side, and process responses.
  *
@@ -161,10 +172,14 @@ export interface ITestServer {
     readonly onDataReceived: Event<DataReceivedEvent>;
     sendCommand(options: TestCommandOptions): Promise<void>;
     serverReady(): Promise<void>;
+    getPort(): number;
+    createUUID(cwd: string): string;
 }
 
 export interface ITestDiscoveryAdapter {
+    // ** Uncomment second line and comment out first line to use the new discovery method.
     discoverTests(uri: Uri): Promise<DiscoveredTestPayload>;
+    // discoverTests(uri: Uri, executionFactory: IPythonExecutionFactory): Promise<DiscoveredTestPayload>
 }
 
 // interface for execution/runner adapter
diff --git a/src/client/testing/testController/controller.ts b/src/client/testing/testController/controller.ts
@@ -39,8 +39,10 @@ import {
     ITestExecutionAdapter,
 } from './common/types';
 import { UnittestTestDiscoveryAdapter } from './unittest/testDiscoveryAdapter';
-import { WorkspaceTestAdapter } from './workspaceTestAdapter';
 import { UnittestTestExecutionAdapter } from './unittest/testExecutionAdapter';
+import { PytestTestDiscoveryAdapter } from './pytest/pytestDiscoveryAdapter';
+import { PytestTestExecutionAdapter } from './pytest/pytestExecutionAdapter';
+import { WorkspaceTestAdapter } from './workspaceTestAdapter';
 import { ITestDebugLauncher } from '../common/types';
 
 // Types gymnastics to make sure that sendTriggerTelemetry only accepts the correct types.
@@ -141,7 +143,6 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
             });
             return this.refreshTestData(undefined, { forceRefresh: true });
         };
-
         this.pythonTestServer = new PythonTestServer(this.pythonExecFactory, this.debugLauncher);
     }
 
@@ -161,10 +162,8 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                 executionAdapter = new UnittestTestExecutionAdapter(this.pythonTestServer, this.configSettings);
                 testProvider = UNITTEST_PROVIDER;
             } else {
-                // TODO: PYTEST DISCOVERY ADAPTER
-                // this is a placeholder for now
-                discoveryAdapter = new UnittestTestDiscoveryAdapter(this.pythonTestServer, { ...this.configSettings });
-                executionAdapter = new UnittestTestExecutionAdapter(this.pythonTestServer, this.configSettings);
+                discoveryAdapter = new PytestTestDiscoveryAdapter(this.pythonTestServer, { ...this.configSettings });
+                executionAdapter = new PytestTestExecutionAdapter(this.pythonTestServer, this.configSettings);
                 testProvider = PYTEST_PROVIDER;
             }
 
@@ -224,18 +223,30 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
         this.refreshingStartedEvent.fire();
         if (uri) {
             const settings = this.configSettings.getSettings(uri);
+            traceVerbose(`Testing: Refreshing test data for ${uri.fsPath}`);
             if (settings.testing.pytestEnabled) {
-                traceVerbose(`Testing: Refreshing test data for ${uri.fsPath}`);
-
                 // Ensure we send test telemetry if it gets disabled again
                 this.sendTestDisabledTelemetry = true;
-
+                // ** uncomment ~231 - 241 to NEW new test discovery mechanism
+                // const workspace = this.workspaceService.getWorkspaceFolder(uri);
+                // traceVerbose(`Discover tests for workspace name: ${workspace?.name} - uri: ${uri.fsPath}`);
+                // const testAdapter =
+                //     this.testAdapters.get(uri) || (this.testAdapters.values().next().value as WorkspaceTestAdapter);
+                // testAdapter.discoverTests(
+                //     this.testController,
+                //     this.refreshCancellation.token,
+                //     this.testAdapters.size > 1,
+                //     this.workspaceService.workspaceFile?.fsPath,
+                //     this.pythonExecFactory,
+                // );
+                // uncomment ~243 to use OLD test discovery mechanism
                 await this.pytest.refreshTestData(this.testController, uri, this.refreshCancellation.token);
             } else if (settings.testing.unittestEnabled) {
-                // TODO: Use new test discovery mechanism
-                // traceVerbose(`Testing: Refreshing test data for ${uri.fsPath}`);
+                // ** Ensure we send test telemetry if it gets disabled again
+                this.sendTestDisabledTelemetry = true;
+                // uncomment ~248 - 258 to NEW new test discovery mechanism
                 // const workspace = this.workspaceService.getWorkspaceFolder(uri);
-                // console.warn(`Discover tests for workspace name: ${workspace?.name} - uri: ${uri.fsPath}`);
+                // traceVerbose(`Discover tests for workspace name: ${workspace?.name} - uri: ${uri.fsPath}`);
                 // const testAdapter =
                 //     this.testAdapters.get(uri) || (this.testAdapters.values().next().value as WorkspaceTestAdapter);
                 // testAdapter.discoverTests(
@@ -244,9 +255,7 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                 //     this.testAdapters.size > 1,
                 //     this.workspaceService.workspaceFile?.fsPath,
                 // );
-                // // Ensure we send test telemetry if it gets disabled again
-                // this.sendTestDisabledTelemetry = true;
-                // comment below 229 to run the new way and uncomment above 212 ~ 227
+                // uncomment ~260 to use OLD test discovery mechanism
                 await this.unittest.refreshTestData(this.testController, uri, this.refreshCancellation.token);
             } else {
                 if (this.sendTestDisabledTelemetry) {
@@ -375,7 +384,7 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                             );
                         }
                         if (settings.testing.unittestEnabled) {
-                            // potentially sqeeze in the new exeuction way here?
+                            // potentially squeeze in the new execution way here?
                             sendTelemetryEvent(EventName.UNITTEST_RUN, undefined, {
                                 tool: 'unittest',
                                 debugging: request.profile?.kind === TestRunProfileKind.Debug,
diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import * as path from 'path';
+import { Uri } from 'vscode';
+import {
+    ExecutionFactoryCreateWithEnvironmentOptions,
+    IPythonExecutionFactory,
+    SpawnOptions,
+} from '../../../common/process/types';
+import { IConfigurationService } from '../../../common/types';
+import { createDeferred, Deferred } from '../../../common/utils/async';
+import { EXTENSION_ROOT_DIR } from '../../../constants';
+import { traceVerbose } from '../../../logging';
+import { DataReceivedEvent, DiscoveredTestPayload, ITestDiscoveryAdapter, ITestServer } from '../common/types';
+
+/**
+ * Wrapper class for unittest test discovery. This is where we call `runTestCommand`. #this seems incorrectly copied
+ */
+export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
+    private deferred: Deferred<DiscoveredTestPayload> | undefined;
+
+    private cwd: string | undefined;
+
+    constructor(public testServer: ITestServer, public configSettings: IConfigurationService) {
+        testServer.onDataReceived(this.onDataReceivedHandler, this);
+    }
+
+    public onDataReceivedHandler({ cwd, data }: DataReceivedEvent): void {
+        if (this.deferred && cwd === this.cwd) {
+            const testData: DiscoveredTestPayload = JSON.parse(data);
+
+            this.deferred.resolve(testData);
+            this.deferred = undefined;
+        }
+    }
+
+    // ** Old version of discover tests.
+    discoverTests(uri: Uri): Promise<DiscoveredTestPayload> {
+        traceVerbose(uri);
+        this.deferred = createDeferred<DiscoveredTestPayload>();
+        return this.deferred.promise;
+    }
+    // Uncomment this version of the function discoverTests to use the new discovery method.
+    // public async discoverTests(uri: Uri, executionFactory: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
+    //     const settings = this.configSettings.getSettings(uri);
+    //     const { pytestArgs } = settings.testing;
+    //     traceVerbose(pytestArgs);
+
+    //     this.cwd = uri.fsPath;
+    //     return this.runPytestDiscovery(uri, executionFactory);
+    // }
+
+    async runPytestDiscovery(uri: Uri, executionFactory: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
+        if (!this.deferred) {
+            this.deferred = createDeferred<DiscoveredTestPayload>();
+            const relativePathToPytest = 'pythonFiles';
+            const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
+            const uuid = this.testServer.createUUID(uri.fsPath);
+            const settings = this.configSettings.getSettings(uri);
+            const { pytestArgs } = settings.testing;
+
+            const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
+            const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
+
+            const spawnOptions: SpawnOptions = {
+                cwd: uri.fsPath,
+                throwOnStdErr: true,
+                extraVariables: {
+                    PYTHONPATH: pythonPathCommand,
+                    TEST_UUID: uuid.toString(),
+                    TEST_PORT: this.testServer.getPort().toString(),
+                },
+            };
+
+            // Create the Python environment in which to execute the command.
+            const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
+                allowEnvironmentFetchExceptions: false,
+                resource: uri,
+            };
+            const execService = await executionFactory.createActivatedEnvironment(creationOptions);
+
+            try {
+                execService.exec(
+                    ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs),
+                    spawnOptions,
+                );
+            } catch (ex) {
+                console.error(ex);
+            }
+        }
+        return this.deferred.promise;
+    }
+}
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -0,0 +1,73 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import * as path from 'path';
+import { Uri } from 'vscode';
+import { IConfigurationService } from '../../../common/types';
+import { createDeferred, Deferred } from '../../../common/utils/async';
+import { EXTENSION_ROOT_DIR } from '../../../constants';
+import {
+    DataReceivedEvent,
+    ExecutionTestPayload,
+    ITestExecutionAdapter,
+    ITestServer,
+    TestCommandOptions,
+    TestExecutionCommand,
+} from '../common/types';
+
+/**
+ * Wrapper Class for unittest test execution. This is where we call `runTestCommand`?
+ */
+
+export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
+    private deferred: Deferred<ExecutionTestPayload> | undefined;
+
+    private cwd: string | undefined;
+
+    constructor(public testServer: ITestServer, public configSettings: IConfigurationService) {
+        testServer.onDataReceived(this.onDataReceivedHandler, this);
+    }
+
+    public onDataReceivedHandler({ cwd, data }: DataReceivedEvent): void {
+        if (this.deferred && cwd === this.cwd) {
+            const testData: ExecutionTestPayload = JSON.parse(data);
+
+            this.deferred.resolve(testData);
+            this.deferred = undefined;
+        }
+    }
+
+    public async runTests(uri: Uri, testIds: string[], debugBool?: boolean): Promise<ExecutionTestPayload> {
+        if (!this.deferred) {
+            const settings = this.configSettings.getSettings(uri);
+            const { unittestArgs } = settings.testing;
+
+            const command = buildExecutionCommand(unittestArgs);
+            this.cwd = uri.fsPath;
+
+            const options: TestCommandOptions = {
+                workspaceFolder: uri,
+                command,
+                cwd: this.cwd,
+                debugBool,
+                testIds,
+            };
+
+            this.deferred = createDeferred<ExecutionTestPayload>();
+
+            // send test command to server
+            // server fire onDataReceived event once it gets response
+            this.testServer.sendCommand(options);
+        }
+        return this.deferred.promise;
+    }
+}
+
+function buildExecutionCommand(args: string[]): TestExecutionCommand {
+    const executionScript = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'unittestadapter', 'execution.py');
+
+    return {
+        script: executionScript,
+        args: ['--udiscovery', ...args],
+    };
+}
diff --git a/src/client/testing/testController/workspaceTestAdapter.ts b/src/client/testing/testController/workspaceTestAdapter.ts
@@ -196,6 +196,7 @@ export class WorkspaceTestAdapter {
         return Promise.resolve();
     }
 
+    // add `executionFactory?: IPythonExecutionFactory,` to the function for new pytest method
     public async discoverTests(
         testController: TestController,
         token?: CancellationToken,
@@ -216,8 +217,13 @@ export class WorkspaceTestAdapter {
 
         let rawTestData;
         try {
+            // ** First line is old way, section with if statement below is new way.
             rawTestData = await this.discoveryAdapter.discoverTests(this.workspaceUri);
-
+            // if (executionFactory !== undefined) {
+            //     rawTestData = await this.discoveryAdapter.discoverTests(this.workspaceUri, executionFactory);
+            // } else {
+            //     traceVerbose('executionFactory is undefined');
+            // }
             deferred.resolve();
         } catch (ex) {
             sendTelemetryEvent(EventName.UNITTEST_DISCOVERY_DONE, undefined, { tool: this.testProvider, failed: true });
@@ -352,6 +358,7 @@ function populateTestTree(
                 testItem.canResolveChildren = false;
                 testItem.range = range;
                 testItem.tags = [RunTestTag, DebugTestTag];
+
                 testRoot!.children.add(testItem);
                 // add to our map
                 wstAdapter.runIdToTestItem.set(child.runID, testItem);
@@ -365,7 +372,6 @@ function populateTestTree(
 
                     node.canResolveChildren = true;
                     node.tags = [RunTestTag, DebugTestTag];
-
                     testRoot!.children.add(node);
                 }
                 populateTestTree(testController, child, node, wstAdapter, token);
