Skip to content

Commit 237f82b

Browse files
authored
Fix UUID and disposing to resolve race condition (#21667)
fixes #21599 and #21507
1 parent d9e368f commit 237f82b

11 files changed

+664
-157
lines changed

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

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,11 @@
33

44
import * as net from 'net';
55
import * as crypto from 'crypto';
6-
import { Disposable, Event, EventEmitter } from 'vscode';
6+
import { Disposable, Event, EventEmitter, TestRun } from 'vscode';
77
import * as path from 'path';
88
import {
99
ExecutionFactoryCreateWithEnvironmentOptions,
10+
ExecutionResult,
1011
IPythonExecutionFactory,
1112
SpawnOptions,
1213
} from '../../../common/process/types';
@@ -15,6 +16,7 @@ import { DataReceivedEvent, ITestServer, TestCommandOptions } from './types';
1516
import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
1617
import { UNITTEST_PROVIDER } from '../../common/constants';
1718
import { jsonRPCHeaders, jsonRPCContent, JSONRPC_UUID_HEADER } from './utils';
19+
import { createDeferred } from '../../../common/utils/async';
1820

1921
export class PythonTestServer implements ITestServer, Disposable {
2022
private _onDataReceived: EventEmitter<DataReceivedEvent> = new EventEmitter<DataReceivedEvent>();
@@ -140,7 +142,12 @@ export class PythonTestServer implements ITestServer, Disposable {
140142
return this._onDataReceived.event;
141143
}
142144

143-
async sendCommand(options: TestCommandOptions, runTestIdPort?: string, callback?: () => void): Promise<void> {
145+
async sendCommand(
146+
options: TestCommandOptions,
147+
runTestIdPort?: string,
148+
runInstance?: TestRun,
149+
callback?: () => void,
150+
): Promise<void> {
144151
const { uuid } = options;
145152

146153
const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
@@ -154,7 +161,7 @@ export class PythonTestServer implements ITestServer, Disposable {
154161
};
155162

156163
if (spawnOptions.extraVariables) spawnOptions.extraVariables.RUN_TEST_IDS_PORT = runTestIdPort;
157-
const isRun = !options.testIds;
164+
const isRun = runTestIdPort !== undefined;
158165
// Create the Python environment in which to execute the command.
159166
const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
160167
allowEnvironmentFetchExceptions: false,
@@ -195,7 +202,19 @@ export class PythonTestServer implements ITestServer, Disposable {
195202
// This means it is running discovery
196203
traceLog(`Discovering unittest tests with arguments: ${args}\r\n`);
197204
}
198-
await execService.exec(args, spawnOptions);
205+
const deferred = createDeferred<ExecutionResult<string>>();
206+
207+
const result = execService.execObservable(args, spawnOptions);
208+
209+
runInstance?.token.onCancellationRequested(() => {
210+
result?.proc?.kill();
211+
});
212+
result?.proc?.on('close', () => {
213+
traceLog('Exec server closed.', uuid);
214+
deferred.resolve({ stdout: '', stderr: '' });
215+
callback?.();
216+
});
217+
await deferred.promise;
199218
}
200219
} catch (ex) {
201220
this.uuids = this.uuids.filter((u) => u !== uuid);

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,12 @@ export interface ITestServer {
174174
readonly onDataReceived: Event<DataReceivedEvent>;
175175
readonly onRunDataReceived: Event<DataReceivedEvent>;
176176
readonly onDiscoveryDataReceived: Event<DataReceivedEvent>;
177-
sendCommand(options: TestCommandOptions, runTestIdsPort?: string, callback?: () => void): Promise<void>;
177+
sendCommand(
178+
options: TestCommandOptions,
179+
runTestIdsPort?: string,
180+
runInstance?: TestRun,
181+
callback?: () => void,
182+
): Promise<void>;
178183
serverReady(): Promise<void>;
179184
getPort(): number;
180185
createUUID(cwd: string): string;

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

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,14 @@ import * as path from 'path';
44
import { Uri } from 'vscode';
55
import {
66
ExecutionFactoryCreateWithEnvironmentOptions,
7+
ExecutionResult,
78
IPythonExecutionFactory,
89
SpawnOptions,
910
} from '../../../common/process/types';
1011
import { IConfigurationService, ITestOutputChannel } from '../../../common/types';
1112
import { createDeferred } from '../../../common/utils/async';
1213
import { EXTENSION_ROOT_DIR } from '../../../constants';
13-
import { traceError, traceVerbose } from '../../../logging';
14+
import { traceVerbose } from '../../../logging';
1415
import {
1516
DataReceivedEvent,
1617
DiscoveredTestPayload,
@@ -48,7 +49,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
4849
return discoveryPayload;
4950
}
5051

51-
async runPytestDiscovery(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
52+
async runPytestDiscovery(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<void> {
5253
const deferred = createDeferred<DiscoveredTestPayload>();
5354
const relativePathToPytest = 'pythonFiles';
5455
const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
@@ -78,17 +79,15 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
7879
};
7980
const execService = await executionFactory?.createActivatedEnvironment(creationOptions);
8081
// delete UUID following entire discovery finishing.
81-
execService
82-
?.exec(['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs), spawnOptions)
83-
.then(() => {
84-
this.testServer.deleteUUID(uuid);
85-
return deferred.resolve();
86-
})
87-
.catch((err) => {
88-
traceError(`Error while trying to run pytest discovery, \n${err}\r\n\r\n`);
89-
this.testServer.deleteUUID(uuid);
90-
return deferred.reject(err);
91-
});
92-
return deferred.promise;
82+
const deferredExec = createDeferred<ExecutionResult<string>>();
83+
const execArgs = ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs);
84+
const result = execService?.execObservable(execArgs, spawnOptions);
85+
86+
result?.proc?.on('close', () => {
87+
deferredExec.resolve({ stdout: '', stderr: '' });
88+
this.testServer.deleteUUID(uuid);
89+
deferred.resolve();
90+
});
91+
await deferredExec.promise;
9392
}
9493
}

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

Lines changed: 28 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,15 @@ import {
1515
} from '../common/types';
1616
import {
1717
ExecutionFactoryCreateWithEnvironmentOptions,
18+
ExecutionResult,
1819
IPythonExecutionFactory,
1920
SpawnOptions,
2021
} from '../../../common/process/types';
2122
import { removePositionalFoldersAndFiles } from './arguments';
2223
import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
2324
import { PYTEST_PROVIDER } from '../../common/constants';
2425
import { EXTENSION_ROOT_DIR } from '../../../common/constants';
25-
import { startTestIdServer } from '../common/utils';
26-
27-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
28-
// (global as any).EXTENSION_ROOT_DIR = EXTENSION_ROOT_DIR;
29-
/**
30-
* Wrapper Class for pytest test execution..
31-
*/
26+
import * as utils from '../common/utils';
3227

3328
export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
3429
constructor(
@@ -48,18 +43,20 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
4843
): Promise<ExecutionTestPayload> {
4944
const uuid = this.testServer.createUUID(uri.fsPath);
5045
traceVerbose(uri, testIds, debugBool);
51-
const disposable = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
46+
const disposedDataReceived = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
5247
if (runInstance) {
5348
this.resultResolver?.resolveExecution(JSON.parse(e.data), runInstance);
5449
}
5550
});
56-
try {
57-
await this.runTestsNew(uri, testIds, uuid, debugBool, executionFactory, debugLauncher);
58-
} finally {
59-
this.testServer.deleteUUID(uuid);
60-
disposable.dispose();
61-
// confirm with testing that this gets called (it must clean this up)
62-
}
51+
const dispose = function (testServer: ITestServer) {
52+
testServer.deleteUUID(uuid);
53+
disposedDataReceived.dispose();
54+
};
55+
runInstance?.token.onCancellationRequested(() => {
56+
dispose(this.testServer);
57+
});
58+
await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, executionFactory, debugLauncher);
59+
6360
// placeholder until after the rewrite is adopted
6461
// TODO: remove after adoption.
6562
const executionPayload: ExecutionTestPayload = {
@@ -74,6 +71,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
7471
uri: Uri,
7572
testIds: string[],
7673
uuid: string,
74+
runInstance?: TestRun,
7775
debugBool?: boolean,
7876
executionFactory?: IPythonExecutionFactory,
7977
debugLauncher?: ITestDebugLauncher,
@@ -124,7 +122,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
124122
}
125123
traceLog(`Running PYTEST execution for the following test ids: ${testIds}`);
126124

127-
const pytestRunTestIdsPort = await startTestIdServer(testIds);
125+
const pytestRunTestIdsPort = await utils.startTestIdServer(testIds);
128126
if (spawnOptions.extraVariables)
129127
spawnOptions.extraVariables.RUN_TEST_IDS_PORT = pytestRunTestIdsPort.toString();
130128

@@ -143,14 +141,27 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
143141
traceInfo(`Running DEBUG pytest with arguments: ${testArgs.join(' ')}\r\n`);
144142
await debugLauncher!.launchDebugger(launchOptions, () => {
145143
deferred.resolve();
144+
this.testServer.deleteUUID(uuid);
146145
});
147146
} else {
148147
// combine path to run script with run args
149148
const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py');
150149
const runArgs = [scriptPath, ...testArgs];
151150
traceInfo(`Running pytests with arguments: ${runArgs.join(' ')}\r\n`);
152151

153-
await execService?.exec(runArgs, spawnOptions);
152+
const deferredExec = createDeferred<ExecutionResult<string>>();
153+
const result = execService?.execObservable(runArgs, spawnOptions);
154+
155+
runInstance?.token.onCancellationRequested(() => {
156+
result?.proc?.kill();
157+
});
158+
159+
result?.proc?.on('close', () => {
160+
deferredExec.resolve({ stdout: '', stderr: '' });
161+
this.testServer.deleteUUID(uuid);
162+
deferred.resolve();
163+
});
164+
await deferredExec.promise;
154165
}
155166
} catch (ex) {
156167
traceError(`Error while running tests: ${testIds}\r\n${ex}\r\n\r\n`);

src/client/testing/testController/unittest/testDiscoveryAdapter.ts

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,12 +46,11 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
4646
const disposable = this.testServer.onDiscoveryDataReceived((e: DataReceivedEvent) => {
4747
this.resultResolver?.resolveDiscovery(JSON.parse(e.data));
4848
});
49-
try {
50-
await this.callSendCommand(options);
51-
} finally {
49+
50+
await this.callSendCommand(options, () => {
5251
this.testServer.deleteUUID(uuid);
5352
disposable.dispose();
54-
}
53+
});
5554
// placeholder until after the rewrite is adopted
5655
// TODO: remove after adoption.
5756
const discoveryPayload: DiscoveredTestPayload = {
@@ -61,8 +60,8 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
6160
return discoveryPayload;
6261
}
6362

64-
private async callSendCommand(options: TestCommandOptions): Promise<DiscoveredTestPayload> {
65-
await this.testServer.sendCommand(options);
63+
private async callSendCommand(options: TestCommandOptions, callback: () => void): Promise<DiscoveredTestPayload> {
64+
await this.testServer.sendCommand(options, undefined, undefined, callback);
6665
const discoveryPayload: DiscoveredTestPayload = { cwd: '', status: 'success' };
6766
return discoveryPayload;
6867
}

src/client/testing/testController/unittest/testExecutionAdapter.ts

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -37,18 +37,19 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
3737
runInstance?: TestRun,
3838
): Promise<ExecutionTestPayload> {
3939
const uuid = this.testServer.createUUID(uri.fsPath);
40-
const disposable = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
40+
const disposedDataReceived = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
4141
if (runInstance) {
4242
this.resultResolver?.resolveExecution(JSON.parse(e.data), runInstance);
4343
}
4444
});
45-
try {
46-
await this.runTestsNew(uri, testIds, uuid, debugBool);
47-
} finally {
45+
const dispose = function () {
46+
disposedDataReceived.dispose();
47+
};
48+
runInstance?.token.onCancellationRequested(() => {
4849
this.testServer.deleteUUID(uuid);
49-
disposable.dispose();
50-
// confirm with testing that this gets called (it must clean this up)
51-
}
50+
dispose();
51+
});
52+
await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, dispose);
5253
const executionPayload: ExecutionTestPayload = { cwd: uri.fsPath, status: 'success', error: '' };
5354
return executionPayload;
5455
}
@@ -57,7 +58,9 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
5758
uri: Uri,
5859
testIds: string[],
5960
uuid: string,
61+
runInstance?: TestRun,
6062
debugBool?: boolean,
63+
dispose?: () => void,
6164
): Promise<ExecutionTestPayload> {
6265
const settings = this.configSettings.getSettings(uri);
6366
const { unittestArgs } = settings.testing;
@@ -80,8 +83,10 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
8083

8184
const runTestIdsPort = await startTestIdServer(testIds);
8285

83-
await this.testServer.sendCommand(options, runTestIdsPort.toString(), () => {
86+
await this.testServer.sendCommand(options, runTestIdsPort.toString(), runInstance, () => {
87+
this.testServer.deleteUUID(uuid);
8488
deferred.resolve();
89+
dispose?.();
8590
});
8691
// placeholder until after the rewrite is adopted
8792
// TODO: remove after adoption.

0 commit comments

Comments
 (0)