switch to using tcp for comm with server

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

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import * as http from 'http';
 import * as net from 'net';
 import * as crypto from 'crypto';
 import { Disposable, Event, EventEmitter } from 'vscode';
@@ -14,53 +13,63 @@ import { traceLog } from '../../../logging';
 import { DataReceivedEvent, ITestServer, TestCommandOptions } from './types';
 import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
 import { UNITTEST_PROVIDER } from '../../common/constants';
+import { jsonRPCProcessor } from './utils';
 
 export class PythonTestServer implements ITestServer, Disposable {
     private _onDataReceived: EventEmitter<DataReceivedEvent> = new EventEmitter<DataReceivedEvent>();
 
     private uuids: Map<string, string>;
 
-    private server: http.Server;
+    private server: net.Server;
 
     private ready: Promise<void>;
 
     constructor(private executionFactory: IPythonExecutionFactory, private debugLauncher: ITestDebugLauncher) {
         this.uuids = new Map();
 
-        const requestListener: http.RequestListener = async (request, response) => {
-            const buffers = [];
-
-            try {
-                for await (const chunk of request) {
-                    buffers.push(chunk);
-                }
-
-                const data = Buffer.concat(buffers).toString();
-                // grab the uuid from the header
-                const indexRequestuuid = request.rawHeaders.indexOf('Request-uuid');
-                const uuid = request.rawHeaders[indexRequestuuid + 1];
-                response.end();
-
-                JSON.parse(data);
-                // Check if the uuid we received exists in the list of active ones.
-                // If yes, process the response, if not, ignore it.
-                const cwd = this.uuids.get(uuid);
-                if (cwd) {
-                    this._onDataReceived.fire({ cwd, data });
-                    this.uuids.delete(uuid);
+        this.server = net.createServer((socket: net.Socket) => {
+            socket.on('data', (data: Buffer) => {
+                try {
+                    let rawData: string = data.toString();
+
+                    while (rawData.length > 0) {
+                        const { headers, extractedData, remainingRawData } = jsonRPCProcessor(rawData);
+                        rawData = remainingRawData;
+                        const uuid = headers.get('Request-uuid');
+                        if (uuid) {
+                            const cwd = this.uuids.get(uuid);
+                            if (cwd) {
+                                this._onDataReceived.fire({ uuid, data: extractedData });
+                                this.uuids.delete(uuid);
+                            }
+                        } else {
+                            traceLog(`Error processing test server request: uuid not found`);
+                            this._onDataReceived.fire({ uuid: '', data: '' });
+                        }
+                    }
+                } catch (ex) {
+                    traceLog(`Error processing test server request: ${ex} observe`);
+                    this._onDataReceived.fire({ uuid: '', data: '' });
                 }
-            } catch (ex) {
-                traceLog(`Error processing test server request: ${ex} observe`);
-                this._onDataReceived.fire({ cwd: '', data: '' });
-            }
-        };
-
-        this.server = http.createServer(requestListener);
+            });
+        });
         this.ready = new Promise((resolve, _reject) => {
             this.server.listen(undefined, 'localhost', () => {
                 resolve();
             });
         });
+        this.server.on('error', (ex) => {
+            traceLog(`Error starting test server: ${ex}`);
+        });
+        this.server.on('close', () => {
+            traceLog('Test server closed');
+        });
+        this.server.on('listening', () => {
+            traceLog('Test server listening');
+        });
+        this.server.on('connection', () => {
+            traceLog('Test server connection');
+        });
     }
 
     public serverReady(): Promise<void> {
@@ -142,7 +151,7 @@ export class PythonTestServer implements ITestServer, Disposable {
         } catch (ex) {
             this.uuids.delete(uuid);
             this._onDataReceived.fire({
-                cwd: options.cwd,
+                uuid,
                 data: JSON.stringify({
                     status: 'error',
                     errors: [(ex as Error).message],

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

@@ -128,7 +128,7 @@ export type RawDiscoveredTests = {
 // New test discovery adapter types
 
 export type DataReceivedEvent = {
-    cwd: string;
+    uuid: string;
     data: string;
 };
 

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

@@ -5,3 +5,34 @@ export function fixLogLines(content: string): string {
     const lines = content.split(/\r?\n/g);
     return `${lines.join('\r\n')}\r\n`;
 }
+export interface IJSONRPCMessage {
+    headers: Map<string, string>;
+    extractedData: string;
+    remainingRawData: string;
+}
+
+export function jsonRPCProcessor(rawData: string): IJSONRPCMessage {
+    const lines = rawData.split('\n');
+    let remainingRawData = '';
+    const headerMap = new Map<string, string>();
+    for (let i = 0; i < lines.length; i += 1) {
+        const line = lines[i];
+        if (line === '') {
+            remainingRawData = lines.slice(i + 1).join('\n');
+            break;
+        }
+        const [key, value] = line.split(':');
+        if (['Content-Length', 'Content-Type', 'Request-uuid'].includes(key)) {
+            headerMap.set(key.trim(), value.trim());
+        }
+    }
+
+    const length = parseInt(headerMap.get('Content-Length') ?? '0', 10);
+    const data = remainingRawData.slice(0, length);
+    remainingRawData = remainingRawData.slice(length);
+    return {
+        headers: headerMap,
+        extractedData: data,
+        remainingRawData,
+    };
+}

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

@@ -17,20 +17,19 @@ import { DataReceivedEvent, DiscoveredTestPayload, ITestDiscoveryAdapter, ITestS
  * Wrapper class for unittest test discovery. This is where we call `runTestCommand`. #this seems incorrectly copied
  */
 export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
-    private deferred: Deferred<DiscoveredTestPayload> | undefined;
+    private promiseMap: Map<string, Deferred<DiscoveredTestPayload | undefined>> = new Map();
 
-    private cwd: string | undefined;
+    private deferred: Deferred<DiscoveredTestPayload> | undefined;
 
     constructor(public testServer: ITestServer, public configSettings: IConfigurationService) {
         testServer.onDataReceived(this.onDataReceivedHandler, this);
     }
 
-    public onDataReceivedHandler({ cwd, data }: DataReceivedEvent): void {
-        if (this.deferred && cwd === this.cwd) {
-            const testData: DiscoveredTestPayload = JSON.parse(data);
-
-            this.deferred.resolve(testData);
-            this.deferred = undefined;
+    public onDataReceivedHandler({ uuid, data }: DataReceivedEvent): void {
+        const deferred = this.promiseMap.get(uuid);
+        if (deferred) {
+            deferred.resolve(JSON.parse(data));
+            this.promiseMap.delete(uuid);
         }
     }
 
@@ -51,43 +50,43 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
     // }
 
     async runPytestDiscovery(uri: Uri, executionFactory: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
-        if (!this.deferred) {
-            this.deferred = createDeferred<DiscoveredTestPayload>();
-            const relativePathToPytest = 'pythonFiles';
-            const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
-            const uuid = this.testServer.createUUID(uri.fsPath);
-            const settings = this.configSettings.getSettings(uri);
-            const { pytestArgs } = settings.testing;
+        const deferred = createDeferred<DiscoveredTestPayload>();
+        this.deferred = createDeferred<DiscoveredTestPayload>();
+        const relativePathToPytest = 'pythonFiles';
+        const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
+        const uuid = this.testServer.createUUID(uri.fsPath);
+        const settings = this.configSettings.getSettings(uri);
+        const { pytestArgs } = settings.testing;
 
-            const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
-            const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
+        const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
+        const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
 
-            const spawnOptions: SpawnOptions = {
-                cwd: uri.fsPath,
-                throwOnStdErr: true,
-                extraVariables: {
-                    PYTHONPATH: pythonPathCommand,
-                    TEST_UUID: uuid.toString(),
-                    TEST_PORT: this.testServer.getPort().toString(),
-                },
-            };
+        const spawnOptions: SpawnOptions = {
+            cwd: uri.fsPath,
+            throwOnStdErr: true,
+            extraVariables: {
+                PYTHONPATH: pythonPathCommand,
+                TEST_UUID: uuid.toString(),
+                TEST_PORT: this.testServer.getPort().toString(),
+            },
+        };
 
-            // Create the Python environment in which to execute the command.
-            const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
-                allowEnvironmentFetchExceptions: false,
-                resource: uri,
-            };
-            const execService = await executionFactory.createActivatedEnvironment(creationOptions);
+        // Create the Python environment in which to execute the command.
+        const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
+            allowEnvironmentFetchExceptions: false,
+            resource: uri,
+        };
+        const execService = await executionFactory.createActivatedEnvironment(creationOptions);
 
-            try {
-                execService.exec(
-                    ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs),
-                    spawnOptions,
-                );
-            } catch (ex) {
-                console.error(ex);
-            }
+        try {
+            execService.exec(
+                ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs),
+                spawnOptions,
+            );
+        } catch (ex) {
+            console.error(ex);
         }
-        return this.deferred.promise;
+
+        return deferred.promise;
     }
 }

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

@@ -12,20 +12,19 @@ import { DataReceivedEvent, ExecutionTestPayload, ITestExecutionAdapter, ITestSe
  */
 
 export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
-    private deferred: Deferred<ExecutionTestPayload> | undefined;
+    private promiseMap: Map<string, Deferred<ExecutionTestPayload | undefined>> = new Map();
 
-    private cwd: string | undefined;
+    private deferred: Deferred<ExecutionTestPayload> | undefined;
 
     constructor(public testServer: ITestServer, public configSettings: IConfigurationService) {
         testServer.onDataReceived(this.onDataReceivedHandler, this);
     }
 
-    public onDataReceivedHandler({ cwd, data }: DataReceivedEvent): void {
-        if (this.deferred && cwd === this.cwd) {
-            const testData: ExecutionTestPayload = JSON.parse(data);
-
-            this.deferred.resolve(testData);
-            this.deferred = undefined;
+    public onDataReceivedHandler({ uuid, data }: DataReceivedEvent): void {
+        const deferred = this.promiseMap.get(uuid);
+        if (deferred) {
+            deferred.resolve(JSON.parse(data));
+            this.promiseMap.delete(uuid);
         }
     }
 
@@ -42,48 +41,45 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
     //         debugBool?: boolean,
     //         executionFactory?: IPythonExecutionFactory,
     //     ): Promise<ExecutionTestPayload> {
-    //         if (!this.deferred) {
-    //             this.deferred = createDeferred<ExecutionTestPayload>();
-    //             const relativePathToPytest = 'pythonFiles';
-    //             const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
-    //             this.configSettings.isTestExecution();
-    //             const uuid = this.testServer.createUUID(uri.fsPath);
-    //             const settings = this.configSettings.getSettings(uri);
-    //             const { pytestArgs } = settings.testing;
+    //         const deferred = createDeferred<ExecutionTestPayload>();
+    //         const relativePathToPytest = 'pythonFiles';
+    //         const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
+    //         this.configSettings.isTestExecution();
+    //         const uuid = this.testServer.createUUID(uri.fsPath);
+    //         this.promiseMap.set(uuid, deferred);
+    //         const settings = this.configSettings.getSettings(uri);
+    //         const { pytestArgs } = settings.testing;
 
-    //             const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
-    //             const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
+    //         const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
+    //         const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
 
-    //             const spawnOptions: SpawnOptions = {
-    //                 cwd: uri.fsPath,
-    //                 throwOnStdErr: true,
-    //                 extraVariables: {
-    //                     PYTHONPATH: pythonPathCommand,
-    //                     TEST_UUID: uuid.toString(),
-    //                     TEST_PORT: this.testServer.getPort().toString(),
-    //                 },
-    //             };
+    //         const spawnOptions: SpawnOptions = {
+    //             cwd: uri.fsPath,
+    //             throwOnStdErr: true,
+    //             extraVariables: {
+    //                 PYTHONPATH: pythonPathCommand,
+    //                 TEST_UUID: uuid.toString(),
+    //                 TEST_PORT: this.testServer.getPort().toString(),
+    //             },
+    //         };
 
-    //             // Create the Python environment in which to execute the command.
-    //             const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
-    //                 allowEnvironmentFetchExceptions: false,
-    //                 resource: uri,
-    //             };
-    //             // need to check what will happen in the exec service is NOT defined and is null
-    //             const execService = await executionFactory?.createActivatedEnvironment(creationOptions);
+    //         // Create the Python environment in which to execute the command.
+    //         const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
+    //             allowEnvironmentFetchExceptions: false,
+    //             resource: uri,
+    //         };
+    //         // need to check what will happen in the exec service is NOT defined and is null
+    //         const execService = await executionFactory?.createActivatedEnvironment(creationOptions);
 
-    //             const testIdsString = testIds.join(' ');
-    //             console.debug('what to do with debug bool?', debugBool);
-    //             try {
-    //                 execService?.exec(
-    //                     ['-m', 'pytest', '-p', 'vscode_pytest', testIdsString].concat(pytestArgs),
-    //                     spawnOptions,
-    //                 );
-    //             } catch (ex) {
-    //                 console.error(ex);
-    //             }
+    //         const testIdsString = testIds.join(' ');
+    //         console.debug('what to do with debug bool?', debugBool);
+    //         try {
+    //             execService?.exec(['-m', 'pytest', '-p', 'vscode_pytest', testIdsString].concat(pytestArgs), spawnOptions);
+    //         } catch (ex) {
+    //             console.error(ex);
     //         }
-    //         return this.deferred.promise;
+
+    //         return deferred.promise;
     //     }
     // }