From cc4b112c097af5271d1fcb8c2f90298e8c8af1d1 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Thu, 18 May 2023 10:44:54 -0700
Subject: [PATCH 01/12] allow large scale testing

---
 .../testing/testController/common/server.ts   | 33 +++++++++++++------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/client/testing/testController/common/server.ts b/src/client/testing/testController/common/server.ts
index 6849f0f8969a..a97eee70c5d8 100644
--- a/src/client/testing/testController/common/server.ts
+++ b/src/client/testing/testController/common/server.ts
@@ -26,23 +26,36 @@ 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;
+                        }
+                        // do I need to check that the UUIDs are the same or smth before adding them?
                         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) {

From 972d70856ddeb2c8cae32a3d228fef0e9274ff5f Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Thu, 18 May 2023 11:02:19 -0700
Subject: [PATCH 02/12] remove comment

---
 src/client/testing/testController/common/server.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/client/testing/testController/common/server.ts b/src/client/testing/testController/common/server.ts
index a97eee70c5d8..a00623aa33c7 100644
--- a/src/client/testing/testController/common/server.ts
+++ b/src/client/testing/testController/common/server.ts
@@ -45,7 +45,6 @@ export class PythonTestServer implements ITestServer, Disposable {
                             this._onDataReceived.fire({ uuid: '', data: '' });
                             return;
                         }
-                        // do I need to check that the UUIDs are the same or smth before adding them?
                         rawData = rpcHeaders.remainingRawData;
                         const rpcContent = jsonRPCContent(rpcHeaders.headers, rawData);
                         const extractedData = rpcContent.extractedJSON;

From e90decabb0d723602f55431903dc94f33b090b47 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Fri, 19 May 2023 11:57:50 -0700
Subject: [PATCH 03/12] pass test ids through stdin

---
 .../vscode_pytest/run_pytest_script.py        | 26 +++++++++++++++++++
 src/client/common/process/rawProcessApis.ts   |  4 +++
 src/client/common/process/types.ts            |  1 +
 .../pytest/pytestExecutionAdapter.ts          | 12 ++++++---
 .../testController/workspaceTestAdapter.ts    |  3 +--
 5 files changed, 40 insertions(+), 6 deletions(-)
 create mode 100644 pythonFiles/vscode_pytest/run_pytest_script.py

diff --git a/pythonFiles/vscode_pytest/run_pytest_script.py b/pythonFiles/vscode_pytest/run_pytest_script.py
new file mode 100644
index 000000000000..86756fa7b4b3
--- /dev/null
+++ b/pythonFiles/vscode_pytest/run_pytest_script.py
@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+import json
+import os
+import pathlib
+import sys
+
+import pytest
+
+# 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:]
+    try:
+        # Load test_ids from stdin.
+        test_ids = json.loads(sys.stdin.read())
+        arg_array = ["-p", "vscode_pytest"] + args + test_ids
+        pytest.main(arg_array)
+    except json.JSONDecodeError:
+        print("Error: Could not parse test ids from stdin")
diff --git a/src/client/common/process/rawProcessApis.ts b/src/client/common/process/rawProcessApis.ts
index 025e5b607229..f12058ab6d35 100644
--- a/src/client/common/process/rawProcessApis.ts
+++ b/src/client/common/process/rawProcessApis.ts
@@ -93,6 +93,10 @@ export function plainExec(
     const spawnOptions = getDefaultOptions(options, defaultEnv);
     const encoding = spawnOptions.encoding ? spawnOptions.encoding : 'utf8';
     const proc = spawn(file, args, spawnOptions);
+    if (options.pytestExecutionTestIds) {
+        proc.stdin?.write(JSON.stringify(options.pytestExecutionTestIds));
+        proc.stdin?.end();
+    }
     // Listen to these errors (unhandled errors in streams tears down the process).
     // Errors will be bubbled up to the `error` event in `proc`, hence no need to log.
     proc.stdout?.on('error', noop);
diff --git a/src/client/common/process/types.ts b/src/client/common/process/types.ts
index 8298957285e8..ae787bf37d56 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;
+    pytestExecutionTestIds?: string[];
 };
 
 export type ShellOptions = ExecOptions & { throwOnStdErr?: boolean };
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
index 623fd1ff3a8c..ee98f820069c 100644
--- a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -90,6 +90,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 TEST_PORT: this.testServer.getPort().toString(),
             },
             outputChannel: this.outputChannel,
+            pytestExecutionTestIds: testIds,
         };
 
         // Create the Python environment in which to execute the command.
@@ -114,7 +115,7 @@ 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);
             if (debugBool) {
                 const pytestPort = this.testServer.getPort().toString();
                 const pytestUUID = uuid.toString();
@@ -129,9 +130,12 @@ 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);
+                const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py');
+                const runArgs = [scriptPath, ...testArgs];
+                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..9c4be2900c1e 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'));
 
             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) {

From 807fa21ca06346cae3a130472a68aeff1b4a325e Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Fri, 19 May 2023 16:09:37 -0700
Subject: [PATCH 04/12] working but need to clean up

---
 .../vscode_pytest/run_pytest_script.py        |  44 +++++++-
 .../testing/testController/common/utils.ts    |  41 +++++++
 .../pytest/pytestExecutionAdapter.ts          | 101 +++++++++++++++++-
 3 files changed, 183 insertions(+), 3 deletions(-)

diff --git a/pythonFiles/vscode_pytest/run_pytest_script.py b/pythonFiles/vscode_pytest/run_pytest_script.py
index 86756fa7b4b3..8be5eb960aae 100644
--- a/pythonFiles/vscode_pytest/run_pytest_script.py
+++ b/pythonFiles/vscode_pytest/run_pytest_script.py
@@ -1,12 +1,27 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
+# Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 import json
 import os
 import pathlib
+import socket
 import sys
 
 import pytest
 
+# Path to the directory you want to append
+directory_path = pathlib.Path(__file__).parent.parent / "lib" / "python"
+import os
+
+sys.path.append(os.fspath(directory_path))
+print(sys.path)
+
+# Disable Flake8 rule for the next import statement
+# flake8: noqa: E402
+import debugpy
+
+debugpy.connect(5678)
+debugpy.breakpoint()
+
 # 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.
@@ -17,6 +32,31 @@
     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 = 0
+    if run_test_ids_port is not None:
+        run_test_ids_port_int = int(run_test_ids_port)
+    try:
+        client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+        # Bind the socket to a specific address and port
+        client_socket.connect(("localhost", run_test_ids_port_int))
+
+        # Listen for incoming connections
+        # client_socket.listen(1)
+        print(f"Server listening on port {run_test_ids_port_int}...")
+
+        while True:
+            # Receive the data from the client
+            # data = client_socket.recv(1024).decode("utf-8")
+            data: bytes = client_socket.recv(1024 * 1024)
+            print(f"Received data: {data}")
+
+            # Close the client connection
+            client_socket.close()
+    except socket.error as e:
+        print(f"Error: Could not connect to runTestIdsPort: {e}")
+        print("Error: Could not connect to runTestIdsPort")
     try:
         # Load test_ids from stdin.
         test_ids = json.loads(sys.stdin.read())
@@ -24,3 +64,5 @@
         pytest.main(arg_array)
     except json.JSONDecodeError:
         print("Error: Could not parse test ids from stdin")
+    finally:
+        print("PSTHAFS", os.environ.get("TEST_PORT"))
diff --git a/src/client/testing/testController/common/utils.ts b/src/client/testing/testController/common/utils.ts
index e0bad383d695..31676fc228f8 100644
--- a/src/client/testing/testController/common/utils.ts
+++ b/src/client/testing/testController/common/utils.ts
@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+import * as net from 'net';
 
 export function fixLogLines(content: string): string {
     const lines = content.split(/\r?\n/g);
@@ -50,3 +51,43 @@ export function jsonRPCContent(headers: Map<string, string>, rawData: string): I
         remainingRawData,
     };
 }
+export const startServer = (test_ids: string): Promise<number> =>
+    new Promise((resolve, reject) => {
+        const server = net.createServer((socket: net.Socket) => {
+            console.log('Client connected');
+
+            // Convert the test_ids array to JSON
+            const testData = JSON.stringify(test_ids);
+
+            // 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);
+
+            // Store the port of the socket as test_run_socket
+            const test_run_socket = socket.localPort;
+
+            // Handle socket events
+            socket.on('data', (data) => {
+                console.log('Received data:', data.toString());
+            });
+
+            socket.on('end', () => {
+                console.log('Client disconnected');
+            });
+        });
+
+        server.listen(0, () => {
+            const { port } = server.address() as net.AddressInfo;
+            console.log(`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 ee98f820069c..2a6150f1a287 100644
--- a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -3,6 +3,7 @@
 
 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';
@@ -116,6 +117,104 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 testArgs.push('--capture', 'no');
             }
             const pluginArgs = ['-p', 'vscode_pytest'].concat(testArgs).concat(testIds);
+            const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py');
+            const runArgs = [scriptPath, ...testArgs];
+            // let testRunSocket = '5';
+            // // let s = net.Server;
+            // const serverS = net.createServer((socket: net.Socket) => {
+            //     console.log('Client connected');
+
+            //     // 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);
+
+            //     // Store the port of the socket as test_run_socket
+            //     testRunSocket = socket.localPort ? socket.localPort.toString() : '0';
+            //     if (spawnOptions.extraVariables) spawnOptions.extraVariables.RUN_TEST_IDS_PORT = testRunSocket;
+
+            //     socket.on('end', () => {
+            //         console.log('Client disconnected');
+            //     });
+            // });
+            // const portP = (serverS.address() as net.AddressInfo).port;
+            // await serverS.listen(undefined, 'localhost', () => {
+            //     resolve();
+            // });
+            // console.log(`Test ID Server: listening on port ${portP}`);
+            // serverS.listen(0, () => {
+            //     const { port: assignedPort } = serverS.address() as net.AddressInfo;
+            //     console.log(`Server listening on port ${assignedPort}`);
+            // });
+            // Start the server and wait until it is listening
+            // 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}`;
+
+            const startServer = (): Promise<number> =>
+                new Promise((resolve, reject) => {
+                    const server = net.createServer((socket: net.Socket) => {
+                        console.log('Client connected');
+
+                        // Send the payload to the socket
+                        socket.write(payload);
+                        console.log('payload sent', payload);
+
+                        // Store the port of the socket as test_run_socket
+                        // const test_run_socket = socket.localPort;
+
+                        // Handle socket events
+                        socket.on('data', (data) => {
+                            console.log('Received data:', data.toString());
+                        });
+
+                        socket.on('end', () => {
+                            console.log('Client disconnected');
+                        });
+                    });
+
+                    server.listen(0, () => {
+                        const { port } = server.address() as net.AddressInfo;
+                        console.log(`Server listening on port ${port}`);
+                        resolve(port);
+                    });
+
+                    server.on('error', (error: Error) => {
+                        reject(error);
+                    });
+                    server.on('connection', (socket: net.Socket) => {
+                        socket.write();
+                        console.log('Test server connected to a client.');
+                    });
+                });
+
+            // Start the server and wait until it is listening
+            await startServer()
+                .then((assignedPort) => {
+                    console.log(`Server started and listening on port ${assignedPort}`);
+                    if (spawnOptions.extraVariables)
+                        spawnOptions.extraVariables.RUN_TEST_IDS_PORT = assignedPort.toString();
+                    // Continue with your code here
+                })
+                .catch((error) => {
+                    console.error('Error starting server:', error);
+                });
+
             if (debugBool) {
                 const pytestPort = this.testServer.getPort().toString();
                 const pytestUUID = uuid.toString();
@@ -130,8 +229,6 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 console.debug(`Running debug test with arguments: ${pluginArgs.join(' ')}\r\n`);
                 await debugLauncher!.launchDebugger(launchOptions);
             } else {
-                const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py');
-                const runArgs = [scriptPath, ...testArgs];
                 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);

From cf0cde1928da1ce467991183f18cef0f6e71a3fc Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Fri, 19 May 2023 16:36:26 -0700
Subject: [PATCH 05/12] working with process json

---
 .../vscode_pytest/run_pytest_script.py        | 80 ++++++++++++++-----
 .../pytest/pytestExecutionAdapter.ts          | 61 +-------------
 2 files changed, 63 insertions(+), 78 deletions(-)

diff --git a/pythonFiles/vscode_pytest/run_pytest_script.py b/pythonFiles/vscode_pytest/run_pytest_script.py
index 8be5eb960aae..0c198f28aade 100644
--- a/pythonFiles/vscode_pytest/run_pytest_script.py
+++ b/pythonFiles/vscode_pytest/run_pytest_script.py
@@ -1,10 +1,17 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+import contextlib
+import io
 import json
 import os
 import pathlib
+import random
 import socket
+import subprocess
 import sys
+import threading
+import uuid
+from typing import Any, Dict, List, Optional, Union
 
 import pytest
 
@@ -21,6 +28,32 @@
 
 debugpy.connect(5678)
 debugpy.breakpoint()
+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
@@ -33,36 +66,43 @@
     # 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 = 0
-    if run_test_ids_port is not None:
-        run_test_ids_port_int = int(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)
-
-        # Bind the socket to a specific address and port
         client_socket.connect(("localhost", run_test_ids_port_int))
-
-        # Listen for incoming connections
-        # client_socket.listen(1)
-        print(f"Server listening on port {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).decode("utf-8")
-            data: bytes = client_socket.recv(1024 * 1024)
-            print(f"Received data: {data}")
+            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""
 
-            # Close the client connection
-            client_socket.close()
+                # 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:
-        # Load test_ids from stdin.
-        test_ids = json.loads(sys.stdin.read())
-        arg_array = ["-p", "vscode_pytest"] + args + test_ids
-        pytest.main(arg_array)
+        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")
-    finally:
-        print("PSTHAFS", os.environ.get("TEST_PORT"))
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
index 2a6150f1a287..b6c46b51f6ba 100644
--- a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -119,51 +119,9 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
             const pluginArgs = ['-p', 'vscode_pytest'].concat(testArgs).concat(testIds);
             const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py');
             const runArgs = [scriptPath, ...testArgs];
-            // let testRunSocket = '5';
-            // // let s = net.Server;
-            // const serverS = net.createServer((socket: net.Socket) => {
-            //     console.log('Client connected');
-
-            //     // 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);
-
-            //     // Store the port of the socket as test_run_socket
-            //     testRunSocket = socket.localPort ? socket.localPort.toString() : '0';
-            //     if (spawnOptions.extraVariables) spawnOptions.extraVariables.RUN_TEST_IDS_PORT = testRunSocket;
-
-            //     socket.on('end', () => {
-            //         console.log('Client disconnected');
-            //     });
-            // });
-            // const portP = (serverS.address() as net.AddressInfo).port;
-            // await serverS.listen(undefined, 'localhost', () => {
-            //     resolve();
-            // });
-            // console.log(`Test ID Server: listening on port ${portP}`);
-            // serverS.listen(0, () => {
-            //     const { port: assignedPort } = serverS.address() as net.AddressInfo;
-            //     console.log(`Server listening on port ${assignedPort}`);
-            // });
-            // Start the server and wait until it is listening
-            // Convert the test_ids array to JSON
-            const testData = JSON.stringify(testIds);
 
-            // Create the headers
+            const testData = JSON.stringify(testIds);
             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}`;
 
             const startServer = (): Promise<number> =>
@@ -171,18 +129,6 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                     const server = net.createServer((socket: net.Socket) => {
                         console.log('Client connected');
 
-                        // Send the payload to the socket
-                        socket.write(payload);
-                        console.log('payload sent', payload);
-
-                        // Store the port of the socket as test_run_socket
-                        // const test_run_socket = socket.localPort;
-
-                        // Handle socket events
-                        socket.on('data', (data) => {
-                            console.log('Received data:', data.toString());
-                        });
-
                         socket.on('end', () => {
                             console.log('Client disconnected');
                         });
@@ -198,8 +144,8 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                         reject(error);
                     });
                     server.on('connection', (socket: net.Socket) => {
-                        socket.write();
-                        console.log('Test server connected to a client.');
+                        socket.write(payload);
+                        console.log('payload sent', payload);
                     });
                 });
 
@@ -209,7 +155,6 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                     console.log(`Server started and listening on port ${assignedPort}`);
                     if (spawnOptions.extraVariables)
                         spawnOptions.extraVariables.RUN_TEST_IDS_PORT = assignedPort.toString();
-                    // Continue with your code here
                 })
                 .catch((error) => {
                     console.error('Error starting server:', error);

From 4eba88189d51af28ac5ef202dfe1e574d5ec8c93 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Mon, 22 May 2023 10:50:52 -0700
Subject: [PATCH 06/12] clean up

---
 src/client/common/process/rawProcessApis.ts               | 4 ----
 src/client/testing/testController/common/utils.ts         | 3 ---
 src/client/testing/testController/workspaceTestAdapter.ts | 2 +-
 3 files changed, 1 insertion(+), 8 deletions(-)

diff --git a/src/client/common/process/rawProcessApis.ts b/src/client/common/process/rawProcessApis.ts
index f12058ab6d35..025e5b607229 100644
--- a/src/client/common/process/rawProcessApis.ts
+++ b/src/client/common/process/rawProcessApis.ts
@@ -93,10 +93,6 @@ export function plainExec(
     const spawnOptions = getDefaultOptions(options, defaultEnv);
     const encoding = spawnOptions.encoding ? spawnOptions.encoding : 'utf8';
     const proc = spawn(file, args, spawnOptions);
-    if (options.pytestExecutionTestIds) {
-        proc.stdin?.write(JSON.stringify(options.pytestExecutionTestIds));
-        proc.stdin?.end();
-    }
     // Listen to these errors (unhandled errors in streams tears down the process).
     // Errors will be bubbled up to the `error` event in `proc`, hence no need to log.
     proc.stdout?.on('error', noop);
diff --git a/src/client/testing/testController/common/utils.ts b/src/client/testing/testController/common/utils.ts
index 31676fc228f8..c5c189f5930f 100644
--- a/src/client/testing/testController/common/utils.ts
+++ b/src/client/testing/testController/common/utils.ts
@@ -68,9 +68,6 @@ export const startServer = (test_ids: string): Promise<number> =>
             // Send the payload to the socket
             socket.write(payload);
 
-            // Store the port of the socket as test_run_socket
-            const test_run_socket = socket.localPort;
-
             // Handle socket events
             socket.on('data', (data) => {
                 console.log('Received data:', data.toString());
diff --git a/src/client/testing/testController/workspaceTestAdapter.ts b/src/client/testing/testController/workspaceTestAdapter.ts
index 9c4be2900c1e..0b19d4d87d6f 100644
--- a/src/client/testing/testController/workspaceTestAdapter.ts
+++ b/src/client/testing/testController/workspaceTestAdapter.ts
@@ -348,7 +348,7 @@ 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'));
             let errorNode = testController.items.get(`DiscoveryError:${workspacePath}`);
             const message = util.format(
                 `${testingErrorConst} ${Testing.seePythonOutput}\r\n`,

From 9730dcb53ff54655dcb6a08412701354e9c368bd Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Tue, 23 May 2023 14:45:55 -0700
Subject: [PATCH 07/12] type edits

---
 .github/workflows/build.yml                       | 1 +
 src/client/testing/testController/common/utils.ts | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index f418a802ff8e..2511eb02aecc 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -101,6 +101,7 @@ jobs:
           python -m pip install --upgrade -r build/test-requirements.txt
 
       - name: Run Pyright
+        version: 1.1.308
         uses: jakebailey/pyright-action@v1
         with:
           version: 1.1.308
diff --git a/src/client/testing/testController/common/utils.ts b/src/client/testing/testController/common/utils.ts
index c5c189f5930f..e056166ad608 100644
--- a/src/client/testing/testController/common/utils.ts
+++ b/src/client/testing/testController/common/utils.ts
@@ -51,13 +51,13 @@ export function jsonRPCContent(headers: Map<string, string>, rawData: string): I
         remainingRawData,
     };
 }
-export const startServer = (test_ids: string): Promise<number> =>
+export const startServer = (testIds: string): Promise<number> =>
     new Promise((resolve, reject) => {
         const server = net.createServer((socket: net.Socket) => {
             console.log('Client connected');
 
             // Convert the test_ids array to JSON
-            const testData = JSON.stringify(test_ids);
+            const testData = JSON.stringify(testIds);
 
             // Create the headers
             const headers = [`Content-Length: ${Buffer.byteLength(testData)}`, 'Content-Type: application/json'];

From 9db9cd27d40f2eca427a57bfd083db3675ae3d27 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Tue, 23 May 2023 14:54:48 -0700
Subject: [PATCH 08/12] fix logging

---
 .vscode/launch.json                           |  3 +--
 .../vscode_pytest/run_pytest_script.py        | 19 -------------------
 .../testing/testController/common/utils.ts    |  9 ++++-----
 .../pytest/pytestExecutionAdapter.ts          | 12 +++++-------
 4 files changed, 10 insertions(+), 33 deletions(-)

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/vscode_pytest/run_pytest_script.py b/pythonFiles/vscode_pytest/run_pytest_script.py
index 0c198f28aade..fd458d2eb470 100644
--- a/pythonFiles/vscode_pytest/run_pytest_script.py
+++ b/pythonFiles/vscode_pytest/run_pytest_script.py
@@ -1,33 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-import contextlib
 import io
 import json
 import os
 import pathlib
-import random
 import socket
-import subprocess
 import sys
-import threading
-import uuid
-from typing import Any, Dict, List, Optional, Union
 
 import pytest
 
-# Path to the directory you want to append
-directory_path = pathlib.Path(__file__).parent.parent / "lib" / "python"
-import os
-
-sys.path.append(os.fspath(directory_path))
-print(sys.path)
-
-# Disable Flake8 rule for the next import statement
-# flake8: noqa: E402
-import debugpy
-
-debugpy.connect(5678)
-debugpy.breakpoint()
 CONTENT_LENGTH: str = "Content-Length:"
 
 
diff --git a/src/client/testing/testController/common/utils.ts b/src/client/testing/testController/common/utils.ts
index e056166ad608..88e3450d35dc 100644
--- a/src/client/testing/testController/common/utils.ts
+++ b/src/client/testing/testController/common/utils.ts
@@ -1,6 +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);
@@ -54,8 +55,6 @@ export function jsonRPCContent(headers: Map<string, string>, rawData: string): I
 export const startServer = (testIds: string): Promise<number> =>
     new Promise((resolve, reject) => {
         const server = net.createServer((socket: net.Socket) => {
-            console.log('Client connected');
-
             // Convert the test_ids array to JSON
             const testData = JSON.stringify(testIds);
 
@@ -70,17 +69,17 @@ export const startServer = (testIds: string): Promise<number> =>
 
             // Handle socket events
             socket.on('data', (data) => {
-                console.log('Received data:', data.toString());
+                traceLog('Received data:', data.toString());
             });
 
             socket.on('end', () => {
-                console.log('Client disconnected');
+                traceLog('Client disconnected');
             });
         });
 
         server.listen(0, () => {
             const { port } = server.address() as net.AddressInfo;
-            console.log(`Server listening on port ${port}`);
+            traceLog(`Server listening on port ${port}`);
             resolve(port);
         });
 
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
index b6c46b51f6ba..f90edf3c91dc 100644
--- a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -6,7 +6,7 @@ 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,
@@ -127,16 +127,14 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
             const startServer = (): Promise<number> =>
                 new Promise((resolve, reject) => {
                     const server = net.createServer((socket: net.Socket) => {
-                        console.log('Client connected');
-
                         socket.on('end', () => {
-                            console.log('Client disconnected');
+                            traceLog('Client disconnected');
                         });
                     });
 
                     server.listen(0, () => {
                         const { port } = server.address() as net.AddressInfo;
-                        console.log(`Server listening on port ${port}`);
+                        traceLog(`Server listening on port ${port}`);
                         resolve(port);
                     });
 
@@ -145,14 +143,14 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                     });
                     server.on('connection', (socket: net.Socket) => {
                         socket.write(payload);
-                        console.log('payload sent', payload);
+                        traceLog('payload sent', payload);
                     });
                 });
 
             // Start the server and wait until it is listening
             await startServer()
                 .then((assignedPort) => {
-                    console.log(`Server started and listening on port ${assignedPort}`);
+                    traceLog(`Server started and listening on port ${assignedPort}`);
                     if (spawnOptions.extraVariables)
                         spawnOptions.extraVariables.RUN_TEST_IDS_PORT = assignedPort.toString();
                 })

From 83256ebe47103733d32f6c86d50d4985c50c7ac8 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Tue, 23 May 2023 15:02:45 -0700
Subject: [PATCH 09/12] fix pyright for pr check

---
 .github/workflows/pr-check.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/pr-check.yml b/.github/workflows/pr-check.yml
index 84903463e204..054eb2b989c3 100644
--- a/.github/workflows/pr-check.yml
+++ b/.github/workflows/pr-check.yml
@@ -75,6 +75,7 @@ jobs:
           python -m pip install --upgrade -r build/test-requirements.txt
 
       - name: Run Pyright
+        version: 1.1.308
         uses: jakebailey/pyright-action@v1
         with:
           version: 1.1.308

From 2090f2372e5c4fc22c77f88a0eb06ddee7f7c140 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Tue, 23 May 2023 15:39:24 -0700
Subject: [PATCH 10/12] fix typing

---
 pythonFiles/unittestadapter/discovery.py       | 8 +++++++-
 pythonFiles/unittestadapter/utils.py           | 9 ++++++++-
 pythonFiles/vscode_pytest/run_pytest_script.py | 3 ++-
 3 files changed, 17 insertions(+), 3 deletions(-)

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
index fd458d2eb470..f6d6bdcafd5f 100644
--- a/pythonFiles/vscode_pytest/run_pytest_script.py
+++ b/pythonFiles/vscode_pytest/run_pytest_script.py
@@ -6,13 +6,14 @@
 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]:
+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)
 

From 81c59716e339361f5d53e7439bfea33648adae27 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Wed, 24 May 2023 09:51:49 -0700
Subject: [PATCH 11/12] switch to stdinStr

---
 src/client/common/process/types.ts                              | 2 +-
 .../testing/testController/pytest/pytestExecutionAdapter.ts     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/client/common/process/types.ts b/src/client/common/process/types.ts
index ae787bf37d56..62e787b694b5 100644
--- a/src/client/common/process/types.ts
+++ b/src/client/common/process/types.ts
@@ -25,7 +25,7 @@ export type SpawnOptions = ChildProcessSpawnOptions & {
     throwOnStdErr?: boolean;
     extraVariables?: NodeJS.ProcessEnv;
     outputChannel?: OutputChannel;
-    pytestExecutionTestIds?: string[];
+    stdinStr?: string;
 };
 
 export type ShellOptions = ExecOptions & { throwOnStdErr?: boolean };
diff --git a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
index f90edf3c91dc..4c6dcd9cdbee 100644
--- a/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestExecutionAdapter.ts
@@ -91,7 +91,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 TEST_PORT: this.testServer.getPort().toString(),
             },
             outputChannel: this.outputChannel,
-            pytestExecutionTestIds: testIds,
+            stdinStr: testIds.toString(),
         };
 
         // Create the Python environment in which to execute the command.

From 88b0389b6ca3ea62031378f73e4626577147f817 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <eleanorboyd@microsoft.com>
Date: Wed, 24 May 2023 13:06:51 -0700
Subject: [PATCH 12/12] remove extra version pyright note

---
 .github/workflows/build.yml    | 1 -
 .github/workflows/pr-check.yml | 1 -
 2 files changed, 2 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 2511eb02aecc..f418a802ff8e 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -101,7 +101,6 @@ jobs:
           python -m pip install --upgrade -r build/test-requirements.txt
 
       - name: Run Pyright
-        version: 1.1.308
         uses: jakebailey/pyright-action@v1
         with:
           version: 1.1.308
diff --git a/.github/workflows/pr-check.yml b/.github/workflows/pr-check.yml
index 054eb2b989c3..84903463e204 100644
--- a/.github/workflows/pr-check.yml
+++ b/.github/workflows/pr-check.yml
@@ -75,7 +75,6 @@ jobs:
           python -m pip install --upgrade -r build/test-requirements.txt
 
       - name: Run Pyright
-        version: 1.1.308
         uses: jakebailey/pyright-action@v1
         with:
           version: 1.1.308