allow large scale testing

.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"
             }
         },
         {

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__)))

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
 

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")

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 };

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) {

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<string, string>, rawData: string): I
         remainingRawData,
     };
 }
+export const startServer = (testIds: string): Promise<number> =>
+    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);
+        });
+    });

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<number> =>
+                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`);

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) {