Consistent disposal of receivers across adapters

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

@@ -33,27 +33,30 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
     async discoverTests(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
         const settings = this.configSettings.getSettings(uri);
+        const uuid = this.testServer.createUUID(uri.fsPath);
         const { pytestArgs } = settings.testing;
         traceVerbose(pytestArgs);
-        const disposable = this.testServer.onDiscoveryDataReceived((e: DataReceivedEvent) => {
-            // cancelation token ?
+        const dataReceivedDisposable = this.testServer.onDiscoveryDataReceived((e: DataReceivedEvent) => {
             this.resultResolver?.resolveDiscovery(JSON.parse(e.data));
         });
+        const disposeDataReceiver = function (testServer: ITestServer) {
+            testServer.deleteUUID(uuid);
+            dataReceivedDisposable.dispose();
+        };
         try {
-            await this.runPytestDiscovery(uri, executionFactory);
+            await this.runPytestDiscovery(uri, uuid, executionFactory);
         } finally {
-            disposable.dispose();
+            disposeDataReceiver(this.testServer);
         }
         // this is only a placeholder to handle function overloading until rewrite is finished
         const discoveryPayload: DiscoveredTestPayload = { cwd: uri.fsPath, status: 'success' };
         return discoveryPayload;
     }
 
-    async runPytestDiscovery(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<void> {
+    async runPytestDiscovery(uri: Uri, uuid: string, executionFactory?: IPythonExecutionFactory): Promise<void> {
         const 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 cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
@@ -85,7 +88,6 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
         result?.proc?.on('close', () => {
             deferredExec.resolve({ stdout: '', stderr: '' });
-            this.testServer.deleteUUID(uuid);
             deferred.resolve();
         });
         await deferredExec.promise;

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

@@ -43,19 +43,28 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
     ): Promise<ExecutionTestPayload> {
         const uuid = this.testServer.createUUID(uri.fsPath);
         traceVerbose(uri, testIds, debugBool);
-        const disposedDataReceived = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
+        const dataReceivedDisposable = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
             if (runInstance) {
                 this.resultResolver?.resolveExecution(JSON.parse(e.data), runInstance);
             }
         });
-        const dispose = function (testServer: ITestServer) {
+        const disposeDataReceiver = function (testServer: ITestServer) {
             testServer.deleteUUID(uuid);
-            disposedDataReceived.dispose();
+            dataReceivedDisposable.dispose();
         };
         runInstance?.token.onCancellationRequested(() => {
-            dispose(this.testServer);
+            disposeDataReceiver(this.testServer);
         });
-        await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, executionFactory, debugLauncher);
+        await this.runTestsNew(
+            uri,
+            testIds,
+            uuid,
+            runInstance,
+            debugBool,
+            executionFactory,
+            debugLauncher,
+            disposeDataReceiver,
+        );
 
         // placeholder until after the rewrite is adopted
         // TODO: remove after adoption.
@@ -75,6 +84,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         debugBool?: boolean,
         executionFactory?: IPythonExecutionFactory,
         debugLauncher?: ITestDebugLauncher,
+        disposeDataReceiver?: (testServer: ITestServer) => void,
     ): Promise<ExecutionTestPayload> {
         const deferred = createDeferred<ExecutionTestPayload>();
         const relativePathToPytest = 'pythonFiles';
@@ -158,8 +168,8 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
 
                 result?.proc?.on('close', () => {
                     deferredExec.resolve({ stdout: '', stderr: '' });
-                    this.testServer.deleteUUID(uuid);
                     deferred.resolve();
+                    disposeDataReceiver?.(this.testServer);
                 });
                 await deferredExec.promise;
             }

src/client/testing/testController/unittest/testDiscoveryAdapter.ts

@@ -43,13 +43,16 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             outChannel: this.outputChannel,
         };
 
-        const disposable = this.testServer.onDiscoveryDataReceived((e: DataReceivedEvent) => {
+        const dataReceivedDisposable = this.testServer.onDiscoveryDataReceived((e: DataReceivedEvent) => {
             this.resultResolver?.resolveDiscovery(JSON.parse(e.data));
         });
+        const disposeDataReceiver = function (testServer: ITestServer) {
+            testServer.deleteUUID(uuid);
+            dataReceivedDisposable.dispose();
+        };
 
         await this.callSendCommand(options, () => {
-            this.testServer.deleteUUID(uuid);
-            disposable.dispose();
+            disposeDataReceiver(this.testServer);
         });
         // placeholder until after the rewrite is adopted
         // TODO: remove after adoption.

src/client/testing/testController/unittest/testExecutionAdapter.ts

@@ -42,14 +42,14 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
                 this.resultResolver?.resolveExecution(JSON.parse(e.data), runInstance);
             }
         });
-        const dispose = function () {
+        const disposeDataReceiver = function (testServer: ITestServer) {
+            testServer.deleteUUID(uuid);
             disposedDataReceived.dispose();
         };
         runInstance?.token.onCancellationRequested(() => {
-            this.testServer.deleteUUID(uuid);
-            dispose();
+            disposeDataReceiver(this.testServer);
         });
-        await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, dispose);
+        await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, disposeDataReceiver);
         const executionPayload: ExecutionTestPayload = { cwd: uri.fsPath, status: 'success', error: '' };
         return executionPayload;
     }
@@ -60,7 +60,7 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         uuid: string,
         runInstance?: TestRun,
         debugBool?: boolean,
-        dispose?: () => void,
+        disposeDataReceiver?: (testServer: ITestServer) => void,
     ): Promise<ExecutionTestPayload> {
         const settings = this.configSettings.getSettings(uri);
         const { unittestArgs } = settings.testing;
@@ -84,9 +84,8 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         const runTestIdsPort = await startTestIdServer(testIds);
 
         await this.testServer.sendCommand(options, runTestIdsPort.toString(), runInstance, () => {
-            this.testServer.deleteUUID(uuid);
             deferred.resolve();
-            dispose?.();
+            disposeDataReceiver?.(this.testServer);
         });
         // placeholder until after the rewrite is adopted
         // TODO: remove after adoption.