Unittest for large workspaces (#21351)

follows the same steps as making pytest compatible with large workspaces
with many tests. Now test_ids are sent over a port as a json instead of
in the exec function which can hit a cap on # of characters. Should fix
#21339.

Author: eleanorjboyd

pythonFiles/testing_tools/process_json_util.py

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+import io
+import json
+from typing import List
+
+CONTENT_LENGTH: str = "Content-Length:"
+
+
+def process_rpc_json(data: str) -> List[str]:
+    """Process the JSON data which comes from the server."""
+    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)

pythonFiles/tests/unittestadapter/test_execution.py

@@ -20,27 +20,23 @@
                 "111",
                 "--uuid",
                 "fake-uuid",
-                "--testids",
-                "test_file.test_class.test_method",
             ],
-            (111, "fake-uuid", ["test_file.test_class.test_method"]),
+            (111, "fake-uuid"),
         ),
         (
-            ["--port", "111", "--uuid", "fake-uuid", "--testids", ""],
-            (111, "fake-uuid", [""]),
+            ["--port", "111", "--uuid", "fake-uuid"],
+            (111, "fake-uuid"),
         ),
         (
             [
                 "--port",
                 "111",
                 "--uuid",
                 "fake-uuid",
-                "--testids",
-                "test_file.test_class.test_method",
                 "-v",
                 "-s",
             ],
-            (111, "fake-uuid", ["test_file.test_class.test_method"]),
+            (111, "fake-uuid"),
         ),
     ],
 )

pythonFiles/unittestadapter/execution.py

@@ -5,12 +5,19 @@
 import enum
 import json
 import os
+import pathlib
+import socket
 import sys
 import traceback
 import unittest
 from types import TracebackType
 from typing import Dict, List, Optional, Tuple, Type, 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 testing_tools import process_json_util
+
 # 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__)))
 sys.path.insert(0, PYTHON_FILES)
@@ -25,7 +32,7 @@
 
 def parse_execution_cli_args(
     args: List[str],
-) -> Tuple[int, Union[str, None], List[str]]:
+) -> Tuple[int, Union[str, None]]:
     """Parse command-line arguments that should be processed by the script.
 
     So far this includes the port number that it needs to connect to, the uuid passed by the TS side,
@@ -39,10 +46,9 @@ def parse_execution_cli_args(
     arg_parser = argparse.ArgumentParser()
     arg_parser.add_argument("--port", default=DEFAULT_PORT)
     arg_parser.add_argument("--uuid")
-    arg_parser.add_argument("--testids", nargs="+")
     parsed_args, _ = arg_parser.parse_known_args(args)
 
-    return (int(parsed_args.port), parsed_args.uuid, parsed_args.testids)
+    return (int(parsed_args.port), parsed_args.uuid)
 
 
 ErrorType = Union[
@@ -226,11 +232,62 @@ def run_tests(
 
     start_dir, pattern, top_level_dir = parse_unittest_args(argv[index + 1 :])
 
-    # Perform test execution.
-    port, uuid, testids = parse_execution_cli_args(argv[:index])
-    payload = run_tests(start_dir, testids, pattern, top_level_dir, uuid)
+    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
+    )
+
+    # get data from socket
+    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_json_util.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")
+
+    port, uuid = parse_execution_cli_args(argv[:index])
+    if test_ids_from_buffer:
+        # Perform test execution.
+        payload = run_tests(
+            start_dir, test_ids_from_buffer, pattern, top_level_dir, uuid
+        )
+    else:
+        cwd = os.path.abspath(start_dir)
+        status = TestExecutionStatus.error
+        payload: PayloadDict = {
+            "cwd": cwd,
+            "status": status,
+            "error": "No test ids received from buffer",
+        }
 
-    # Build the request data (it has to be a POST request or the Node side will not process it), and send it.
+    # Build the request data and send it.
     addr = ("localhost", port)
     data = json.dumps(payload)
     request = f"""Content-Length: {len(data)}

pythonFiles/vscode_pytest/run_pytest_script.py

@@ -1,41 +1,17 @@
 # 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)
-
+script_dir = pathlib.Path(__file__).parent.parent
+sys.path.append(os.fspath(script_dir))
+sys.path.append(os.fspath(script_dir / "lib" / "python"))
+from testing_tools import process_json_util
 
 # 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
@@ -69,7 +45,9 @@ def process_rpc_json(data: str) -> List[str]:
 
             try:
                 # Try to parse the buffer as JSON
-                test_ids_from_buffer = process_rpc_json(buffer.decode("utf-8"))
+                test_ids_from_buffer = process_json_util.process_rpc_json(
+                    buffer.decode("utf-8")
+                )
                 # Clear the buffer as complete JSON object is received
                 buffer = b""
 

src/client/testing/common/debugLauncher.ts

@@ -202,13 +202,20 @@ export class DebugLauncher implements ITestDebugLauncher {
             throw Error(`Invalid debug config "${debugConfig.name}"`);
         }
         launchArgs.request = 'launch';
+
+        // Both types of tests need to have the port for the test result server.
+        if (options.runTestIdsPort) {
+            launchArgs.env = {
+                ...launchArgs.env,
+                RUN_TEST_IDS_PORT: options.runTestIdsPort,
+            };
+        }
         if (options.testProvider === 'pytest' && pythonTestAdapterRewriteExperiment) {
             if (options.pytestPort && options.pytestUUID) {
                 launchArgs.env = {
                     ...launchArgs.env,
                     TEST_PORT: options.pytestPort,
                     TEST_UUID: options.pytestUUID,
-                    RUN_TEST_IDS_PORT: options.pytestRunTestIdsPort,
                 };
             } else {
                 throw Error(

src/client/testing/common/types.ts

@@ -27,7 +27,7 @@ export type LaunchOptions = {
     outChannel?: OutputChannel;
     pytestPort?: string;
     pytestUUID?: string;
-    pytestRunTestIdsPort?: string;
+    runTestIdsPort?: string;
 };
 
 export type ParserOptions = TestDiscoveryOptions;

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

@@ -106,17 +106,18 @@ export class PythonTestServer implements ITestServer, Disposable {
         return this._onDataReceived.event;
     }
 
-    async sendCommand(options: TestCommandOptions): Promise<void> {
+    async sendCommand(options: TestCommandOptions, runTestIdPort?: string): Promise<void> {
         const { uuid } = options;
         const spawnOptions: SpawnOptions = {
             token: options.token,
             cwd: options.cwd,
             throwOnStdErr: true,
             outputChannel: options.outChannel,
+            extraVariables: {},
         };
 
+        if (spawnOptions.extraVariables) spawnOptions.extraVariables.RUN_TEST_IDS_PORT = runTestIdPort;
         const isRun = !options.testIds;
-
         // Create the Python environment in which to execute the command.
         const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
             allowEnvironmentFetchExceptions: false,
@@ -127,23 +128,9 @@ export class PythonTestServer implements ITestServer, Disposable {
         // Add the generated UUID to the data to be sent (expecting to receive it back).
         // first check if we have testIds passed in (in case of execution) and
         // insert appropriate flag and test id array
-        let args = [];
-        if (options.testIds) {
-            args = [
-                options.command.script,
-                '--port',
-                this.getPort().toString(),
-                '--uuid',
-                uuid,
-                '--testids',
-                ...options.testIds,
-            ].concat(options.command.args);
-        } else {
-            // if not case of execution, go with the normal args
-            args = [options.command.script, '--port', this.getPort().toString(), '--uuid', uuid].concat(
-                options.command.args,
-            );
-        }
+        const args = [options.command.script, '--port', this.getPort().toString(), '--uuid', uuid].concat(
+            options.command.args,
+        );
 
         if (options.outChannel) {
             options.outChannel.appendLine(`python ${args.join(' ')}`);
@@ -156,6 +143,7 @@ export class PythonTestServer implements ITestServer, Disposable {
                     args,
                     token: options.token,
                     testProvider: UNITTEST_PROVIDER,
+                    runTestIdsPort: runTestIdPort,
                 };
                 traceInfo(`Running DEBUG unittest with arguments: ${args}\r\n`);
                 await this.debugLauncher!.launchDebugger(launchOptions);

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

@@ -172,7 +172,7 @@ export type TestCommandOptionsPytest = {
  */
 export interface ITestServer {
     readonly onDataReceived: Event<DataReceivedEvent>;
-    sendCommand(options: TestCommandOptions): Promise<void>;
+    sendCommand(options: TestCommandOptions, runTestIdsPort?: string): Promise<void>;
     serverReady(): Promise<void>;
     getPort(): number;
     createUUID(cwd: string): string;

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

@@ -168,7 +168,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                     testProvider: PYTEST_PROVIDER,
                     pytestPort,
                     pytestUUID,
-                    pytestRunTestIdsPort,
+                    runTestIdsPort: pytestRunTestIdsPort,
                 };
                 traceInfo(`Running DEBUG pytest with arguments: ${testArgs.join(' ')}\r\n`);
                 await debugLauncher!.launchDebugger(launchOptions);