fix to get user defined env and use it to set up testing subprocess (#…

…22165) fixes #21642 and #22166
microsoft · Oct 6, 2023 · e7dfef8 · e7dfef8
1 parent 514bce6
commit e7dfef8
Show file tree

Hide file tree

Showing 13 changed files with 171 additions and 94 deletions.
diff --git a/src/client/testing/testController/common/server.ts b/src/client/testing/testController/common/server.ts
@@ -23,6 +23,7 @@ import {
     extractJsonPayload,
 } from './utils';
 import { createDeferred } from '../../../common/utils/async';
+import { EnvironmentVariables } from '../../../api/types';
 
 export class PythonTestServer implements ITestServer, Disposable {
     private _onDataReceived: EventEmitter<DataReceivedEvent> = new EventEmitter<DataReceivedEvent>();
@@ -165,28 +166,29 @@ export class PythonTestServer implements ITestServer, Disposable {
 
     async sendCommand(
         options: TestCommandOptions,
+        env: EnvironmentVariables,
         runTestIdPort?: string,
         runInstance?: TestRun,
         testIds?: string[],
         callback?: () => void,
     ): Promise<void> {
         const { uuid } = options;
-
+        // get and edit env vars
+        const mutableEnv = { ...env };
         const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
         const pythonPathCommand = [options.cwd, ...pythonPathParts].join(path.delimiter);
+        mutableEnv.PYTHONPATH = pythonPathCommand;
+        mutableEnv.TEST_UUID = uuid.toString();
+        mutableEnv.TEST_PORT = this.getPort().toString();
+        mutableEnv.RUN_TEST_IDS_PORT = runTestIdPort;
+
         const spawnOptions: SpawnOptions = {
             token: options.token,
             cwd: options.cwd,
             throwOnStdErr: true,
             outputChannel: options.outChannel,
-            extraVariables: {
-                PYTHONPATH: pythonPathCommand,
-                TEST_UUID: uuid.toString(),
-                TEST_PORT: this.getPort().toString(),
-            },
+            env: mutableEnv,
         };
-
-        if (spawnOptions.extraVariables) spawnOptions.extraVariables.RUN_TEST_IDS_PORT = runTestIdPort;
         const isRun = runTestIdPort !== undefined;
         // Create the Python environment in which to execute the command.
         const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {

diff --git a/src/client/testing/testController/common/types.ts b/src/client/testing/testController/common/types.ts
@@ -15,6 +15,7 @@ import {
 import { ITestDebugLauncher, TestDiscoveryOptions } from '../../common/types';
 import { IPythonExecutionFactory } from '../../../common/process/types';
 import { Deferred } from '../../../common/utils/async';
+import { EnvironmentVariables } from '../../../common/variables/types';
 
 export type TestRunInstanceOptions = TestRunOptions & {
     exclude?: readonly TestItem[];
@@ -177,6 +178,7 @@ export interface ITestServer {
     readonly onDiscoveryDataReceived: Event<DataReceivedEvent>;
     sendCommand(
         options: TestCommandOptions,
+        env: EnvironmentVariables,
         runTestIdsPort?: string,
         runInstance?: TestRun,
         testIds?: string[],

diff --git a/src/client/testing/testController/controller.ts b/src/client/testing/testController/controller.ts
@@ -50,6 +50,7 @@ import { ITestDebugLauncher } from '../common/types';
 import { IServiceContainer } from '../../ioc/types';
 import { PythonResultResolver } from './common/resultResolver';
 import { onDidSaveTextDocument } from '../../common/vscodeApis/workspaceApis';
+import { IEnvironmentVariablesProvider } from '../../common/variables/types';
 
 // Types gymnastics to make sure that sendTriggerTelemetry only accepts the correct types.
 type EventPropertyType = IEventNamePropertyMapping[EventName.UNITTEST_DISCOVERY_TRIGGER];
@@ -100,6 +101,7 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
         @inject(ITestDebugLauncher) private readonly debugLauncher: ITestDebugLauncher,
         @inject(ITestOutputChannel) private readonly testOutputChannel: ITestOutputChannel,
         @inject(IServiceContainer) private readonly serviceContainer: IServiceContainer,
+        @inject(IEnvironmentVariablesProvider) private readonly envVarsService: IEnvironmentVariablesProvider,
     ) {
         this.refreshCancellation = new CancellationTokenSource();
 
@@ -174,12 +176,14 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                     this.configSettings,
                     this.testOutputChannel,
                     resultResolver,
+                    this.envVarsService,
                 );
                 executionAdapter = new UnittestTestExecutionAdapter(
                     this.pythonTestServer,
                     this.configSettings,
                     this.testOutputChannel,
                     resultResolver,
+                    this.envVarsService,
                 );
             } else {
                 testProvider = PYTEST_PROVIDER;
@@ -189,12 +193,14 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                     this.configSettings,
                     this.testOutputChannel,
                     resultResolver,
+                    this.envVarsService,
                 );
                 executionAdapter = new PytestTestExecutionAdapter(
                     this.pythonTestServer,
                     this.configSettings,
                     this.testOutputChannel,
                     resultResolver,
+                    this.envVarsService,
                 );
             }
 

diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -19,6 +19,7 @@ import {
     ITestServer,
 } from '../common/types';
 import { createDiscoveryErrorPayload, createEOTPayload, createTestingDeferred } from '../common/utils';
+import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 
 /**
  * Wrapper class for unittest test discovery. This is where we call `runTestCommand`. #this seems incorrectly copied
@@ -29,6 +30,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         public configSettings: IConfigurationService,
         private readonly outputChannel: ITestOutputChannel,
         private readonly resultResolver?: ITestResultResolver,
+        private readonly envVarsService?: IEnvironmentVariablesProvider,
     ) {}
 
     async discoverTests(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
@@ -61,18 +63,21 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         const { pytestArgs } = settings.testing;
         const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
 
+        // get and edit env vars
+        const mutableEnv = {
+            ...(await this.envVarsService?.getEnvironmentVariables(uri)),
+        };
         const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
         const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
+        mutableEnv.PYTHONPATH = pythonPathCommand;
+        mutableEnv.TEST_UUID = uuid.toString();
+        mutableEnv.TEST_PORT = this.testServer.getPort().toString();
 
         const spawnOptions: SpawnOptions = {
             cwd,
             throwOnStdErr: true,
-            extraVariables: {
-                PYTHONPATH: pythonPathCommand,
-                TEST_UUID: uuid.toString(),
-                TEST_PORT: this.testServer.getPort().toString(),
-            },
             outputChannel: this.outputChannel,
+            env: mutableEnv,
         };
 
         // Create the Python environment in which to execute the command.

diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -24,13 +24,15 @@ import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
 import { PYTEST_PROVIDER } from '../../common/constants';
 import { EXTENSION_ROOT_DIR } from '../../../common/constants';
 import * as utils from '../common/utils';
+import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 
 export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
     constructor(
         public testServer: ITestServer,
         public configSettings: IConfigurationService,
         private readonly outputChannel: ITestOutputChannel,
         private readonly resultResolver?: ITestResultResolver,
+        private readonly envVarsService?: IEnvironmentVariablesProvider,
     ) {}
 
     async runTests(
@@ -46,6 +48,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         const deferredTillEOT: Deferred<void> = utils.createTestingDeferred();
 
         const dataReceivedDisposable = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
+            runInstance?.token.isCancellationRequested;
             if (runInstance) {
                 const eParsed = JSON.parse(e.data);
                 this.resultResolver?.resolveExecution(eParsed, runInstance, deferredTillEOT);
@@ -105,20 +108,13 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         const settings = this.configSettings.getSettings(uri);
         const { pytestArgs } = settings.testing;
         const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
-
+        // get and edit env vars
+        const mutableEnv = { ...(await this.envVarsService?.getEnvironmentVariables(uri)) };
         const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
         const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
-        const spawnOptions: SpawnOptions = {
-            cwd,
-            throwOnStdErr: true,
-            extraVariables: {
-                PYTHONPATH: pythonPathCommand,
-                TEST_UUID: uuid.toString(),
-                TEST_PORT: this.testServer.getPort().toString(),
-            },
-            outputChannel: this.outputChannel,
-            stdinStr: testIds.toString(),
-        };
+        mutableEnv.PYTHONPATH = pythonPathCommand;
+        mutableEnv.TEST_UUID = uuid.toString();
+        mutableEnv.TEST_PORT = this.testServer.getPort().toString();
 
         // Create the Python environment in which to execute the command.
         const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
@@ -141,9 +137,17 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 testArgs.push('--capture', 'no');
             }
 
+            // add port with run test ids to env vars
             const pytestRunTestIdsPort = await utils.startTestIdServer(testIds);
-            if (spawnOptions.extraVariables)
-                spawnOptions.extraVariables.RUN_TEST_IDS_PORT = pytestRunTestIdsPort.toString();
+            mutableEnv.RUN_TEST_IDS_PORT = pytestRunTestIdsPort.toString();
+
+            const spawnOptions: SpawnOptions = {
+                cwd,
+                throwOnStdErr: true,
+                outputChannel: this.outputChannel,
+                stdinStr: testIds.toString(),
+                env: mutableEnv,
+            };
 
             if (debugBool) {
                 const pytestPort = this.testServer.getPort().toString();

diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -15,6 +15,7 @@ import {
     TestDiscoveryCommand,
 } from '../common/types';
 import { Deferred, createDeferred } from '../../../common/utils/async';
+import { EnvironmentVariables, IEnvironmentVariablesProvider } from '../../../common/variables/types';
 
 /**
  * Wrapper class for unittest test discovery. This is where we call `runTestCommand`.
@@ -25,13 +26,17 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         public configSettings: IConfigurationService,
         private readonly outputChannel: ITestOutputChannel,
         private readonly resultResolver?: ITestResultResolver,
+        private readonly envVarsService?: IEnvironmentVariablesProvider,
     ) {}
 
     public async discoverTests(uri: Uri): 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;
-
+        let env: EnvironmentVariables | undefined = await this.envVarsService?.getEnvironmentVariables(uri);
+        if (env === undefined) {
+            env = {} as EnvironmentVariables;
+        }
         const command = buildDiscoveryCommand(unittestArgs);
 
         const uuid = this.testServer.createUUID(uri.fsPath);
@@ -52,7 +57,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             dataReceivedDisposable.dispose();
         };
 
-        await this.callSendCommand(options, () => {
+        await this.callSendCommand(options, env, () => {
             disposeDataReceiver?.(this.testServer);
         });
         await deferredTillEOT.promise;
@@ -66,8 +71,12 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         return discoveryPayload;
     }
 
-    private async callSendCommand(options: TestCommandOptions, callback: () => void): Promise<DiscoveredTestPayload> {
-        await this.testServer.sendCommand(options, undefined, undefined, [], callback);
+    private async callSendCommand(
+        options: TestCommandOptions,
+        env: EnvironmentVariables,
+        callback: () => void,
+    ): Promise<DiscoveredTestPayload> {
+        await this.testServer.sendCommand(options, env, undefined, undefined, [], callback);
         const discoveryPayload: DiscoveredTestPayload = { cwd: '', status: 'success' };
         return discoveryPayload;
     }

diff --git a/src/client/testing/testController/unittest/testExecutionAdapter.ts b/src/client/testing/testController/unittest/testExecutionAdapter.ts
@@ -17,6 +17,7 @@ import {
 } from '../common/types';
 import { traceError, traceInfo, traceLog } from '../../../logging';
 import { startTestIdServer } from '../common/utils';
+import { EnvironmentVariables, IEnvironmentVariablesProvider } from '../../../common/variables/types';
 
 /**
  * Wrapper Class for unittest test execution. This is where we call `runTestCommand`?
@@ -28,6 +29,7 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         public configSettings: IConfigurationService,
         private readonly outputChannel: ITestOutputChannel,
         private readonly resultResolver?: ITestResultResolver,
+        private readonly envVarsService?: IEnvironmentVariablesProvider,
     ) {}
 
     public async runTests(
@@ -78,6 +80,10 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
 
         const command = buildExecutionCommand(unittestArgs);
+        let env: EnvironmentVariables | undefined = await this.envVarsService?.getEnvironmentVariables(uri);
+        if (env === undefined) {
+            env = {} as EnvironmentVariables;
+        }
 
         const options: TestCommandOptions = {
             workspaceFolder: uri,
@@ -92,7 +98,7 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
 
         const runTestIdsPort = await startTestIdServer(testIds);
 
-        await this.testServer.sendCommand(options, runTestIdsPort.toString(), runInstance, testIds, () => {
+        await this.testServer.sendCommand(options, env, runTestIdsPort.toString(), runInstance, testIds, () => {
             deferredTillEOT?.resolve();
         });
         // placeholder until after the rewrite is adopted