From e2a9cecf9cec0f3a42d2723927a83d486c9b5bd8 Mon Sep 17 00:00:00 2001 From: Eleanor Boyd Date: Wed, 24 May 2023 16:42:24 -0700 Subject: [PATCH] allow large scale testing (#21269) allows new testing rewrite to handle 500+ tests and load and run these tests. High limit tested was 10,000 tests. --- .vscode/launch.json | 3 +- pythonFiles/unittestadapter/discovery.py | 8 +- pythonFiles/unittestadapter/utils.py | 9 +- .../vscode_pytest/run_pytest_script.py | 90 +++++++++++++++++++ src/client/common/process/types.ts | 1 + .../testing/testController/common/server.ts | 32 ++++--- .../testing/testController/common/utils.ts | 37 ++++++++ .../pytest/pytestExecutionAdapter.ts | 54 +++++++++-- .../testController/workspaceTestAdapter.ts | 5 +- 9 files changed, 217 insertions(+), 22 deletions(-) create mode 100644 pythonFiles/vscode_pytest/run_pytest_script.py diff --git a/.vscode/launch.json b/.vscode/launch.json index 1ca0db3dc858..82981a93305d 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -22,8 +22,7 @@ // Enable this to log telemetry to the output during debugging "XVSC_PYTHON_LOG_TELEMETRY": "1", // Enable this to log debugger output. Directory must exist ahead of time - "XDEBUGPY_LOG_DIR": "${workspaceRoot}/tmp/Debug_Output_Ex", - "ENABLE_PYTHON_TESTING_REWRITE": "1" + "XDEBUGPY_LOG_DIR": "${workspaceRoot}/tmp/Debug_Output_Ex" } }, { diff --git a/pythonFiles/unittestadapter/discovery.py b/pythonFiles/unittestadapter/discovery.py index dc0a139ed5a2..bcc2fd967f78 100644 --- a/pythonFiles/unittestadapter/discovery.py +++ b/pythonFiles/unittestadapter/discovery.py @@ -8,7 +8,13 @@ import sys import traceback import unittest -from typing import List, Literal, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union + +script_dir = pathlib.Path(__file__).parent.parent +sys.path.append(os.fspath(script_dir)) +sys.path.append(os.fspath(script_dir / "lib" / "python")) + +from typing_extensions import Literal # Add the path to pythonFiles to sys.path to find testing_tools.socket_manager. PYTHON_FILES = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) diff --git a/pythonFiles/unittestadapter/utils.py b/pythonFiles/unittestadapter/utils.py index 568ff30ee92d..9c8b896a8d6e 100644 --- a/pythonFiles/unittestadapter/utils.py +++ b/pythonFiles/unittestadapter/utils.py @@ -6,8 +6,15 @@ import inspect import os import pathlib +import sys import unittest -from typing import List, Tuple, TypedDict, Union +from typing import List, Tuple, Union + +script_dir = pathlib.Path(__file__).parent.parent +sys.path.append(os.fspath(script_dir)) +sys.path.append(os.fspath(script_dir / "lib" / "python")) + +from typing_extensions import TypedDict # Types diff --git a/pythonFiles/vscode_pytest/run_pytest_script.py b/pythonFiles/vscode_pytest/run_pytest_script.py new file mode 100644 index 000000000000..f6d6bdcafd5f --- /dev/null +++ b/pythonFiles/vscode_pytest/run_pytest_script.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +import io +import json +import os +import pathlib +import socket +import sys +from typing import List + +import pytest + +CONTENT_LENGTH: str = "Content-Length:" + + +def process_rpc_json(data: str) -> List[str]: + """Process the JSON data which comes from the server which runs the pytest discovery.""" + str_stream: io.StringIO = io.StringIO(data) + + length: int = 0 + + while True: + line: str = str_stream.readline() + if CONTENT_LENGTH.lower() in line.lower(): + length = int(line[len(CONTENT_LENGTH) :]) + break + + if not line or line.isspace(): + raise ValueError("Header does not contain Content-Length") + + while True: + line: str = str_stream.readline() + if not line or line.isspace(): + break + + raw_json: str = str_stream.read(length) + return json.loads(raw_json) + + +# This script handles running pytest via pytest.main(). It is called via run in the +# pytest execution adapter and gets the test_ids to run via stdin and the rest of the +# args through sys.argv. It then runs pytest.main() with the args and test_ids. + +if __name__ == "__main__": + # Add the root directory to the path so that we can import the plugin. + directory_path = pathlib.Path(__file__).parent.parent + sys.path.append(os.fspath(directory_path)) + # Get the rest of the args to run with pytest. + args = sys.argv[1:] + run_test_ids_port = os.environ.get("RUN_TEST_IDS_PORT") + run_test_ids_port_int = ( + int(run_test_ids_port) if run_test_ids_port is not None else 0 + ) + test_ids_from_buffer = [] + try: + client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + client_socket.connect(("localhost", run_test_ids_port_int)) + print(f"CLIENT: Server listening on port {run_test_ids_port_int}...") + buffer = b"" + + while True: + # Receive the data from the client + data = client_socket.recv(1024 * 1024) + if not data: + break + + # Append the received data to the buffer + buffer += data + + try: + # Try to parse the buffer as JSON + test_ids_from_buffer = process_rpc_json(buffer.decode("utf-8")) + # Clear the buffer as complete JSON object is received + buffer = b"" + + # Process the JSON data + print(f"Received JSON data: {test_ids_from_buffer}") + break + except json.JSONDecodeError: + # JSON decoding error, the complete JSON object is not yet received + continue + except socket.error as e: + print(f"Error: Could not connect to runTestIdsPort: {e}") + print("Error: Could not connect to runTestIdsPort") + try: + if test_ids_from_buffer: + arg_array = ["-p", "vscode_pytest"] + args + test_ids_from_buffer + pytest.main(arg_array) + except json.JSONDecodeError: + print("Error: Could not parse test ids from stdin") diff --git a/src/client/common/process/types.ts b/src/client/common/process/types.ts index 8298957285e8..62e787b694b5 100644 --- a/src/client/common/process/types.ts +++ b/src/client/common/process/types.ts @@ -25,6 +25,7 @@ export type SpawnOptions = ChildProcessSpawnOptions & { throwOnStdErr?: boolean; extraVariables?: NodeJS.ProcessEnv; outputChannel?: OutputChannel; + stdinStr?: string; }; export type ShellOptions = ExecOptions & { throwOnStdErr?: boolean }; diff --git a/src/client/testing/testController/common/server.ts b/src/client/testing/testController/common/server.ts index 6849f0f8969a..a00623aa33c7 100644 --- a/src/client/testing/testController/common/server.ts +++ b/src/client/testing/testController/common/server.ts @@ -26,23 +26,35 @@ export class PythonTestServer implements ITestServer, Disposable { constructor(private executionFactory: IPythonExecutionFactory, private debugLauncher: ITestDebugLauncher) { this.server = net.createServer((socket: net.Socket) => { + let buffer: Buffer = Buffer.alloc(0); // Buffer to accumulate received data socket.on('data', (data: Buffer) => { try { let rawData: string = data.toString(); - - while (rawData.length > 0) { - const rpcHeaders = jsonRPCHeaders(rawData); + buffer = Buffer.concat([buffer, data]); + while (buffer.length > 0) { + const rpcHeaders = jsonRPCHeaders(buffer.toString()); const uuid = rpcHeaders.headers.get(JSONRPC_UUID_HEADER); + const totalContentLength = rpcHeaders.headers.get('Content-Length'); + if (!uuid) { + traceLog('On data received: Error occurred because payload UUID is undefined'); + this._onDataReceived.fire({ uuid: '', data: '' }); + return; + } + if (!this.uuids.includes(uuid)) { + traceLog('On data received: Error occurred because the payload UUID is not recognized'); + this._onDataReceived.fire({ uuid: '', data: '' }); + return; + } rawData = rpcHeaders.remainingRawData; - if (uuid && this.uuids.includes(uuid)) { - const rpcContent = jsonRPCContent(rpcHeaders.headers, rawData); - rawData = rpcContent.remainingRawData; - this._onDataReceived.fire({ uuid, data: rpcContent.extractedJSON }); + const rpcContent = jsonRPCContent(rpcHeaders.headers, rawData); + const extractedData = rpcContent.extractedJSON; + if (extractedData.length === Number(totalContentLength)) { + // do not send until we have the full content + this._onDataReceived.fire({ uuid, data: extractedData }); this.uuids = this.uuids.filter((u) => u !== uuid); + buffer = Buffer.alloc(0); } else { - traceLog(`Error processing test server request: uuid not found`); - this._onDataReceived.fire({ uuid: '', data: '' }); - return; + break; } } } catch (ex) { diff --git a/src/client/testing/testController/common/utils.ts b/src/client/testing/testController/common/utils.ts index e0bad383d695..88e3450d35dc 100644 --- a/src/client/testing/testController/common/utils.ts +++ b/src/client/testing/testController/common/utils.ts @@ -1,5 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +import * as net from 'net'; +import { traceLog } from '../../../logging'; export function fixLogLines(content: string): string { const lines = content.split(/\r?\n/g); @@ -50,3 +52,38 @@ export function jsonRPCContent(headers: Map, rawData: string): I remainingRawData, }; } +export const startServer = (testIds: string): Promise => + new Promise((resolve, reject) => { + const server = net.createServer((socket: net.Socket) => { + // Convert the test_ids array to JSON + const testData = JSON.stringify(testIds); + + // Create the headers + const headers = [`Content-Length: ${Buffer.byteLength(testData)}`, 'Content-Type: application/json']; + + // Create the payload by concatenating the headers and the test data + const payload = `${headers.join('\r\n')}\r\n\r\n${testData}`; + + // Send the payload to the socket + socket.write(payload); + + // Handle socket events + socket.on('data', (data) => { + traceLog('Received data:', data.toString()); + }); + + socket.on('end', () => { + traceLog('Client disconnected'); + }); + }); + + server.listen(0, () => { + const { port } = server.address() as net.AddressInfo; + traceLog(`Server listening on port ${port}`); + resolve(port); + }); + + server.on('error', (error: Error) => { + reject(error); + }); + }); diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts index 623fd1ff3a8c..4c6dcd9cdbee 100644 --- a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts +++ b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts @@ -3,9 +3,10 @@ import { Uri } from 'vscode'; import * as path from 'path'; +import * as net from 'net'; import { IConfigurationService, ITestOutputChannel } from '../../../common/types'; import { createDeferred, Deferred } from '../../../common/utils/async'; -import { traceVerbose } from '../../../logging'; +import { traceLog, traceVerbose } from '../../../logging'; import { DataReceivedEvent, ExecutionTestPayload, ITestExecutionAdapter, ITestServer } from '../common/types'; import { ExecutionFactoryCreateWithEnvironmentOptions, @@ -90,6 +91,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter { TEST_PORT: this.testServer.getPort().toString(), }, outputChannel: this.outputChannel, + stdinStr: testIds.toString(), }; // Create the Python environment in which to execute the command. @@ -114,7 +116,48 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter { if (debugBool && !testArgs.some((a) => a.startsWith('--capture') || a === '-s')) { testArgs.push('--capture', 'no'); } - const pluginArgs = ['-p', 'vscode_pytest', '-v'].concat(testArgs).concat(testIds); + const pluginArgs = ['-p', 'vscode_pytest'].concat(testArgs).concat(testIds); + const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py'); + const runArgs = [scriptPath, ...testArgs]; + + const testData = JSON.stringify(testIds); + const headers = [`Content-Length: ${Buffer.byteLength(testData)}`, 'Content-Type: application/json']; + const payload = `${headers.join('\r\n')}\r\n\r\n${testData}`; + + const startServer = (): Promise => + new Promise((resolve, reject) => { + const server = net.createServer((socket: net.Socket) => { + socket.on('end', () => { + traceLog('Client disconnected'); + }); + }); + + server.listen(0, () => { + const { port } = server.address() as net.AddressInfo; + traceLog(`Server listening on port ${port}`); + resolve(port); + }); + + server.on('error', (error: Error) => { + reject(error); + }); + server.on('connection', (socket: net.Socket) => { + socket.write(payload); + traceLog('payload sent', payload); + }); + }); + + // Start the server and wait until it is listening + await startServer() + .then((assignedPort) => { + traceLog(`Server started and listening on port ${assignedPort}`); + if (spawnOptions.extraVariables) + spawnOptions.extraVariables.RUN_TEST_IDS_PORT = assignedPort.toString(); + }) + .catch((error) => { + console.error('Error starting server:', error); + }); + if (debugBool) { const pytestPort = this.testServer.getPort().toString(); const pytestUUID = uuid.toString(); @@ -129,9 +172,10 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter { console.debug(`Running debug test with arguments: ${pluginArgs.join(' ')}\r\n`); await debugLauncher!.launchDebugger(launchOptions); } else { - const runArgs = ['-m', 'pytest'].concat(pluginArgs); - console.debug(`Running test with arguments: ${runArgs.join(' ')}\r\n`); - execService?.exec(runArgs, spawnOptions); + await execService?.exec(runArgs, spawnOptions).catch((ex) => { + console.debug(`Error while running tests: ${testIds}\r\n${ex}\r\n\r\n`); + return Promise.reject(ex); + }); } } catch (ex) { console.debug(`Error while running tests: ${testIds}\r\n${ex}\r\n\r\n`); diff --git a/src/client/testing/testController/workspaceTestAdapter.ts b/src/client/testing/testController/workspaceTestAdapter.ts index b22fee69d295..0b19d4d87d6f 100644 --- a/src/client/testing/testController/workspaceTestAdapter.ts +++ b/src/client/testing/testController/workspaceTestAdapter.ts @@ -348,12 +348,11 @@ export class WorkspaceTestAdapter { const testingErrorConst = this.testProvider === 'pytest' ? Testing.errorPytestDiscovery : Testing.errorUnittestDiscovery; const { errors } = rawTestData; - traceError(testingErrorConst, '\r\n', errors!.join('\r\n\r\n')); - + traceError(testingErrorConst, '\r\n', errors?.join('\r\n\r\n')); let errorNode = testController.items.get(`DiscoveryError:${workspacePath}`); const message = util.format( `${testingErrorConst} ${Testing.seePythonOutput}\r\n`, - errors!.join('\r\n\r\n'), + errors?.join('\r\n\r\n'), ); if (errorNode === undefined) {