switch to using temp file for test_ids (#24054)

eleanorjboyd · karthiknadig · web-flow · commit ef7c7e3aedd2 · 2024-09-05T13:02:38.000-07:00
first step in work on #23279 --------- Co-authored-by: Karthik Nadig <kanadig@microsoft.com>
diff --git a/python_files/tests/unittestadapter/conftest.py b/python_files/tests/unittestadapter/conftest.py
@@ -1,8 +1,2 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
-
-import sys
-
-# Ignore the contents of this folder for Python 2 tests.
-if sys.version_info[0] < 3:
-    collect_ignore_glob = ["*.py"]
diff --git a/python_files/unittestadapter/execution.py b/python_files/unittestadapter/execution.py
@@ -3,7 +3,6 @@
 
 import atexit
 import enum
-import json
 import os
 import pathlib
 import sys
@@ -24,7 +23,6 @@
 
 from django_handler import django_execution_runner  # noqa: E402
 
-from testing_tools import process_json_util, socket_manager  # noqa: E402
 from unittestadapter.pvsc_utils import (  # noqa: E402
     EOTPayloadDict,
     ExecutionPayloadDict,
@@ -269,8 +267,15 @@ def run_tests(
     return payload
 
 
+def execute_eot_and_cleanup():
+    eot_payload: EOTPayloadDict = {"command_type": "execution", "eot": True}
+    send_post_request(eot_payload, test_run_pipe)
+    if __socket:
+        __socket.close()
+
+
 __socket = None
-atexit.register(lambda: __socket.close() if __socket else None)
+atexit.register(execute_eot_and_cleanup)
 
 
 def send_run_data(raw_data, test_run_pipe):
@@ -306,70 +311,43 @@ def send_run_data(raw_data, test_run_pipe):
     if not test_run_pipe:
         print("Error[vscode-unittest]: TEST_RUN_PIPE env var is not set.")
         raise VSCodeUnittestError("Error[vscode-unittest]: TEST_RUN_PIPE env var is not set.")
-    test_ids_from_buffer = []
-    raw_json = None
-    try:
-        with socket_manager.PipeManager(run_test_ids_pipe) as sock:
-            buffer: str = ""
-            while True:
-                # Receive the data from the client
-                data: str = sock.read()
-                if not data:
-                    break
-
-                # Append the received data to the buffer
-                buffer += data
-
-                try:
-                    # Try to parse the buffer as JSON
-                    raw_json = process_json_util.process_rpc_json(buffer)
-                    # Clear the buffer as complete JSON object is received
-                    buffer = ""
-                    print("Received JSON data in run")
-                    break
-                except json.JSONDecodeError:
-                    # JSON decoding error, the complete JSON object is not yet received
-                    continue
-    except OSError as e:
-        msg = f"Error: Could not connect to RUN_TEST_IDS_PIPE: {e}"
-        print(msg)
-        raise VSCodeUnittestError(msg) from e
-
+    test_ids = []
     try:
-        if raw_json and "params" in raw_json and raw_json["params"]:
-            test_ids_from_buffer = raw_json["params"]
-            # Check to see if we are running django tests.
-            if manage_py_path := os.environ.get("MANAGE_PY_PATH"):
-                args = argv[index + 1 :] or []
-                django_execution_runner(manage_py_path, test_ids_from_buffer, args)
-                # the django run subprocesses sends the eot payload.
-            else:
-                # Perform test execution.
-                payload = run_tests(
-                    start_dir,
-                    test_ids_from_buffer,
-                    pattern,
-                    top_level_dir,
-                    verbosity,
-                    failfast,
-                    locals_,
-                )
-                eot_payload: EOTPayloadDict = {"command_type": "execution", "eot": True}
-                send_post_request(eot_payload, test_run_pipe)
-        else:
-            # No test ids received from buffer
-            cwd = os.path.abspath(start_dir)  # noqa: PTH100
-            status = TestExecutionStatus.error
-            payload: ExecutionPayloadDict = {
-                "cwd": cwd,
-                "status": status,
-                "error": "No test ids received from buffer",
-                "result": None,
-            }
-            send_post_request(payload, test_run_pipe)
-            eot_payload: EOTPayloadDict = {"command_type": "execution", "eot": True}
-            send_post_request(eot_payload, test_run_pipe)
-    except json.JSONDecodeError as exc:
-        msg = "Error: Could not parse test ids from stdin"
-        print(msg)
-        raise VSCodeUnittestError(msg) from exc
+        # Read the test ids from the file, attempt to delete file afterwords.
+        ids_path = pathlib.Path(run_test_ids_pipe)
+        test_ids = ids_path.read_text(encoding="utf-8").splitlines()
+        print("Received test ids from temp file.")
+        try:
+            ids_path.unlink()
+        except Exception as e:
+            print("Error[vscode-pytest]: unable to delete temp file" + str(e))
+
+    except Exception as e:
+        # No test ids received from buffer, return error payload
+        cwd = pathlib.Path(start_dir).absolute()
+        status: TestExecutionStatus = TestExecutionStatus.error
+        payload: ExecutionPayloadDict = {
+            "cwd": str(cwd),
+            "status": status,
+            "result": None,
+            "error": "No test ids read from temp file," + str(e),
+        }
+        send_post_request(payload, test_run_pipe)
+
+    # If no error occurred, we will have test ids to run.
+    if manage_py_path := os.environ.get("MANAGE_PY_PATH"):
+        print("MANAGE_PY_PATH env var set, running Django test suite.")
+        args = argv[index + 1 :] or []
+        django_execution_runner(manage_py_path, test_ids, args)
+        # the django run subprocesses sends the eot payload.
+    else:
+        # Perform regular unittest execution.
+        payload = run_tests(
+            start_dir,
+            test_ids,
+            pattern,
+            top_level_dir,
+            verbosity,
+            failfast,
+            locals_,
+        )
diff --git a/python_files/vscode_pytest/run_pytest_script.py b/python_files/vscode_pytest/run_pytest_script.py
@@ -1,6 +1,5 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-import json
 import os
 import pathlib
 import sys
@@ -17,10 +16,12 @@
 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 (  # noqa: E402
-    process_json_util,
-    socket_manager,
-)
+
+
+def run_pytest(args):
+    arg_array = ["-p", "vscode_pytest", *args]
+    pytest.main(arg_array)
+
 
 # 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
@@ -34,52 +35,21 @@
     # Get the rest of the args to run with pytest.
     args = sys.argv[1:]
     run_test_ids_pipe = os.environ.get("RUN_TEST_IDS_PIPE")
-    if not run_test_ids_pipe:
-        print("Error[vscode-pytest]: RUN_TEST_IDS_PIPE env var is not set.")
-    raw_json = {}
-    try:
-        socket_name = os.environ.get("RUN_TEST_IDS_PIPE")
-        with socket_manager.PipeManager(socket_name) as sock:
-            buffer = ""
-            while True:
-                # Receive the data from the client as a string
-                data = sock.read(3000)
-                if not data:
-                    break
-
-                # Append the received data to the buffer
-                buffer += data
-
-                try:
-                    # Try to parse the buffer as JSON
-                    raw_json = process_json_util.process_rpc_json(buffer)
-                    # Clear the buffer as complete JSON object is received
-                    buffer = ""
-                    print("Received JSON data in run script")
-                    break
-                except json.JSONDecodeError:
-                    # JSON decoding error, the complete JSON object is not yet received
-                    continue
-                except UnicodeDecodeError:
-                    continue
-    except OSError as e:
-        print(f"Error: Could not connect to runTestIdsPort: {e}")
-        print("Error: Could not connect to runTestIdsPort")
-    try:
-        test_ids_from_buffer = raw_json.get("params")
-        if test_ids_from_buffer:
-            arg_array = ["-p", "vscode_pytest", *args, *test_ids_from_buffer]
+    if run_test_ids_pipe:
+        try:
+            # Read the test ids from the file, delete file, and run pytest.
+            ids_path = pathlib.Path(run_test_ids_pipe)
+            ids = ids_path.read_text(encoding="utf-8").splitlines()
+            try:
+                ids_path.unlink()
+            except Exception as e:
+                print("Error[vscode-pytest]: unable to delete temp file" + str(e))
+            arg_array = ["-p", "vscode_pytest", *args, *ids]
             print("Running pytest with args: " + str(arg_array))
             pytest.main(arg_array)
-        else:
-            print(
-                "Error: No test ids received from stdin, could be an error or a run request without ids provided.",
-            )
-            print("Running pytest with no test ids as args. Args being used: ", args)
-            arg_array = ["-p", "vscode_pytest", *args]
-            pytest.main(arg_array)
-    except json.JSONDecodeError:
-        print(
-            "Error: Could not parse test ids from stdin. Raw json received from socket: \n",
-            raw_json,
-        )
+        except Exception as e:
+            print("Error[vscode-pytest]: unable to read testIds from temp file" + str(e))
+            run_pytest(args)
+    else:
+        print("Error[vscode-pytest]: RUN_TEST_IDS_PIPE env var is not set.")
+        run_pytest(args)
diff --git a/src/client/testing/testController/common/utils.ts b/src/client/testing/testController/common/utils.ts
@@ -3,6 +3,8 @@
 import * as net from 'net';
 import * as path from 'path';
 import * as fs from 'fs';
+import * as os from 'os';
+import * as crypto from 'crypto';
 import { CancellationToken, Position, TestController, TestItem, Uri, Range, Disposable } from 'vscode';
 import { Message } from 'vscode-jsonrpc';
 import { traceError, traceInfo, traceLog, traceVerbose } from '../../../logging';
@@ -20,6 +22,7 @@ import {
 } from './types';
 import { Deferred, createDeferred } from '../../../common/utils/async';
 import { createNamedPipeServer, generateRandomPipeName } from '../../../common/pipes/namedPipes';
+import { EXTENSION_ROOT_DIR } from '../../../constants';
 
 export function fixLogLines(content: string): string {
     const lines = content.split(/\r?\n/g);
@@ -193,6 +196,36 @@ interface ExecutionResultMessage extends Message {
     params: ExecutionTestPayload | EOTTestPayload;
 }
 
+/**
+ * Writes an array of test IDs to a temporary file.
+ *
+ * @param testIds - The array of test IDs to write.
+ * @returns A promise that resolves to the file name of the temporary file.
+ */
+export async function writeTestIdsFile(testIds: string[]): Promise<string> {
+    // temp file name in format of test-ids-<randomSuffix>.txt
+    const randomSuffix = crypto.randomBytes(10).toString('hex');
+    const tempName = `test-ids-${randomSuffix}.txt`;
+    // create temp file
+    let tempFileName: string;
+    try {
+        traceLog('Attempting to use temp directory for test ids file, file name:', tempName);
+        tempFileName = path.join(os.tmpdir(), tempName);
+    } catch (error) {
+        // Handle the error when accessing the temp directory
+        traceError('Error accessing temp directory:', error, ' Attempt to use extension root dir instead');
+        // Make new temp directory in extension root dir
+        const tempDir = path.join(EXTENSION_ROOT_DIR, '.temp');
+        await fs.promises.mkdir(tempDir, { recursive: true });
+        tempFileName = path.join(EXTENSION_ROOT_DIR, '.temp', tempName);
+        traceLog('New temp file:', tempFileName);
+    }
+    // write test ids to file
+    await fs.promises.writeFile(tempFileName, testIds.join('\n'));
+    // return file name
+    return tempFileName;
+}
+
 export async function startRunResultNamedPipe(
     dataReceivedCallback: (payload: ExecutionTestPayload | EOTTestPayload) => void,
     deferredTillServerClose: Deferred<void>,
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -142,9 +142,9 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 testArgs = utils.addValueIfKeyNotExist(testArgs, '--capture', 'no');
             }
 
-            // add port with run test ids to env vars
-            const testIdsPipeName = await utils.startTestIdsNamedPipe(testIds);
-            mutableEnv.RUN_TEST_IDS_PIPE = testIdsPipeName;
+            // create a file with the test ids and set the environment variable to the file name
+            const testIdsFileName = await utils.writeTestIdsFile(testIds);
+            mutableEnv.RUN_TEST_IDS_PIPE = testIdsFileName;
             traceInfo(`All environment variables set for pytest execution: ${JSON.stringify(mutableEnv)}`);
 
             const spawnOptions: SpawnOptions = {
@@ -162,7 +162,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                     args: testArgs,
                     token: runInstance?.token,
                     testProvider: PYTEST_PROVIDER,
-                    runTestIdsPort: testIdsPipeName,
+                    runTestIdsPort: testIdsFileName,
                     pytestPort: resultNamedPipeName,
                 };
                 traceInfo(`Running DEBUG pytest with arguments: ${testArgs} for workspace ${uri.fsPath} \r\n`);
diff --git a/src/client/testing/testController/unittest/testExecutionAdapter.ts b/src/client/testing/testController/unittest/testExecutionAdapter.ts
@@ -137,8 +137,8 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         traceLog(`Running UNITTEST execution for the following test ids: ${testIds}`);
 
         // create named pipe server to send test ids
-        const testIdsPipeName = await utils.startTestIdsNamedPipe(testIds);
-        mutableEnv.RUN_TEST_IDS_PIPE = testIdsPipeName;
+        const testIdsFileName = await utils.writeTestIdsFile(testIds);
+        mutableEnv.RUN_TEST_IDS_PIPE = testIdsFileName;
         traceInfo(`All environment variables set for pytest execution: ${JSON.stringify(mutableEnv)}`);
 
         const spawnOptions: SpawnOptions = {
@@ -167,7 +167,7 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
                     args,
                     token: options.token,
                     testProvider: UNITTEST_PROVIDER,
-                    runTestIdsPort: testIdsPipeName,
+                    runTestIdsPort: testIdsFileName,
                     pytestPort: resultNamedPipeName, // change this from pytest
                 };
                 traceInfo(`Running DEBUG unittest for workspace ${options.cwd} with arguments: ${args}\r\n`);
diff --git a/src/test/testing/testController/pytest/pytestExecutionAdapter.unit.test.ts b/src/test/testing/testController/pytest/pytestExecutionAdapter.unit.test.ts
diff --git a/src/test/testing/testController/testCancellationRunAdapters.unit.test.ts b/src/test/testing/testController/testCancellationRunAdapters.unit.test.ts
diff --git a/src/test/testing/testController/unittest/testExecutionAdapter.unit.test.ts b/src/test/testing/testController/unittest/testExecutionAdapter.unit.test.ts
