From e2f4691ea07bb3c9b57060a88026040a1a0d98e9 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Thu, 20 Nov 2025 17:20:04 -0800
Subject: [PATCH 1/8] Refactor pytest test discovery

---
 .../pytest/pytestDiscoveryAdapter.ts          | 284 +++++++-----------
 .../testController/pytest/pytestHelpers.ts    | 134 +++++++++
 2 files changed, 244 insertions(+), 174 deletions(-)
 create mode 100644 src/client/testing/testController/pytest/pytestHelpers.ts

diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
index 308c9ba1f9bc..f33dee90d7a3 100644
--- a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 import * as path from 'path';
 import { CancellationToken, CancellationTokenSource, Uri } from 'vscode';
-import * as fs from 'fs';
 import { ChildProcess } from 'child_process';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
@@ -10,21 +9,20 @@ import {
     SpawnOptions,
 } from '../../../common/process/types';
 import { IConfigurationService } from '../../../common/types';
-import { createDeferred, Deferred } from '../../../common/utils/async';
+import { Deferred } from '../../../common/utils/async';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
-import { traceError, traceInfo, traceVerbose, traceWarn } from '../../../logging';
+import { traceError, traceInfo, traceVerbose } from '../../../logging';
 import { DiscoveredTestPayload, ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
-import {
-    createDiscoveryErrorPayload,
-    createTestingDeferred,
-    fixLogLinesNoTrailing,
-    startDiscoveryNamedPipe,
-    addValueIfKeyNotExist,
-    hasSymlinkParent,
-} from '../common/utils';
+import { createTestingDeferred, startDiscoveryNamedPipe } from '../common/utils';
 import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 import { PythonEnvironment } from '../../../pythonEnvironments/info';
 import { useEnvExtension, getEnvironment, runInBackground } from '../../../envExt/api.internal';
+import {
+    cleanupOnCancellation,
+    buildPytestEnv as configureSubprocessEnv,
+    createProcessHandlers,
+    handleSymlinkAndRootDir,
+} from './pytestHelpers';
 
 /**
  * Wrapper class for unittest test discovery. This is where we call `runTestCommand`. #this seems incorrectly copied
@@ -42,185 +40,123 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         token?: CancellationToken,
         interpreter?: PythonEnvironment,
     ): Promise<void> {
-        const cSource = new CancellationTokenSource();
-        const deferredReturn = createDeferred<void>();
-
-        token?.onCancellationRequested(() => {
-            traceInfo(`Test discovery cancelled.`);
-            cSource.cancel();
-            deferredReturn.resolve();
-        });
-
-        const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-            // if the token is cancelled, we don't want process the data
-            if (!token?.isCancellationRequested) {
-                this.resultResolver?.resolveDiscovery(data);
-            }
-        }, cSource.token);
-
-        this.runPytestDiscovery(uri, name, cSource, executionFactory, interpreter, token).then(() => {
-            deferredReturn.resolve();
-        });
-
-        return deferredReturn.promise;
-    }
-
-    async runPytestDiscovery(
-        uri: Uri,
-        discoveryPipeName: string,
-        cSource: CancellationTokenSource,
-        executionFactory: IPythonExecutionFactory,
-        interpreter?: PythonEnvironment,
-        token?: CancellationToken,
-    ): Promise<void> {
-        const relativePathToPytest = 'python_files';
-        const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
-        const settings = this.configSettings.getSettings(uri);
-        let { pytestArgs } = settings.testing;
-        const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
-
-        // check for symbolic path
-        const stats = await fs.promises.lstat(cwd);
-        const resolvedPath = await fs.promises.realpath(cwd);
-        let isSymbolicLink = false;
-        if (stats.isSymbolicLink()) {
-            isSymbolicLink = true;
-            traceWarn('The cwd is a symbolic link.');
-        } else if (resolvedPath !== cwd) {
-            traceWarn(
-                'The cwd resolves to a different path, checking if it has a symbolic link somewhere in its path.',
-            );
-            isSymbolicLink = await hasSymlinkParent(cwd);
-        }
-        if (isSymbolicLink) {
-            traceWarn("Symlink found, adding '--rootdir' to pytestArgs only if it doesn't already exist. cwd: ", cwd);
-            pytestArgs = addValueIfKeyNotExist(pytestArgs, '--rootdir', cwd);
-        }
-        // if user has provided `--rootdir` then use that, otherwise add `cwd`
-        // root dir is required so pytest can find the relative paths and for symlinks
-        addValueIfKeyNotExist(pytestArgs, '--rootdir', cwd);
-
-        // get and edit env vars
-        const mutableEnv = {
-            ...(await this.envVarsService?.getEnvironmentVariables(uri)),
-        };
-        // get python path from mutable env, it contains process.env as well
-        const pythonPathParts: string[] = mutableEnv.PYTHONPATH?.split(path.delimiter) ?? [];
-        const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
-        mutableEnv.PYTHONPATH = pythonPathCommand;
-        mutableEnv.TEST_RUN_PIPE = discoveryPipeName;
-        traceInfo(
-            `Environment variables set for pytest discovery: PYTHONPATH=${mutableEnv.PYTHONPATH}, TEST_RUN_PIPE=${mutableEnv.TEST_RUN_PIPE}`,
-        );
-
-        // delete UUID following entire discovery finishing.
-        const execArgs = ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs);
-        traceVerbose(`Running pytest discovery with command: ${execArgs.join(' ')} for workspace ${uri.fsPath}.`);
-
-        if (useEnvExtension()) {
-            const pythonEnv = await getEnvironment(uri);
-            if (pythonEnv) {
-                const deferredTillExecClose: Deferred<void> = createTestingDeferred();
+        // Setup: cancellation and discovery pipe
+        const discoveryPipeCancellation = new CancellationTokenSource();
+        try {
+            token?.onCancellationRequested(() => {
+                traceInfo(`Test discovery cancelled.`);
+                discoveryPipeCancellation.cancel();
+            });
+
+            const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+                if (!token?.isCancellationRequested) {
+                    this.resultResolver?.resolveDiscovery(data);
+                }
+            }, discoveryPipeCancellation.token);
+            traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+
+            // Build pytest command and arguments
+            const settings = this.configSettings.getSettings(uri);
+            let { pytestArgs } = settings.testing;
+            const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
+            pytestArgs = await handleSymlinkAndRootDir(cwd, pytestArgs);
+            const execArgs = ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs);
+            traceVerbose(`Running pytest discovery with command: ${execArgs.join(' ')} for workspace ${uri.fsPath}.`);
+
+            // Configure subprocess environment
+            const fullPluginPath = path.join(EXTENSION_ROOT_DIR, 'python_files');
+            const envVars = await this.envVarsService?.getEnvironmentVariables(uri);
+            const mutableEnv = await configureSubprocessEnv(envVars, fullPluginPath, discoveryPipeName);
+
+            // Setup process handlers (shared by both execution paths)
+            const deferredTillExecClose: Deferred<void> = createTestingDeferred();
+            const handlers = createProcessHandlers(uri, cwd, this.resultResolver, deferredTillExecClose);
+
+            // Execute using environment extension if available
+            if (useEnvExtension()) {
+                traceInfo(`Using environment extension for pytest discovery in workspace ${uri.fsPath}`);
+                const pythonEnv = await getEnvironment(uri);
+                if (!pythonEnv) {
+                    traceError(
+                        `Python environment not found for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
+                    );
+                    return;
+                }
+                traceVerbose(`Using Python environment: ${JSON.stringify(pythonEnv)}`);
 
                 const proc = await runInBackground(pythonEnv, {
                     cwd,
                     args: execArgs,
                     env: (mutableEnv as unknown) as { [key: string]: string },
                 });
+                traceInfo(`Started pytest discovery subprocess (environment extension) for workspace ${uri.fsPath}`);
+
+                // Wire up cancellation and process events
                 token?.onCancellationRequested(() => {
-                    traceInfo(`Test discovery cancelled, killing pytest subprocess for workspace ${uri.fsPath}`);
-                    proc.kill();
-                    deferredTillExecClose.resolve();
-                    cSource.cancel();
-                });
-                proc.stdout.on('data', (data) => {
-                    const out = fixLogLinesNoTrailing(data.toString());
-                    traceInfo(out);
-                });
-                proc.stderr.on('data', (data) => {
-                    const out = fixLogLinesNoTrailing(data.toString());
-                    traceError(out);
+                    cleanupOnCancellation(proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
+                proc.stdout.on('data', handlers.onStdout);
+                proc.stderr.on('data', handlers.onStderr);
                 proc.onExit((code, signal) => {
-                    if (code !== 0) {
-                        traceError(
-                            `Subprocess exited unsuccessfully with exit code ${code} and signal ${signal} on workspace ${uri.fsPath}`,
-                        );
-                        this.resultResolver?.resolveDiscovery(createDiscoveryErrorPayload(code, signal, cwd));
-                    }
-                    deferredTillExecClose.resolve();
+                    handlers.onExit(code, signal);
+                    handlers.onClose(code, signal);
                 });
+
                 await deferredTillExecClose.promise;
-            } else {
-                traceError(`Python Environment not found for: ${uri.fsPath}`);
+                traceInfo(`Pytest discovery completed for workspace ${uri.fsPath}`);
+                return;
             }
-            return;
-        }
 
-        const spawnOptions: SpawnOptions = {
-            cwd,
-            throwOnStdErr: true,
-            env: mutableEnv,
-            token,
-        };
-
-        // Create the Python environment in which to execute the command.
-        const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
-            allowEnvironmentFetchExceptions: false,
-            resource: uri,
-            interpreter,
-        };
-        const execService = await executionFactory.createActivatedEnvironment(creationOptions);
-
-        const execInfo = await execService?.getExecutablePath();
-        traceVerbose(`Executable path for pytest discovery: ${execInfo}.`);
-
-        const deferredTillExecClose: Deferred<void> = createTestingDeferred();
-
-        let resultProc: ChildProcess | undefined;
-
-        token?.onCancellationRequested(() => {
-            traceInfo(`Test discovery cancelled, killing pytest subprocess for workspace ${uri.fsPath}`);
-            // if the resultProc exists just call kill on it which will handle resolving the ExecClose deferred, otherwise resolve the deferred here.
-            if (resultProc) {
-                resultProc?.kill();
-            } else {
-                deferredTillExecClose.resolve();
-                cSource.cancel();
-            }
-        });
-        const result = execService?.execObservable(execArgs, spawnOptions);
-        resultProc = result?.proc;
-
-        // Take all output from the subprocess and add it to the test output channel. This will be the pytest output.
-        // Displays output to user and ensure the subprocess doesn't run into buffer overflow.
-
-        result?.proc?.stdout?.on('data', (data) => {
-            const out = fixLogLinesNoTrailing(data.toString());
-            traceInfo(out);
-        });
-        result?.proc?.stderr?.on('data', (data) => {
-            const out = fixLogLinesNoTrailing(data.toString());
-            traceError(out);
-        });
-        result?.proc?.on('exit', (code, signal) => {
-            if (code !== 0) {
+            // Execute using execution factory (fallback path)
+            traceInfo(`Using execution factory for pytest discovery in workspace ${uri.fsPath}`);
+            const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
+                allowEnvironmentFetchExceptions: false,
+                resource: uri,
+                interpreter,
+            };
+            const execService = await executionFactory.createActivatedEnvironment(creationOptions);
+            if (!execService) {
                 traceError(
-                    `Subprocess exited unsuccessfully with exit code ${code} and signal ${signal} on workspace ${uri.fsPath}.`,
+                    `Failed to create execution service for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
                 );
+                return;
             }
-        });
-        result?.proc?.on('close', (code, signal) => {
-            // pytest exits with code of 5 when 0 tests are found- this is not a failure for discovery.
-            if (code !== 0 && code !== 5) {
-                traceError(
-                    `Subprocess exited unsuccessfully with exit code ${code} and signal ${signal} on workspace ${uri.fsPath}. Creating and sending error discovery payload`,
-                );
-                this.resultResolver?.resolveDiscovery(createDiscoveryErrorPayload(code, signal, cwd));
+            const execInfo = await execService.getExecutablePath();
+            traceVerbose(`Using Python executable: ${execInfo} for workspace ${uri.fsPath}`);
+
+            const spawnOptions: SpawnOptions = {
+                cwd,
+                throwOnStdErr: true,
+                env: mutableEnv,
+                token,
+            };
+
+            let resultProc: ChildProcess | undefined;
+            const result = execService.execObservable(execArgs, spawnOptions);
+            resultProc = result?.proc;
+
+            if (!resultProc) {
+                traceError(`Failed to spawn pytest discovery subprocess for workspace ${uri.fsPath}`);
+                return;
             }
-            // due to the sync reading of the output.
-            deferredTillExecClose?.resolve();
-        });
-        await deferredTillExecClose.promise;
+            traceInfo(`Started pytest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
+
+            // Wire up cancellation and process events
+            token?.onCancellationRequested(() => {
+                cleanupOnCancellation(resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
+            });
+            resultProc.stdout?.on('data', handlers.onStdout);
+            resultProc.stderr?.on('data', handlers.onStderr);
+            resultProc.on('exit', handlers.onExit);
+            resultProc.on('close', handlers.onClose);
+
+            await deferredTillExecClose.promise;
+            traceInfo(`Pytest discovery completed for workspace ${uri.fsPath}`);
+        } catch (error) {
+            traceError(`Error during pytest discovery for workspace ${uri.fsPath}: ${error}`);
+            throw error;
+        } finally {
+            discoveryPipeCancellation.dispose();
+        }
     }
 }
diff --git a/src/client/testing/testController/pytest/pytestHelpers.ts b/src/client/testing/testController/pytest/pytestHelpers.ts
new file mode 100644
index 000000000000..1f3a415fc8cb
--- /dev/null
+++ b/src/client/testing/testController/pytest/pytestHelpers.ts
@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import * as path from 'path';
+import * as fs from 'fs';
+import { CancellationTokenSource, Uri } from 'vscode';
+import { Deferred } from '../../../common/utils/async';
+import { traceError, traceInfo, traceVerbose, traceWarn } from '../../../logging';
+import {
+    addValueIfKeyNotExist,
+    createDiscoveryErrorPayload,
+    fixLogLinesNoTrailing,
+    hasSymlinkParent,
+} from '../common/utils';
+import { ITestResultResolver } from '../common/types';
+
+/**
+ * Checks if the current working directory contains a symlink and ensures --rootdir is set in pytest args.
+ * This is required for pytest to correctly resolve relative paths in symlinked directories.
+ */
+export async function handleSymlinkAndRootDir(cwd: string, pytestArgs: string[]): Promise<string[]> {
+    const stats = await fs.promises.lstat(cwd);
+    const resolvedPath = await fs.promises.realpath(cwd);
+    let isSymbolicLink = false;
+    if (stats.isSymbolicLink()) {
+        isSymbolicLink = true;
+        traceWarn(`Working directory is a symbolic link: ${cwd} -> ${resolvedPath}`);
+    } else if (resolvedPath !== cwd) {
+        traceWarn(
+            `Working directory resolves to different path: ${cwd} -> ${resolvedPath}. Checking for symlinks in parent directories.`,
+        );
+        isSymbolicLink = await hasSymlinkParent(cwd);
+    }
+    if (isSymbolicLink) {
+        traceWarn(
+            `Symlink detected in path. Adding '--rootdir=${cwd}' to pytest args to ensure correct path resolution.`,
+        );
+        pytestArgs = addValueIfKeyNotExist(pytestArgs, '--rootdir', cwd);
+    }
+    // if user has provided `--rootdir` then use that, otherwise add `cwd`
+    // root dir is required so pytest can find the relative paths and for symlinks
+    addValueIfKeyNotExist(pytestArgs, '--rootdir', cwd);
+    return pytestArgs;
+}
+
+/**
+ * Builds the environment variables required for pytest discovery.
+ * Sets PYTHONPATH to include the plugin path and TEST_RUN_PIPE for communication.
+ */
+export async function buildPytestEnv(
+    envVars: { [key: string]: string | undefined } | undefined,
+    fullPluginPath: string,
+    discoveryPipeName: string,
+): Promise<{ [key: string]: string | undefined }> {
+    const mutableEnv = {
+        ...envVars,
+    };
+    // get python path from mutable env, it contains process.env as well
+    const pythonPathParts: string[] = mutableEnv.PYTHONPATH?.split(path.delimiter) ?? [];
+    const pythonPathCommand = [fullPluginPath, ...pythonPathParts].join(path.delimiter);
+    mutableEnv.PYTHONPATH = pythonPathCommand;
+    mutableEnv.TEST_RUN_PIPE = discoveryPipeName;
+    traceInfo(
+        `Environment variables set for pytest discovery: PYTHONPATH=${mutableEnv.PYTHONPATH}, TEST_RUN_PIPE=${mutableEnv.TEST_RUN_PIPE}`,
+    );
+    return mutableEnv;
+}
+
+/**
+ * Creates standard process event handlers for pytest discovery subprocess.
+ * Handles stdout/stderr logging and error reporting on process exit.
+ */
+export function createProcessHandlers(
+    uri: Uri,
+    cwd: string,
+    resultResolver: ITestResultResolver | undefined,
+    deferredTillExecClose: Deferred<void>,
+): {
+    onStdout: (data: any) => void;
+    onStderr: (data: any) => void;
+    onExit: (code: number | null, signal: NodeJS.Signals | null) => void;
+    onClose: (code: number | null, signal: NodeJS.Signals | null) => void;
+} {
+    return {
+        onStdout: (data: any) => {
+            const out = fixLogLinesNoTrailing(data.toString());
+            traceInfo(out);
+        },
+        onStderr: (data: any) => {
+            const out = fixLogLinesNoTrailing(data.toString());
+            traceError(out);
+        },
+        onExit: (code: number | null, signal: NodeJS.Signals | null) => {
+            // The 'exit' event fires when the process terminates, but streams may still be open.
+            if (code !== 0 && code !== 5) {
+                traceError(
+                    `Pytest discovery subprocess exited with code ${code} and signal ${signal} for workspace ${uri.fsPath}. Note: Exit code 5 (no tests found) is expected for empty test suites.`,
+                );
+            } else if (code === 0) {
+                traceVerbose(`Pytest discovery subprocess exited successfully for workspace ${uri.fsPath}`);
+            }
+        },
+        onClose: (code: number | null, signal: NodeJS.Signals | null) => {
+            // We resolve the deferred here to ensure all output has been captured.
+            // pytest exits with code of 5 when 0 tests are found- this is not a failure for discovery.
+            if (code !== 0 && code !== 5) {
+                traceError(
+                    `Pytest discovery failed with exit code ${code} and signal ${signal} for workspace ${uri.fsPath}. Creating error payload.`,
+                );
+                resultResolver?.resolveDiscovery(createDiscoveryErrorPayload(code, signal, cwd));
+            } else {
+                traceVerbose(`Pytest discovery subprocess streams closed for workspace ${uri.fsPath}`);
+            }
+            deferredTillExecClose?.resolve();
+        },
+    };
+}
+
+/**
+ * Handles cleanup when test discovery is cancelled.
+ * Kills the subprocess (if running), resolves the completion deferred, and cancels the discovery pipe.
+ */
+export function cleanupOnCancellation(
+    proc: { kill: () => void } | undefined,
+    processCompletion: Deferred<void>,
+    pipeCancellation: CancellationTokenSource,
+    uri: Uri,
+): void {
+    traceInfo(`Test discovery cancelled, killing pytest subprocess for workspace ${uri.fsPath}`);
+    if (proc) {
+        proc.kill();
+    }
+    processCompletion.resolve();
+    pipeCancellation.cancel();
+}

From 48a62d2b203d2aba0e6b6cd80287bb8980913ca2 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Thu, 20 Nov 2025 17:28:58 -0800
Subject: [PATCH 2/8] add unittest too

---
 .../testController/common/discoveryHelpers.ts | 103 +++++++
 .../pytest/pytestDiscoveryAdapter.ts          |  14 +-
 .../testController/pytest/pytestHelpers.ts    |  80 +-----
 .../unittest/testDiscoveryAdapter.ts          | 268 +++++++-----------
 .../unittest/unittestHelpers.ts               |  28 ++
 5 files changed, 242 insertions(+), 251 deletions(-)
 create mode 100644 src/client/testing/testController/common/discoveryHelpers.ts
 create mode 100644 src/client/testing/testController/unittest/unittestHelpers.ts

diff --git a/src/client/testing/testController/common/discoveryHelpers.ts b/src/client/testing/testController/common/discoveryHelpers.ts
new file mode 100644
index 000000000000..0e595dc5ea43
--- /dev/null
+++ b/src/client/testing/testController/common/discoveryHelpers.ts
@@ -0,0 +1,103 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import { CancellationTokenSource, Uri } from 'vscode';
+import { Deferred } from '../../../common/utils/async';
+import { traceError, traceInfo, traceVerbose } from '../../../logging';
+import { createDiscoveryErrorPayload, fixLogLinesNoTrailing } from './utils';
+import { ITestResultResolver } from './types';
+
+/**
+ * Test provider type for logging purposes.
+ */
+export type TestProvider = 'pytest' | 'unittest';
+
+/**
+ * Creates standard process event handlers for test discovery subprocess.
+ * Handles stdout/stderr logging and error reporting on process exit.
+ *
+ * @param testProvider - The test framework being used ('pytest' or 'unittest')
+ * @param uri - The workspace URI
+ * @param cwd - The current working directory
+ * @param resultResolver - Resolver for test discovery results
+ * @param deferredTillExecClose - Deferred to resolve when process closes
+ * @param allowedSuccessCodes - Additional exit codes to treat as success (e.g., pytest exit code 5 for no tests found)
+ */
+export function createProcessHandlers(
+    testProvider: TestProvider,
+    uri: Uri,
+    cwd: string,
+    resultResolver: ITestResultResolver | undefined,
+    deferredTillExecClose: Deferred<void>,
+    allowedSuccessCodes: number[] = [],
+): {
+    onStdout: (data: any) => void;
+    onStderr: (data: any) => void;
+    onExit: (code: number | null, signal: NodeJS.Signals | null) => void;
+    onClose: (code: number | null, signal: NodeJS.Signals | null) => void;
+} {
+    const isSuccessCode = (code: number | null): boolean => {
+        return code === 0 || (code !== null && allowedSuccessCodes.includes(code));
+    };
+
+    return {
+        onStdout: (data: any) => {
+            const out = fixLogLinesNoTrailing(data.toString());
+            traceInfo(out);
+        },
+        onStderr: (data: any) => {
+            const out = fixLogLinesNoTrailing(data.toString());
+            traceError(out);
+        },
+        onExit: (code: number | null, signal: NodeJS.Signals | null) => {
+            // The 'exit' event fires when the process terminates, but streams may still be open.
+            if (!isSuccessCode(code)) {
+                const exitCodeNote =
+                    allowedSuccessCodes.length > 0
+                        ? ` Note: Exit codes ${allowedSuccessCodes.join(', ')} are also treated as success.`
+                        : '';
+                traceError(
+                    `${testProvider} discovery subprocess exited with code ${code} and signal ${signal} for workspace ${uri.fsPath}.${exitCodeNote}`,
+                );
+            } else if (code === 0) {
+                traceVerbose(`${testProvider} discovery subprocess exited successfully for workspace ${uri.fsPath}`);
+            }
+        },
+        onClose: (code: number | null, signal: NodeJS.Signals | null) => {
+            // We resolve the deferred here to ensure all output has been captured.
+            if (!isSuccessCode(code)) {
+                traceError(
+                    `${testProvider} discovery failed with exit code ${code} and signal ${signal} for workspace ${uri.fsPath}. Creating error payload.`,
+                );
+                resultResolver?.resolveDiscovery(createDiscoveryErrorPayload(code, signal, cwd));
+            } else {
+                traceVerbose(`${testProvider} discovery subprocess streams closed for workspace ${uri.fsPath}`);
+            }
+            deferredTillExecClose?.resolve();
+        },
+    };
+}
+
+/**
+ * Handles cleanup when test discovery is cancelled.
+ * Kills the subprocess (if running), resolves the completion deferred, and cancels the discovery pipe.
+ *
+ * @param testProvider - The test framework being used ('pytest' or 'unittest')
+ * @param proc - The process to kill
+ * @param processCompletion - Deferred to resolve
+ * @param pipeCancellation - Cancellation token source to cancel
+ * @param uri - The workspace URI
+ */
+export function cleanupOnCancellation(
+    testProvider: TestProvider,
+    proc: { kill: () => void } | undefined,
+    processCompletion: Deferred<void>,
+    pipeCancellation: CancellationTokenSource,
+    uri: Uri,
+): void {
+    traceInfo(`Test discovery cancelled, killing ${testProvider} subprocess for workspace ${uri.fsPath}`);
+    if (proc) {
+        proc.kill();
+    }
+    processCompletion.resolve();
+    pipeCancellation.cancel();
+}
diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
index f33dee90d7a3..a7cfa5ea64d4 100644
--- a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -17,12 +17,8 @@ import { createTestingDeferred, startDiscoveryNamedPipe } from '../common/utils'
 import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 import { PythonEnvironment } from '../../../pythonEnvironments/info';
 import { useEnvExtension, getEnvironment, runInBackground } from '../../../envExt/api.internal';
-import {
-    cleanupOnCancellation,
-    buildPytestEnv as configureSubprocessEnv,
-    createProcessHandlers,
-    handleSymlinkAndRootDir,
-} from './pytestHelpers';
+import { buildPytestEnv as configureSubprocessEnv, handleSymlinkAndRootDir } from './pytestHelpers';
+import { cleanupOnCancellation, createProcessHandlers } from '../common/discoveryHelpers';
 
 /**
  * Wrapper class for unittest test discovery. This is where we call `runTestCommand`. #this seems incorrectly copied
@@ -70,7 +66,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
             // Setup process handlers (shared by both execution paths)
             const deferredTillExecClose: Deferred<void> = createTestingDeferred();
-            const handlers = createProcessHandlers(uri, cwd, this.resultResolver, deferredTillExecClose);
+            const handlers = createProcessHandlers('pytest', uri, cwd, this.resultResolver, deferredTillExecClose, [5]);
 
             // Execute using environment extension if available
             if (useEnvExtension()) {
@@ -93,7 +89,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
                 // Wire up cancellation and process events
                 token?.onCancellationRequested(() => {
-                    cleanupOnCancellation(proc, deferredTillExecClose, discoveryPipeCancellation, uri);
+                    cleanupOnCancellation('pytest', proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
                 proc.stdout.on('data', handlers.onStdout);
                 proc.stderr.on('data', handlers.onStderr);
@@ -143,7 +139,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
             // Wire up cancellation and process events
             token?.onCancellationRequested(() => {
-                cleanupOnCancellation(resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
+                cleanupOnCancellation('pytest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
             resultProc.stdout?.on('data', handlers.onStdout);
             resultProc.stderr?.on('data', handlers.onStderr);
diff --git a/src/client/testing/testController/pytest/pytestHelpers.ts b/src/client/testing/testController/pytest/pytestHelpers.ts
index 1f3a415fc8cb..3a2e4385d8e5 100644
--- a/src/client/testing/testController/pytest/pytestHelpers.ts
+++ b/src/client/testing/testController/pytest/pytestHelpers.ts
@@ -2,16 +2,8 @@
 // Licensed under the MIT License.
 import * as path from 'path';
 import * as fs from 'fs';
-import { CancellationTokenSource, Uri } from 'vscode';
-import { Deferred } from '../../../common/utils/async';
-import { traceError, traceInfo, traceVerbose, traceWarn } from '../../../logging';
-import {
-    addValueIfKeyNotExist,
-    createDiscoveryErrorPayload,
-    fixLogLinesNoTrailing,
-    hasSymlinkParent,
-} from '../common/utils';
-import { ITestResultResolver } from '../common/types';
+import { traceInfo, traceWarn } from '../../../logging';
+import { addValueIfKeyNotExist, hasSymlinkParent } from '../common/utils';
 
 /**
  * Checks if the current working directory contains a symlink and ensures --rootdir is set in pytest args.
@@ -64,71 +56,3 @@ export async function buildPytestEnv(
     );
     return mutableEnv;
 }
-
-/**
- * Creates standard process event handlers for pytest discovery subprocess.
- * Handles stdout/stderr logging and error reporting on process exit.
- */
-export function createProcessHandlers(
-    uri: Uri,
-    cwd: string,
-    resultResolver: ITestResultResolver | undefined,
-    deferredTillExecClose: Deferred<void>,
-): {
-    onStdout: (data: any) => void;
-    onStderr: (data: any) => void;
-    onExit: (code: number | null, signal: NodeJS.Signals | null) => void;
-    onClose: (code: number | null, signal: NodeJS.Signals | null) => void;
-} {
-    return {
-        onStdout: (data: any) => {
-            const out = fixLogLinesNoTrailing(data.toString());
-            traceInfo(out);
-        },
-        onStderr: (data: any) => {
-            const out = fixLogLinesNoTrailing(data.toString());
-            traceError(out);
-        },
-        onExit: (code: number | null, signal: NodeJS.Signals | null) => {
-            // The 'exit' event fires when the process terminates, but streams may still be open.
-            if (code !== 0 && code !== 5) {
-                traceError(
-                    `Pytest discovery subprocess exited with code ${code} and signal ${signal} for workspace ${uri.fsPath}. Note: Exit code 5 (no tests found) is expected for empty test suites.`,
-                );
-            } else if (code === 0) {
-                traceVerbose(`Pytest discovery subprocess exited successfully for workspace ${uri.fsPath}`);
-            }
-        },
-        onClose: (code: number | null, signal: NodeJS.Signals | null) => {
-            // We resolve the deferred here to ensure all output has been captured.
-            // pytest exits with code of 5 when 0 tests are found- this is not a failure for discovery.
-            if (code !== 0 && code !== 5) {
-                traceError(
-                    `Pytest discovery failed with exit code ${code} and signal ${signal} for workspace ${uri.fsPath}. Creating error payload.`,
-                );
-                resultResolver?.resolveDiscovery(createDiscoveryErrorPayload(code, signal, cwd));
-            } else {
-                traceVerbose(`Pytest discovery subprocess streams closed for workspace ${uri.fsPath}`);
-            }
-            deferredTillExecClose?.resolve();
-        },
-    };
-}
-
-/**
- * Handles cleanup when test discovery is cancelled.
- * Kills the subprocess (if running), resolves the completion deferred, and cancels the discovery pipe.
- */
-export function cleanupOnCancellation(
-    proc: { kill: () => void } | undefined,
-    processCompletion: Deferred<void>,
-    pipeCancellation: CancellationTokenSource,
-    uri: Uri,
-): void {
-    traceInfo(`Test discovery cancelled, killing pytest subprocess for workspace ${uri.fsPath}`);
-    if (proc) {
-        proc.kill();
-    }
-    processCompletion.resolve();
-    pipeCancellation.cancel();
-}
diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
index a40e25153fbc..fcc88aa39c16 100644
--- a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
+++ b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -1,33 +1,27 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import * as path from 'path';
-import { CancellationTokenSource, Uri } from 'vscode';
-import { CancellationToken } from 'vscode-jsonrpc';
+import { CancellationToken, CancellationTokenSource, Uri } from 'vscode';
 import { ChildProcess } from 'child_process';
 import { IConfigurationService } from '../../../common/types';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
-import {
-    DiscoveredTestPayload,
-    ITestDiscoveryAdapter,
-    ITestResultResolver,
-    TestCommandOptions,
-    TestDiscoveryCommand,
-} from '../common/types';
-import { createDeferred } from '../../../common/utils/async';
-import { EnvironmentVariables, IEnvironmentVariablesProvider } from '../../../common/variables/types';
+import { DiscoveredTestPayload, ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
+import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
-    ExecutionResult,
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
-import { createDiscoveryErrorPayload, fixLogLinesNoTrailing, startDiscoveryNamedPipe } from '../common/utils';
-import { traceError, traceInfo, traceLog, traceVerbose } from '../../../logging';
+import { startDiscoveryNamedPipe } from '../common/utils';
+import { traceError, traceInfo, traceVerbose } from '../../../logging';
 import { getEnvironment, runInBackground, useEnvExtension } from '../../../envExt/api.internal';
+import { PythonEnvironment } from '../../../pythonEnvironments/info';
+import { createTestingDeferred } from '../common/utils';
+import { buildDiscoveryCommand, buildUnittestEnv as configureSubprocessEnv } from './unittestHelpers';
+import { cleanupOnCancellation, createProcessHandlers } from '../common/discoveryHelpers';
 
 /**
- * Wrapper class for unittest test discovery. This is where we call `runTestCommand`.
+ * Wrapper class for unittest test discovery.
  */
 export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
     constructor(
@@ -36,181 +30,127 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         private readonly envVarsService?: IEnvironmentVariablesProvider,
     ) {}
 
-    public async discoverTests(
+    async discoverTests(
         uri: Uri,
         executionFactory: IPythonExecutionFactory,
         token?: CancellationToken,
+        interpreter?: PythonEnvironment,
     ): Promise<void> {
-        const settings = this.configSettings.getSettings(uri);
-        const { unittestArgs } = settings.testing;
-        const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
-
-        const cSource = new CancellationTokenSource();
-        // Create a deferred to return to the caller
-        const deferredReturn = createDeferred<void>();
-
-        token?.onCancellationRequested(() => {
-            traceInfo(`Test discovery cancelled.`);
-            cSource.cancel();
-            deferredReturn.resolve();
-        });
-
-        const name = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-            if (!token?.isCancellationRequested) {
-                this.resultResolver?.resolveDiscovery(data);
-            }
-        }, cSource.token);
-
-        // set up env with the pipe name
-        let env: EnvironmentVariables | undefined = await this.envVarsService?.getEnvironmentVariables(uri);
-        if (env === undefined) {
-            env = {} as EnvironmentVariables;
-        }
-        env.TEST_RUN_PIPE = name;
-
-        const command = buildDiscoveryCommand(unittestArgs);
-        const options: TestCommandOptions = {
-            workspaceFolder: uri,
-            command,
-            cwd,
-            token,
-        };
-
-        this.runDiscovery(uri, options, name, cwd, cSource, executionFactory).then(() => {
-            deferredReturn.resolve();
-        });
-
-        return deferredReturn.promise;
-    }
-
-    async runDiscovery(
-        uri: Uri,
-        options: TestCommandOptions,
-        testRunPipeName: string,
-        cwd: string,
-        cSource: CancellationTokenSource,
-        executionFactory: IPythonExecutionFactory,
-    ): Promise<void> {
-        // get and edit env vars
-        const mutableEnv = {
-            ...(await this.envVarsService?.getEnvironmentVariables(uri)),
-        };
-        mutableEnv.TEST_RUN_PIPE = testRunPipeName;
-        const args = [options.command.script].concat(options.command.args);
-
-        if (options.outChannel) {
-            options.outChannel.appendLine(`python ${args.join(' ')}`);
-        }
+        // Setup: cancellation and discovery pipe
+        const discoveryPipeCancellation = new CancellationTokenSource();
+        try {
+            token?.onCancellationRequested(() => {
+                traceInfo(`Test discovery cancelled.`);
+                discoveryPipeCancellation.cancel();
+            });
 
-        if (useEnvExtension()) {
-            const pythonEnv = await getEnvironment(uri);
-            if (pythonEnv) {
-                const deferredTillExecClose = createDeferred();
+            const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+                if (!token?.isCancellationRequested) {
+                    this.resultResolver?.resolveDiscovery(data);
+                }
+            }, discoveryPipeCancellation.token);
+            traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+
+            // Build unittest command and arguments
+            const settings = this.configSettings.getSettings(uri);
+            const { unittestArgs } = settings.testing;
+            const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
+            const execArgs = buildDiscoveryCommand(unittestArgs, EXTENSION_ROOT_DIR);
+            traceVerbose(`Running unittest discovery with command: ${execArgs.join(' ')} for workspace ${uri.fsPath}.`);
+
+            // Configure subprocess environment
+            const envVars = await this.envVarsService?.getEnvironmentVariables(uri);
+            const mutableEnv = await configureSubprocessEnv(envVars, discoveryPipeName);
+
+            // Setup process handlers (shared by both execution paths)
+            const deferredTillExecClose = createTestingDeferred();
+            const handlers = createProcessHandlers('unittest', uri, cwd, this.resultResolver, deferredTillExecClose);
+
+            // Execute using environment extension if available
+            if (useEnvExtension()) {
+                traceInfo(`Using environment extension for unittest discovery in workspace ${uri.fsPath}`);
+                const pythonEnv = await getEnvironment(uri);
+                if (!pythonEnv) {
+                    traceError(
+                        `Python environment not found for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
+                    );
+                    return;
+                }
+                traceVerbose(`Using Python environment: ${JSON.stringify(pythonEnv)}`);
 
                 const proc = await runInBackground(pythonEnv, {
                     cwd,
-                    args,
+                    args: execArgs,
                     env: (mutableEnv as unknown) as { [key: string]: string },
                 });
-                options.token?.onCancellationRequested(() => {
-                    traceInfo(`Test discovery cancelled, killing unittest subprocess for workspace ${uri.fsPath}`);
-                    proc.kill();
-                    deferredTillExecClose.resolve();
-                    cSource.cancel();
-                });
-                proc.stdout.on('data', (data) => {
-                    const out = fixLogLinesNoTrailing(data.toString());
-                    traceInfo(out);
-                });
-                proc.stderr.on('data', (data) => {
-                    const out = fixLogLinesNoTrailing(data.toString());
-                    traceError(out);
+                traceInfo(`Started unittest discovery subprocess (environment extension) for workspace ${uri.fsPath}`);
+
+                // Wire up cancellation and process events
+                token?.onCancellationRequested(() => {
+                    cleanupOnCancellation('unittest', proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
+                proc.stdout.on('data', handlers.onStdout);
+                proc.stderr.on('data', handlers.onStderr);
                 proc.onExit((code, signal) => {
-                    if (code !== 0) {
-                        traceError(
-                            `Subprocess exited unsuccessfully with exit code ${code} and signal ${signal} on workspace ${uri.fsPath}`,
-                        );
-                    }
-                    deferredTillExecClose.resolve();
+                    handlers.onExit(code, signal);
+                    handlers.onClose(code, signal);
                 });
+
                 await deferredTillExecClose.promise;
-            } else {
-                traceError(`Python Environment not found for: ${uri.fsPath}`);
+                traceInfo(`Unittest discovery completed for workspace ${uri.fsPath}`);
+                return;
             }
-            return;
-        }
-
-        const spawnOptions: SpawnOptions = {
-            token: options.token,
-            cwd: options.cwd,
-            throwOnStdErr: true,
-            env: mutableEnv,
-        };
 
-        try {
-            traceLog(`Discovering unittest tests for workspace ${options.cwd} with arguments: ${args}\r\n`);
-            const deferredTillExecClose = createDeferred<ExecutionResult<string>>();
-
-            // Create the Python environment in which to execute the command.
+            // Execute using execution factory (fallback path)
+            traceInfo(`Using execution factory for unittest discovery in workspace ${uri.fsPath}`);
             const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
                 allowEnvironmentFetchExceptions: false,
-                resource: options.workspaceFolder,
+                resource: uri,
+                interpreter,
             };
             const execService = await executionFactory.createActivatedEnvironment(creationOptions);
-            const execInfo = await execService?.getExecutablePath();
-            traceVerbose(`Executable path for unittest discovery: ${execInfo}.`);
+            if (!execService) {
+                traceError(
+                    `Failed to create execution service for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
+                );
+                return;
+            }
+            const execInfo = await execService.getExecutablePath();
+            traceVerbose(`Using Python executable: ${execInfo} for workspace ${uri.fsPath}`);
+
+            const spawnOptions: SpawnOptions = {
+                cwd,
+                throwOnStdErr: true,
+                env: mutableEnv,
+                token,
+            };
 
             let resultProc: ChildProcess | undefined;
-            options.token?.onCancellationRequested(() => {
-                traceInfo(`Test discovery cancelled, killing unittest subprocess for workspace ${uri.fsPath}`);
-                // if the resultProc exists just call kill on it which will handle resolving the ExecClose deferred, otherwise resolve the deferred here.
-                if (resultProc) {
-                    resultProc?.kill();
-                } else {
-                    deferredTillExecClose.resolve();
-                    cSource.cancel();
-                }
-            });
-            const result = execService?.execObservable(args, spawnOptions);
+            const result = execService.execObservable(execArgs, spawnOptions);
             resultProc = result?.proc;
 
-            // Displays output to user and ensure the subprocess doesn't run into buffer overflow.
-            result?.proc?.stdout?.on('data', (data) => {
-                const out = fixLogLinesNoTrailing(data.toString());
-                traceInfo(out);
-            });
-            result?.proc?.stderr?.on('data', (data) => {
-                const out = fixLogLinesNoTrailing(data.toString());
-                traceError(out);
-            });
-
-            result?.proc?.on('exit', (code, signal) => {
-                // if the child has testIds then this is a run request
+            if (!resultProc) {
+                traceError(`Failed to spawn unittest discovery subprocess for workspace ${uri.fsPath}`);
+                return;
+            }
+            traceInfo(`Started unittest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
 
-                if (code !== 0) {
-                    // This occurs when we are running discovery
-                    traceError(
-                        `Subprocess exited unsuccessfully with exit code ${code} and signal ${signal} on workspace ${options.cwd}. Creating and sending error discovery payload \n`,
-                    );
-                    traceError(
-                        `Subprocess exited unsuccessfully with exit code ${code} and signal ${signal} on workspace ${uri.fsPath}. Creating and sending error discovery payload`,
-                    );
-                    this.resultResolver?.resolveDiscovery(createDiscoveryErrorPayload(code, signal, cwd));
-                }
-                deferredTillExecClose.resolve();
+            // Wire up cancellation and process events
+            token?.onCancellationRequested(() => {
+                cleanupOnCancellation('unittest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
+            resultProc.stdout?.on('data', handlers.onStdout);
+            resultProc.stderr?.on('data', handlers.onStderr);
+            resultProc.on('exit', handlers.onExit);
+            resultProc.on('close', handlers.onClose);
+
             await deferredTillExecClose.promise;
-        } catch (ex) {
-            traceError(`Error while server attempting to run unittest command for workspace ${uri.fsPath}: ${ex}`);
+            traceInfo(`Unittest discovery completed for workspace ${uri.fsPath}`);
+        } catch (error) {
+            traceError(`Error during unittest discovery for workspace ${uri.fsPath}: ${error}`);
+            throw error;
+        } finally {
+            discoveryPipeCancellation.dispose();
         }
     }
 }
-function buildDiscoveryCommand(args: string[]): TestDiscoveryCommand {
-    const discoveryScript = path.join(EXTENSION_ROOT_DIR, 'python_files', 'unittestadapter', 'discovery.py');
-
-    return {
-        script: discoveryScript,
-        args: ['--udiscovery', ...args],
-    };
-}
diff --git a/src/client/testing/testController/unittest/unittestHelpers.ts b/src/client/testing/testController/unittest/unittestHelpers.ts
new file mode 100644
index 000000000000..6668c1df0e65
--- /dev/null
+++ b/src/client/testing/testController/unittest/unittestHelpers.ts
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import * as path from 'path';
+import { traceInfo } from '../../../logging';
+
+/**
+ * Builds the environment variables required for unittest discovery.
+ * Sets TEST_RUN_PIPE for communication.
+ */
+export async function buildUnittestEnv(
+    envVars: { [key: string]: string | undefined } | undefined,
+    discoveryPipeName: string,
+): Promise<{ [key: string]: string | undefined }> {
+    const mutableEnv = {
+        ...envVars,
+    };
+    mutableEnv.TEST_RUN_PIPE = discoveryPipeName;
+    traceInfo(`Environment variables set for unittest discovery: TEST_RUN_PIPE=${mutableEnv.TEST_RUN_PIPE}`);
+    return mutableEnv;
+}
+
+/**
+ * Builds the unittest discovery command.
+ */
+export function buildDiscoveryCommand(args: string[], extensionRootDir: string): string[] {
+    const discoveryScript = path.join(extensionRootDir, 'python_files', 'unittestadapter', 'discovery.py');
+    return [discoveryScript, '--udiscovery', ...args];
+}

From 3d7d0c6ad9cc0a463fef7767e6eef159e4e7594f Mon Sep 17 00:00:00 2001
From: eleanorjboyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Thu, 20 Nov 2025 17:44:09 -0800
Subject: [PATCH 3/8] more refactoring

---
 .../pytest/pytestDiscoveryAdapter.ts          | 90 ++++++++++++++-----
 .../unittest/testDiscoveryAdapter.ts          | 76 ++++++++++++----
 2 files changed, 128 insertions(+), 38 deletions(-)

diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
index a7cfa5ea64d4..c3e9d41a4f0b 100644
--- a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -21,7 +21,60 @@ import { buildPytestEnv as configureSubprocessEnv, handleSymlinkAndRootDir } fro
 import { cleanupOnCancellation, createProcessHandlers } from '../common/discoveryHelpers';
 
 /**
- * Wrapper class for unittest test discovery. This is where we call `runTestCommand`. #this seems incorrectly copied
+ * Sets up the discovery named pipe and wires up cancellation.
+ * @param resultResolver The resolver to handle discovered test data
+ * @param token Optional cancellation token from the caller
+ * @param uri Workspace URI for logging
+ * @returns Object containing the pipe name and cancellation source
+ */
+async function setupDiscoveryPipe(
+    resultResolver: ITestResultResolver | undefined,
+    token: CancellationToken | undefined,
+    uri: Uri,
+): Promise<{ pipeName: string; cancellation: CancellationTokenSource }> {
+    const discoveryPipeCancellation = new CancellationTokenSource();
+
+    // Wire up cancellation from external token
+    token?.onCancellationRequested(() => {
+        traceInfo(`Test discovery cancelled.`);
+        discoveryPipeCancellation.cancel();
+    });
+
+    // Start the named pipe with the discovery listener
+    const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+        if (!token?.isCancellationRequested) {
+            resultResolver?.resolveDiscovery(data);
+        }
+    }, discoveryPipeCancellation.token);
+
+    traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+
+    return {
+        pipeName: discoveryPipeName,
+        cancellation: discoveryPipeCancellation,
+    };
+}
+
+/**
+ * Configures the subprocess environment for pytest discovery.
+ * @param envVarsService Service to retrieve environment variables
+ * @param uri Workspace URI
+ * @param discoveryPipeName Name of the discovery pipe to pass to the subprocess
+ * @returns Configured environment variables for the subprocess
+ */
+async function configureDiscoveryEnv(
+    envVarsService: IEnvironmentVariablesProvider | undefined,
+    uri: Uri,
+    discoveryPipeName: string,
+): Promise<NodeJS.ProcessEnv> {
+    const fullPluginPath = path.join(EXTENSION_ROOT_DIR, 'python_files');
+    const envVars = await envVarsService?.getEnvironmentVariables(uri);
+    const mutableEnv = await configureSubprocessEnv(envVars, fullPluginPath, discoveryPipeName);
+    return mutableEnv;
+}
+
+/**
+ * Wrapper class for pytest test discovery. This is where we call the pytest subprocess.
  */
 export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
     constructor(
@@ -36,33 +89,26 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         token?: CancellationToken,
         interpreter?: PythonEnvironment,
     ): Promise<void> {
-        // Setup: cancellation and discovery pipe
-        const discoveryPipeCancellation = new CancellationTokenSource();
-        try {
-            token?.onCancellationRequested(() => {
-                traceInfo(`Test discovery cancelled.`);
-                discoveryPipeCancellation.cancel();
-            });
-
-            const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-                if (!token?.isCancellationRequested) {
-                    this.resultResolver?.resolveDiscovery(data);
-                }
-            }, discoveryPipeCancellation.token);
-            traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+        // Setup discovery pipe and cancellation
+        const { pipeName: discoveryPipeName, cancellation: discoveryPipeCancellation } = await setupDiscoveryPipe(
+            this.resultResolver,
+            token,
+            uri,
+        );
 
+        try {
             // Build pytest command and arguments
             const settings = this.configSettings.getSettings(uri);
             let { pytestArgs } = settings.testing;
             const cwd = settings.testing.cwd && settings.testing.cwd.length > 0 ? settings.testing.cwd : uri.fsPath;
             pytestArgs = await handleSymlinkAndRootDir(cwd, pytestArgs);
-            const execArgs = ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs);
-            traceVerbose(`Running pytest discovery with command: ${execArgs.join(' ')} for workspace ${uri.fsPath}.`);
+            const commandArgs = ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs);
+            traceVerbose(
+                `Running pytest discovery with command: ${commandArgs.join(' ')} for workspace ${uri.fsPath}.`,
+            );
 
             // Configure subprocess environment
-            const fullPluginPath = path.join(EXTENSION_ROOT_DIR, 'python_files');
-            const envVars = await this.envVarsService?.getEnvironmentVariables(uri);
-            const mutableEnv = await configureSubprocessEnv(envVars, fullPluginPath, discoveryPipeName);
+            const mutableEnv = await configureDiscoveryEnv(this.envVarsService, uri, discoveryPipeName);
 
             // Setup process handlers (shared by both execution paths)
             const deferredTillExecClose: Deferred<void> = createTestingDeferred();
@@ -82,7 +128,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
                 const proc = await runInBackground(pythonEnv, {
                     cwd,
-                    args: execArgs,
+                    args: commandArgs,
                     env: (mutableEnv as unknown) as { [key: string]: string },
                 });
                 traceInfo(`Started pytest discovery subprocess (environment extension) for workspace ${uri.fsPath}`);
@@ -128,7 +174,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
-            const result = execService.execObservable(execArgs, spawnOptions);
+            const result = execService.execObservable(commandArgs, spawnOptions);
             resultProc = result?.proc;
 
             if (!resultProc) {
diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
index fcc88aa39c16..981414feb14d 100644
--- a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
+++ b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -20,6 +20,58 @@ import { createTestingDeferred } from '../common/utils';
 import { buildDiscoveryCommand, buildUnittestEnv as configureSubprocessEnv } from './unittestHelpers';
 import { cleanupOnCancellation, createProcessHandlers } from '../common/discoveryHelpers';
 
+/**
+ * Sets up the discovery named pipe and wires up cancellation.
+ * @param resultResolver The resolver to handle discovered test data
+ * @param token Optional cancellation token from the caller
+ * @param uri Workspace URI for logging
+ * @returns Object containing the pipe name and cancellation source
+ */
+async function setupDiscoveryPipe(
+    resultResolver: ITestResultResolver | undefined,
+    token: CancellationToken | undefined,
+    uri: Uri,
+): Promise<{ pipeName: string; cancellation: CancellationTokenSource }> {
+    const discoveryPipeCancellation = new CancellationTokenSource();
+
+    // Wire up cancellation from external token
+    token?.onCancellationRequested(() => {
+        traceInfo(`Test discovery cancelled.`);
+        discoveryPipeCancellation.cancel();
+    });
+
+    // Start the named pipe with the discovery listener
+    const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+        if (!token?.isCancellationRequested) {
+            resultResolver?.resolveDiscovery(data);
+        }
+    }, discoveryPipeCancellation.token);
+
+    traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+
+    return {
+        pipeName: discoveryPipeName,
+        cancellation: discoveryPipeCancellation,
+    };
+}
+
+/**
+ * Configures the subprocess environment for unittest discovery.
+ * @param envVarsService Service to retrieve environment variables
+ * @param uri Workspace URI
+ * @param discoveryPipeName Name of the discovery pipe to pass to the subprocess
+ * @returns Configured environment variables for the subprocess
+ */
+async function configureDiscoveryEnv(
+    envVarsService: IEnvironmentVariablesProvider | undefined,
+    uri: Uri,
+    discoveryPipeName: string,
+): Promise<NodeJS.ProcessEnv> {
+    const envVars = await envVarsService?.getEnvironmentVariables(uri);
+    const mutableEnv = await configureSubprocessEnv(envVars, discoveryPipeName);
+    return mutableEnv;
+}
+
 /**
  * Wrapper class for unittest test discovery.
  */
@@ -36,21 +88,14 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         token?: CancellationToken,
         interpreter?: PythonEnvironment,
     ): Promise<void> {
-        // Setup: cancellation and discovery pipe
-        const discoveryPipeCancellation = new CancellationTokenSource();
-        try {
-            token?.onCancellationRequested(() => {
-                traceInfo(`Test discovery cancelled.`);
-                discoveryPipeCancellation.cancel();
-            });
-
-            const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-                if (!token?.isCancellationRequested) {
-                    this.resultResolver?.resolveDiscovery(data);
-                }
-            }, discoveryPipeCancellation.token);
-            traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+        // Setup discovery pipe and cancellation
+        const { pipeName: discoveryPipeName, cancellation: discoveryPipeCancellation } = await setupDiscoveryPipe(
+            this.resultResolver,
+            token,
+            uri,
+        );
 
+        try {
             // Build unittest command and arguments
             const settings = this.configSettings.getSettings(uri);
             const { unittestArgs } = settings.testing;
@@ -59,8 +104,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             traceVerbose(`Running unittest discovery with command: ${execArgs.join(' ')} for workspace ${uri.fsPath}.`);
 
             // Configure subprocess environment
-            const envVars = await this.envVarsService?.getEnvironmentVariables(uri);
-            const mutableEnv = await configureSubprocessEnv(envVars, discoveryPipeName);
+            const mutableEnv = await configureDiscoveryEnv(this.envVarsService, uri, discoveryPipeName);
 
             // Setup process handlers (shared by both execution paths)
             const deferredTillExecClose = createTestingDeferred();

From db97ada8eb8d1b1091c068b2ecfdc3fd4d604e61 Mon Sep 17 00:00:00 2001
From: eleanorjboyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Thu, 20 Nov 2025 17:52:42 -0800
Subject: [PATCH 4/8] fix test

---
 .../testController/common/discoveryHelpers.ts |  5 ++
 .../pytest/pytestDiscoveryAdapter.ts          | 54 +++++++++++++++----
 .../unittest/testDiscoveryAdapter.ts          | 54 +++++++++++++++----
 3 files changed, 91 insertions(+), 22 deletions(-)

diff --git a/src/client/testing/testController/common/discoveryHelpers.ts b/src/client/testing/testController/common/discoveryHelpers.ts
index 0e595dc5ea43..b79bedd9a2b5 100644
--- a/src/client/testing/testController/common/discoveryHelpers.ts
+++ b/src/client/testing/testController/common/discoveryHelpers.ts
@@ -96,8 +96,13 @@ export function cleanupOnCancellation(
 ): void {
     traceInfo(`Test discovery cancelled, killing ${testProvider} subprocess for workspace ${uri.fsPath}`);
     if (proc) {
+        traceVerbose(`Killing ${testProvider} subprocess for workspace ${uri.fsPath}`);
         proc.kill();
+    } else {
+        traceVerbose(`No ${testProvider} subprocess to kill for workspace ${uri.fsPath} (proc is undefined)`);
     }
+    traceVerbose(`Resolving process completion deferred for ${testProvider} discovery in workspace ${uri.fsPath}`);
     processCompletion.resolve();
+    traceVerbose(`Cancelling discovery pipe for ${testProvider} discovery in workspace ${uri.fsPath}`);
     pipeCancellation.cancel();
 }
diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
index c3e9d41a4f0b..9756d2783b6b 100644
--- a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -96,6 +96,9 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             uri,
         );
 
+        // Setup process handlers deferred (used by both execution paths)
+        const deferredTillExecClose: Deferred<void> = createTestingDeferred();
+
         try {
             // Build pytest command and arguments
             const settings = this.configSettings.getSettings(uri);
@@ -111,7 +114,6 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             const mutableEnv = await configureDiscoveryEnv(this.envVarsService, uri, discoveryPipeName);
 
             // Setup process handlers (shared by both execution paths)
-            const deferredTillExecClose: Deferred<void> = createTestingDeferred();
             const handlers = createProcessHandlers('pytest', uri, cwd, this.resultResolver, deferredTillExecClose, [5]);
 
             // Execute using environment extension if available
@@ -122,6 +124,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                     traceError(
                         `Python environment not found for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
                     );
+                    deferredTillExecClose.resolve();
                     return;
                 }
                 traceVerbose(`Using Python environment: ${JSON.stringify(pythonEnv)}`);
@@ -161,11 +164,19 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 traceError(
                     `Failed to create execution service for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
                 );
+                deferredTillExecClose.resolve();
                 return;
             }
             const execInfo = await execService.getExecutablePath();
             traceVerbose(`Using Python executable: ${execInfo} for workspace ${uri.fsPath}`);
 
+            // Check for cancellation before spawning process
+            if (token?.isCancellationRequested) {
+                traceInfo(`Pytest discovery cancelled before spawning process for workspace ${uri.fsPath}`);
+                deferredTillExecClose.resolve();
+                return;
+            }
+
             const spawnOptions: SpawnOptions = {
                 cwd,
                 throwOnStdErr: true,
@@ -174,30 +185,51 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
-            const result = execService.execObservable(commandArgs, spawnOptions);
-            resultProc = result?.proc;
 
-            if (!resultProc) {
-                traceError(`Failed to spawn pytest discovery subprocess for workspace ${uri.fsPath}`);
-                return;
-            }
-            traceInfo(`Started pytest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
-
-            // Wire up cancellation and process events
-            token?.onCancellationRequested(() => {
+            // Set up cancellation handler before execObservable to catch early cancellations
+            const cancellationHandler = token?.onCancellationRequested(() => {
+                traceInfo(`Cancellation requested during pytest discovery for workspace ${uri.fsPath}`);
                 cleanupOnCancellation('pytest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
+
+            try {
+                const result = execService.execObservable(commandArgs, spawnOptions);
+                resultProc = result?.proc;
+
+                if (!resultProc) {
+                    traceError(`Failed to spawn pytest discovery subprocess for workspace ${uri.fsPath}`);
+                    deferredTillExecClose.resolve();
+                    return;
+                }
+                traceInfo(`Started pytest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
+            } catch (error) {
+                traceError(`Error spawning pytest discovery subprocess for workspace ${uri.fsPath}: ${error}`);
+                deferredTillExecClose.resolve();
+                cancellationHandler?.dispose();
+                throw error;
+            }
             resultProc.stdout?.on('data', handlers.onStdout);
             resultProc.stderr?.on('data', handlers.onStderr);
             resultProc.on('exit', handlers.onExit);
             resultProc.on('close', handlers.onClose);
 
+            cancellationHandler?.dispose();
+
+            // Check for early cancellation before awaitingre awaiting
+            if (token?.isCancellationRequested) {
+                traceInfo(`Pytest discovery was cancelled before process completion for workspace ${uri.fsPath}`);
+                deferredTillExecClose.resolve();
+            }
+
+            traceVerbose(`Waiting for pytest discovery subprocess to complete for workspace ${uri.fsPath}`);
             await deferredTillExecClose.promise;
             traceInfo(`Pytest discovery completed for workspace ${uri.fsPath}`);
         } catch (error) {
             traceError(`Error during pytest discovery for workspace ${uri.fsPath}: ${error}`);
+            deferredTillExecClose.resolve();
             throw error;
         } finally {
+            traceVerbose(`Cleaning up pytest discovery resources for workspace ${uri.fsPath}`);
             discoveryPipeCancellation.dispose();
         }
     }
diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
index 981414feb14d..d0a992c86a58 100644
--- a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
+++ b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -95,6 +95,9 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             uri,
         );
 
+        // Setup process handlers deferred (used by both execution paths)
+        const deferredTillExecClose = createTestingDeferred();
+
         try {
             // Build unittest command and arguments
             const settings = this.configSettings.getSettings(uri);
@@ -107,7 +110,6 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             const mutableEnv = await configureDiscoveryEnv(this.envVarsService, uri, discoveryPipeName);
 
             // Setup process handlers (shared by both execution paths)
-            const deferredTillExecClose = createTestingDeferred();
             const handlers = createProcessHandlers('unittest', uri, cwd, this.resultResolver, deferredTillExecClose);
 
             // Execute using environment extension if available
@@ -118,6 +120,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                     traceError(
                         `Python environment not found for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
                     );
+                    deferredTillExecClose.resolve();
                     return;
                 }
                 traceVerbose(`Using Python environment: ${JSON.stringify(pythonEnv)}`);
@@ -157,11 +160,19 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 traceError(
                     `Failed to create execution service for workspace ${uri.fsPath}. Cannot proceed with test discovery.`,
                 );
+                deferredTillExecClose.resolve();
                 return;
             }
             const execInfo = await execService.getExecutablePath();
             traceVerbose(`Using Python executable: ${execInfo} for workspace ${uri.fsPath}`);
 
+            // Check for cancellation before spawning process
+            if (token?.isCancellationRequested) {
+                traceInfo(`Unittest discovery cancelled before spawning process for workspace ${uri.fsPath}`);
+                deferredTillExecClose.resolve();
+                return;
+            }
+
             const spawnOptions: SpawnOptions = {
                 cwd,
                 throwOnStdErr: true,
@@ -170,30 +181,51 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
-            const result = execService.execObservable(execArgs, spawnOptions);
-            resultProc = result?.proc;
 
-            if (!resultProc) {
-                traceError(`Failed to spawn unittest discovery subprocess for workspace ${uri.fsPath}`);
-                return;
-            }
-            traceInfo(`Started unittest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
-
-            // Wire up cancellation and process events
-            token?.onCancellationRequested(() => {
+            // Set up cancellation handler before execObservable to catch early cancellations
+            const cancellationHandler = token?.onCancellationRequested(() => {
+                traceInfo(`Cancellation requested during unittest discovery for workspace ${uri.fsPath}`);
                 cleanupOnCancellation('unittest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
+
+            try {
+                const result = execService.execObservable(execArgs, spawnOptions);
+                resultProc = result?.proc;
+
+                if (!resultProc) {
+                    traceError(`Failed to spawn unittest discovery subprocess for workspace ${uri.fsPath}`);
+                    deferredTillExecClose.resolve();
+                    return;
+                }
+                traceInfo(`Started unittest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
+            } catch (error) {
+                traceError(`Error spawning unittest discovery subprocess for workspace ${uri.fsPath}: ${error}`);
+                deferredTillExecClose.resolve();
+                cancellationHandler?.dispose();
+                throw error;
+            }
             resultProc.stdout?.on('data', handlers.onStdout);
             resultProc.stderr?.on('data', handlers.onStderr);
             resultProc.on('exit', handlers.onExit);
             resultProc.on('close', handlers.onClose);
 
+            cancellationHandler?.dispose();
+
+            // Check for early cancellation before awaitingre awaiting
+            if (token?.isCancellationRequested) {
+                traceInfo(`Unittest discovery was cancelled before process completion for workspace ${uri.fsPath}`);
+                deferredTillExecClose.resolve();
+            }
+
+            traceVerbose(`Waiting for unittest discovery subprocess to complete for workspace ${uri.fsPath}`);
             await deferredTillExecClose.promise;
             traceInfo(`Unittest discovery completed for workspace ${uri.fsPath}`);
         } catch (error) {
             traceError(`Error during unittest discovery for workspace ${uri.fsPath}: ${error}`);
+            deferredTillExecClose.resolve();
             throw error;
         } finally {
+            traceVerbose(`Cleaning up unittest discovery resources for workspace ${uri.fsPath}`);
             discoveryPipeCancellation.dispose();
         }
     }

From 0977cf3c9f98a5797227367fa0c261314d7d8d27 Mon Sep 17 00:00:00 2001
From: Eleanor Boyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Thu, 20 Nov 2025 18:01:22 -0800
Subject: [PATCH 5/8] Update
 src/client/testing/testController/unittest/testDiscoveryAdapter.ts

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 .../testing/testController/unittest/testDiscoveryAdapter.ts     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
index d0a992c86a58..45b70cbd3d53 100644
--- a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
+++ b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -211,7 +211,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
             cancellationHandler?.dispose();
 
-            // Check for early cancellation before awaitingre awaiting
+            // Check for early cancellation before awaiting
             if (token?.isCancellationRequested) {
                 traceInfo(`Unittest discovery was cancelled before process completion for workspace ${uri.fsPath}`);
                 deferredTillExecClose.resolve();

From 07bce782547b955bf05b1d1bd2ba673490feeaa8 Mon Sep 17 00:00:00 2001
From: Eleanor Boyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Thu, 20 Nov 2025 18:01:50 -0800
Subject: [PATCH 6/8] Update
 src/client/testing/testController/pytest/pytestHelpers.ts

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 src/client/testing/testController/pytest/pytestHelpers.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/client/testing/testController/pytest/pytestHelpers.ts b/src/client/testing/testController/pytest/pytestHelpers.ts
index 3a2e4385d8e5..ca55cbf778ec 100644
--- a/src/client/testing/testController/pytest/pytestHelpers.ts
+++ b/src/client/testing/testController/pytest/pytestHelpers.ts
@@ -38,11 +38,11 @@ export async function handleSymlinkAndRootDir(cwd: string, pytestArgs: string[])
  * Builds the environment variables required for pytest discovery.
  * Sets PYTHONPATH to include the plugin path and TEST_RUN_PIPE for communication.
  */
-export async function buildPytestEnv(
+export function buildPytestEnv(
     envVars: { [key: string]: string | undefined } | undefined,
     fullPluginPath: string,
     discoveryPipeName: string,
-): Promise<{ [key: string]: string | undefined }> {
+): { [key: string]: string | undefined } {
     const mutableEnv = {
         ...envVars,
     };

From b313fc7e50df7c9a5498805ddb8ab0ef0dbdbd9d Mon Sep 17 00:00:00 2001
From: eleanorjboyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Fri, 21 Nov 2025 08:58:21 -0800
Subject: [PATCH 7/8] fixes based on comments

---
 .../testController/common/discoveryHelpers.ts | 42 +++++++++++-
 .../pytest/pytestDiscoveryAdapter.ts          | 67 +++++--------------
 .../testController/pytest/pytestHelpers.ts    |  2 +-
 .../unittest/testDiscoveryAdapter.ts          | 64 +++++-------------
 .../unittest/unittestHelpers.ts               |  4 +-
 5 files changed, 76 insertions(+), 103 deletions(-)

diff --git a/src/client/testing/testController/common/discoveryHelpers.ts b/src/client/testing/testController/common/discoveryHelpers.ts
index b79bedd9a2b5..4eab97039d2e 100644
--- a/src/client/testing/testController/common/discoveryHelpers.ts
+++ b/src/client/testing/testController/common/discoveryHelpers.ts
@@ -1,16 +1,52 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-import { CancellationTokenSource, Uri } from 'vscode';
+import { CancellationToken, CancellationTokenSource, Disposable, Uri } from 'vscode';
 import { Deferred } from '../../../common/utils/async';
 import { traceError, traceInfo, traceVerbose } from '../../../logging';
-import { createDiscoveryErrorPayload, fixLogLinesNoTrailing } from './utils';
-import { ITestResultResolver } from './types';
+import { createDiscoveryErrorPayload, fixLogLinesNoTrailing, startDiscoveryNamedPipe } from './utils';
+import { DiscoveredTestPayload, ITestResultResolver } from './types';
 
 /**
  * Test provider type for logging purposes.
  */
 export type TestProvider = 'pytest' | 'unittest';
 
+/**
+ * Sets up the discovery named pipe and wires up cancellation.
+ * @param resultResolver The resolver to handle discovered test data
+ * @param token Optional cancellation token from the caller
+ * @param uri Workspace URI for logging
+ * @returns Object containing the pipe name, cancellation source, and disposable for the external token handler
+ */
+export async function setupDiscoveryPipe(
+    resultResolver: ITestResultResolver | undefined,
+    token: CancellationToken | undefined,
+    uri: Uri,
+): Promise<{ pipeName: string; cancellation: CancellationTokenSource; tokenDisposable: Disposable | undefined }> {
+    const discoveryPipeCancellation = new CancellationTokenSource();
+
+    // Wire up cancellation from external token and store the disposable
+    const tokenDisposable = token?.onCancellationRequested(() => {
+        traceInfo(`Test discovery cancelled.`);
+        discoveryPipeCancellation.cancel();
+    });
+
+    // Start the named pipe with the discovery listener
+    const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
+        if (!token?.isCancellationRequested) {
+            resultResolver?.resolveDiscovery(data);
+        }
+    }, discoveryPipeCancellation.token);
+
+    traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
+
+    return {
+        pipeName: discoveryPipeName,
+        cancellation: discoveryPipeCancellation,
+        tokenDisposable,
+    };
+}
+
 /**
  * Creates standard process event handlers for test discovery subprocess.
  * Handles stdout/stderr logging and error reporting on process exit.
diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
index 9756d2783b6b..d69f13c7ee20 100644
--- a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 import * as path from 'path';
-import { CancellationToken, CancellationTokenSource, Uri } from 'vscode';
+import { CancellationToken, Uri } from 'vscode';
 import { ChildProcess } from 'child_process';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
@@ -12,48 +12,13 @@ import { IConfigurationService } from '../../../common/types';
 import { Deferred } from '../../../common/utils/async';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
 import { traceError, traceInfo, traceVerbose } from '../../../logging';
-import { DiscoveredTestPayload, ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
-import { createTestingDeferred, startDiscoveryNamedPipe } from '../common/utils';
+import { ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
+import { createTestingDeferred } from '../common/utils';
 import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 import { PythonEnvironment } from '../../../pythonEnvironments/info';
 import { useEnvExtension, getEnvironment, runInBackground } from '../../../envExt/api.internal';
 import { buildPytestEnv as configureSubprocessEnv, handleSymlinkAndRootDir } from './pytestHelpers';
-import { cleanupOnCancellation, createProcessHandlers } from '../common/discoveryHelpers';
-
-/**
- * Sets up the discovery named pipe and wires up cancellation.
- * @param resultResolver The resolver to handle discovered test data
- * @param token Optional cancellation token from the caller
- * @param uri Workspace URI for logging
- * @returns Object containing the pipe name and cancellation source
- */
-async function setupDiscoveryPipe(
-    resultResolver: ITestResultResolver | undefined,
-    token: CancellationToken | undefined,
-    uri: Uri,
-): Promise<{ pipeName: string; cancellation: CancellationTokenSource }> {
-    const discoveryPipeCancellation = new CancellationTokenSource();
-
-    // Wire up cancellation from external token
-    token?.onCancellationRequested(() => {
-        traceInfo(`Test discovery cancelled.`);
-        discoveryPipeCancellation.cancel();
-    });
-
-    // Start the named pipe with the discovery listener
-    const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-        if (!token?.isCancellationRequested) {
-            resultResolver?.resolveDiscovery(data);
-        }
-    }, discoveryPipeCancellation.token);
-
-    traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
-
-    return {
-        pipeName: discoveryPipeName,
-        cancellation: discoveryPipeCancellation,
-    };
-}
+import { cleanupOnCancellation, createProcessHandlers, setupDiscoveryPipe } from '../common/discoveryHelpers';
 
 /**
  * Configures the subprocess environment for pytest discovery.
@@ -69,7 +34,7 @@ async function configureDiscoveryEnv(
 ): Promise<NodeJS.ProcessEnv> {
     const fullPluginPath = path.join(EXTENSION_ROOT_DIR, 'python_files');
     const envVars = await envVarsService?.getEnvironmentVariables(uri);
-    const mutableEnv = await configureSubprocessEnv(envVars, fullPluginPath, discoveryPipeName);
+    const mutableEnv = configureSubprocessEnv(envVars, fullPluginPath, discoveryPipeName);
     return mutableEnv;
 }
 
@@ -90,11 +55,11 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         interpreter?: PythonEnvironment,
     ): Promise<void> {
         // Setup discovery pipe and cancellation
-        const { pipeName: discoveryPipeName, cancellation: discoveryPipeCancellation } = await setupDiscoveryPipe(
-            this.resultResolver,
-            token,
-            uri,
-        );
+        const {
+            pipeName: discoveryPipeName,
+            cancellation: discoveryPipeCancellation,
+            tokenDisposable,
+        } = await setupDiscoveryPipe(this.resultResolver, token, uri);
 
         // Setup process handlers deferred (used by both execution paths)
         const deferredTillExecClose: Deferred<void> = createTestingDeferred();
@@ -137,7 +102,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 traceInfo(`Started pytest discovery subprocess (environment extension) for workspace ${uri.fsPath}`);
 
                 // Wire up cancellation and process events
-                token?.onCancellationRequested(() => {
+                const envExtCancellationHandler = token?.onCancellationRequested(() => {
                     cleanupOnCancellation('pytest', proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
                 proc.stdout.on('data', handlers.onStdout);
@@ -148,6 +113,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 });
 
                 await deferredTillExecClose.promise;
+                envExtCancellationHandler?.dispose();
                 traceInfo(`Pytest discovery completed for workspace ${uri.fsPath}`);
                 return;
             }
@@ -185,9 +151,10 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
+            let cancellationHandler: import('vscode').Disposable | undefined;
 
-            // Set up cancellation handler before execObservable to catch early cancellations
-            const cancellationHandler = token?.onCancellationRequested(() => {
+            // Set up cancellation handler after all early return checks
+            cancellationHandler = token?.onCancellationRequested(() => {
                 traceInfo(`Cancellation requested during pytest discovery for workspace ${uri.fsPath}`);
                 cleanupOnCancellation('pytest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
@@ -199,6 +166,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 if (!resultProc) {
                     traceError(`Failed to spawn pytest discovery subprocess for workspace ${uri.fsPath}`);
                     deferredTillExecClose.resolve();
+                    cancellationHandler?.dispose();
                     return;
                 }
                 traceInfo(`Started pytest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
@@ -215,7 +183,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
 
             cancellationHandler?.dispose();
 
-            // Check for early cancellation before awaitingre awaiting
+            // Check for early cancellation before awaiting
             if (token?.isCancellationRequested) {
                 traceInfo(`Pytest discovery was cancelled before process completion for workspace ${uri.fsPath}`);
                 deferredTillExecClose.resolve();
@@ -230,6 +198,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             throw error;
         } finally {
             traceVerbose(`Cleaning up pytest discovery resources for workspace ${uri.fsPath}`);
+            tokenDisposable?.dispose();
             discoveryPipeCancellation.dispose();
         }
     }
diff --git a/src/client/testing/testController/pytest/pytestHelpers.ts b/src/client/testing/testController/pytest/pytestHelpers.ts
index ca55cbf778ec..c6e748fb85a7 100644
--- a/src/client/testing/testController/pytest/pytestHelpers.ts
+++ b/src/client/testing/testController/pytest/pytestHelpers.ts
@@ -30,7 +30,7 @@ export async function handleSymlinkAndRootDir(cwd: string, pytestArgs: string[])
     }
     // if user has provided `--rootdir` then use that, otherwise add `cwd`
     // root dir is required so pytest can find the relative paths and for symlinks
-    addValueIfKeyNotExist(pytestArgs, '--rootdir', cwd);
+    pytestArgs = addValueIfKeyNotExist(pytestArgs, '--rootdir', cwd);
     return pytestArgs;
 }
 
diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
index 45b70cbd3d53..7d9a19ae6920 100644
--- a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
+++ b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -1,59 +1,23 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { CancellationToken, CancellationTokenSource, Uri } from 'vscode';
+import { CancellationToken, Uri } from 'vscode';
 import { ChildProcess } from 'child_process';
 import { IConfigurationService } from '../../../common/types';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
-import { DiscoveredTestPayload, ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
+import { ITestDiscoveryAdapter, ITestResultResolver } from '../common/types';
 import { IEnvironmentVariablesProvider } from '../../../common/variables/types';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
-import { startDiscoveryNamedPipe } from '../common/utils';
 import { traceError, traceInfo, traceVerbose } from '../../../logging';
 import { getEnvironment, runInBackground, useEnvExtension } from '../../../envExt/api.internal';
 import { PythonEnvironment } from '../../../pythonEnvironments/info';
 import { createTestingDeferred } from '../common/utils';
 import { buildDiscoveryCommand, buildUnittestEnv as configureSubprocessEnv } from './unittestHelpers';
-import { cleanupOnCancellation, createProcessHandlers } from '../common/discoveryHelpers';
-
-/**
- * Sets up the discovery named pipe and wires up cancellation.
- * @param resultResolver The resolver to handle discovered test data
- * @param token Optional cancellation token from the caller
- * @param uri Workspace URI for logging
- * @returns Object containing the pipe name and cancellation source
- */
-async function setupDiscoveryPipe(
-    resultResolver: ITestResultResolver | undefined,
-    token: CancellationToken | undefined,
-    uri: Uri,
-): Promise<{ pipeName: string; cancellation: CancellationTokenSource }> {
-    const discoveryPipeCancellation = new CancellationTokenSource();
-
-    // Wire up cancellation from external token
-    token?.onCancellationRequested(() => {
-        traceInfo(`Test discovery cancelled.`);
-        discoveryPipeCancellation.cancel();
-    });
-
-    // Start the named pipe with the discovery listener
-    const discoveryPipeName = await startDiscoveryNamedPipe((data: DiscoveredTestPayload) => {
-        if (!token?.isCancellationRequested) {
-            resultResolver?.resolveDiscovery(data);
-        }
-    }, discoveryPipeCancellation.token);
-
-    traceVerbose(`Created discovery pipe: ${discoveryPipeName} for workspace ${uri.fsPath}`);
-
-    return {
-        pipeName: discoveryPipeName,
-        cancellation: discoveryPipeCancellation,
-    };
-}
+import { cleanupOnCancellation, createProcessHandlers, setupDiscoveryPipe } from '../common/discoveryHelpers';
 
 /**
  * Configures the subprocess environment for unittest discovery.
@@ -68,7 +32,7 @@ async function configureDiscoveryEnv(
     discoveryPipeName: string,
 ): Promise<NodeJS.ProcessEnv> {
     const envVars = await envVarsService?.getEnvironmentVariables(uri);
-    const mutableEnv = await configureSubprocessEnv(envVars, discoveryPipeName);
+    const mutableEnv = configureSubprocessEnv(envVars, discoveryPipeName);
     return mutableEnv;
 }
 
@@ -89,11 +53,11 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         interpreter?: PythonEnvironment,
     ): Promise<void> {
         // Setup discovery pipe and cancellation
-        const { pipeName: discoveryPipeName, cancellation: discoveryPipeCancellation } = await setupDiscoveryPipe(
-            this.resultResolver,
-            token,
-            uri,
-        );
+        const {
+            pipeName: discoveryPipeName,
+            cancellation: discoveryPipeCancellation,
+            tokenDisposable,
+        } = await setupDiscoveryPipe(this.resultResolver, token, uri);
 
         // Setup process handlers deferred (used by both execution paths)
         const deferredTillExecClose = createTestingDeferred();
@@ -133,7 +97,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 traceInfo(`Started unittest discovery subprocess (environment extension) for workspace ${uri.fsPath}`);
 
                 // Wire up cancellation and process events
-                token?.onCancellationRequested(() => {
+                const envExtCancellationHandler = token?.onCancellationRequested(() => {
                     cleanupOnCancellation('unittest', proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
                 proc.stdout.on('data', handlers.onStdout);
@@ -144,6 +108,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 });
 
                 await deferredTillExecClose.promise;
+                envExtCancellationHandler?.dispose();
                 traceInfo(`Unittest discovery completed for workspace ${uri.fsPath}`);
                 return;
             }
@@ -181,9 +146,10 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
+            let cancellationHandler: import('vscode').Disposable | undefined;
 
-            // Set up cancellation handler before execObservable to catch early cancellations
-            const cancellationHandler = token?.onCancellationRequested(() => {
+            // Set up cancellation handler after all early return checks
+            cancellationHandler = token?.onCancellationRequested(() => {
                 traceInfo(`Cancellation requested during unittest discovery for workspace ${uri.fsPath}`);
                 cleanupOnCancellation('unittest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
@@ -195,6 +161,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 if (!resultProc) {
                     traceError(`Failed to spawn unittest discovery subprocess for workspace ${uri.fsPath}`);
                     deferredTillExecClose.resolve();
+                    cancellationHandler?.dispose();
                     return;
                 }
                 traceInfo(`Started unittest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
@@ -226,6 +193,7 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             throw error;
         } finally {
             traceVerbose(`Cleaning up unittest discovery resources for workspace ${uri.fsPath}`);
+            tokenDisposable?.dispose();
             discoveryPipeCancellation.dispose();
         }
     }
diff --git a/src/client/testing/testController/unittest/unittestHelpers.ts b/src/client/testing/testController/unittest/unittestHelpers.ts
index 6668c1df0e65..249a78dda7b7 100644
--- a/src/client/testing/testController/unittest/unittestHelpers.ts
+++ b/src/client/testing/testController/unittest/unittestHelpers.ts
@@ -7,10 +7,10 @@ import { traceInfo } from '../../../logging';
  * Builds the environment variables required for unittest discovery.
  * Sets TEST_RUN_PIPE for communication.
  */
-export async function buildUnittestEnv(
+export function buildUnittestEnv(
     envVars: { [key: string]: string | undefined } | undefined,
     discoveryPipeName: string,
-): Promise<{ [key: string]: string | undefined }> {
+): { [key: string]: string | undefined } {
     const mutableEnv = {
         ...envVars,
     };

From 4d4f57ddc0644066ea07f5f22a52edd9397450cc Mon Sep 17 00:00:00 2001
From: eleanorjboyd <26030610+eleanorjboyd@users.noreply.github.com>
Date: Fri, 21 Nov 2025 09:28:44 -0800
Subject: [PATCH 8/8] track all disposables

---
 .../testController/common/discoveryHelpers.ts | 13 ++------
 .../pytest/pytestDiscoveryAdapter.ts          | 33 ++++++++++---------
 .../unittest/testDiscoveryAdapter.ts          | 31 ++++++++---------
 3 files changed, 36 insertions(+), 41 deletions(-)

diff --git a/src/client/testing/testController/common/discoveryHelpers.ts b/src/client/testing/testController/common/discoveryHelpers.ts
index 4eab97039d2e..e170ad576ae8 100644
--- a/src/client/testing/testController/common/discoveryHelpers.ts
+++ b/src/client/testing/testController/common/discoveryHelpers.ts
@@ -84,17 +84,10 @@ export function createProcessHandlers(
             const out = fixLogLinesNoTrailing(data.toString());
             traceError(out);
         },
-        onExit: (code: number | null, signal: NodeJS.Signals | null) => {
+        onExit: (code: number | null, _signal: NodeJS.Signals | null) => {
             // The 'exit' event fires when the process terminates, but streams may still be open.
-            if (!isSuccessCode(code)) {
-                const exitCodeNote =
-                    allowedSuccessCodes.length > 0
-                        ? ` Note: Exit codes ${allowedSuccessCodes.join(', ')} are also treated as success.`
-                        : '';
-                traceError(
-                    `${testProvider} discovery subprocess exited with code ${code} and signal ${signal} for workspace ${uri.fsPath}.${exitCodeNote}`,
-                );
-            } else if (code === 0) {
+            // Only log verbose success message here; error handling happens in onClose.
+            if (isSuccessCode(code)) {
                 traceVerbose(`${testProvider} discovery subprocess exited successfully for workspace ${uri.fsPath}`);
             }
         },
diff --git a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
index d69f13c7ee20..7ad69c71fa0e 100644
--- a/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
+++ b/src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 import * as path from 'path';
-import { CancellationToken, Uri } from 'vscode';
+import { CancellationToken, Disposable, Uri } from 'vscode';
 import { ChildProcess } from 'child_process';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
@@ -64,6 +64,12 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         // Setup process handlers deferred (used by both execution paths)
         const deferredTillExecClose: Deferred<void> = createTestingDeferred();
 
+        // Collect all disposables related to discovery to handle cleanup in finally block
+        const disposables: Disposable[] = [];
+        if (tokenDisposable) {
+            disposables.push(tokenDisposable);
+        }
+
         try {
             // Build pytest command and arguments
             const settings = this.configSettings.getSettings(uri);
@@ -105,6 +111,9 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 const envExtCancellationHandler = token?.onCancellationRequested(() => {
                     cleanupOnCancellation('pytest', proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
+                if (envExtCancellationHandler) {
+                    disposables.push(envExtCancellationHandler);
+                }
                 proc.stdout.on('data', handlers.onStdout);
                 proc.stderr.on('data', handlers.onStderr);
                 proc.onExit((code, signal) => {
@@ -113,7 +122,6 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 });
 
                 await deferredTillExecClose.promise;
-                envExtCancellationHandler?.dispose();
                 traceInfo(`Pytest discovery completed for workspace ${uri.fsPath}`);
                 return;
             }
@@ -151,13 +159,15 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
-            let cancellationHandler: import('vscode').Disposable | undefined;
 
             // Set up cancellation handler after all early return checks
-            cancellationHandler = token?.onCancellationRequested(() => {
+            const cancellationHandler = token?.onCancellationRequested(() => {
                 traceInfo(`Cancellation requested during pytest discovery for workspace ${uri.fsPath}`);
                 cleanupOnCancellation('pytest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
+            if (cancellationHandler) {
+                disposables.push(cancellationHandler);
+            }
 
             try {
                 const result = execService.execObservable(commandArgs, spawnOptions);
@@ -166,14 +176,12 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 if (!resultProc) {
                     traceError(`Failed to spawn pytest discovery subprocess for workspace ${uri.fsPath}`);
                     deferredTillExecClose.resolve();
-                    cancellationHandler?.dispose();
                     return;
                 }
                 traceInfo(`Started pytest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
             } catch (error) {
                 traceError(`Error spawning pytest discovery subprocess for workspace ${uri.fsPath}: ${error}`);
                 deferredTillExecClose.resolve();
-                cancellationHandler?.dispose();
                 throw error;
             }
             resultProc.stdout?.on('data', handlers.onStdout);
@@ -181,14 +189,6 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             resultProc.on('exit', handlers.onExit);
             resultProc.on('close', handlers.onClose);
 
-            cancellationHandler?.dispose();
-
-            // Check for early cancellation before awaiting
-            if (token?.isCancellationRequested) {
-                traceInfo(`Pytest discovery was cancelled before process completion for workspace ${uri.fsPath}`);
-                deferredTillExecClose.resolve();
-            }
-
             traceVerbose(`Waiting for pytest discovery subprocess to complete for workspace ${uri.fsPath}`);
             await deferredTillExecClose.promise;
             traceInfo(`Pytest discovery completed for workspace ${uri.fsPath}`);
@@ -197,8 +197,9 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             deferredTillExecClose.resolve();
             throw error;
         } finally {
-            traceVerbose(`Cleaning up pytest discovery resources for workspace ${uri.fsPath}`);
-            tokenDisposable?.dispose();
+            // Dispose all cancellation handlers and event subscriptions
+            disposables.forEach((d) => d.dispose());
+            // Dispose the discovery pipe cancellation token
             discoveryPipeCancellation.dispose();
         }
     }
diff --git a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
index 7d9a19ae6920..7c986e95a449 100644
--- a/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
+++ b/src/client/testing/testController/unittest/testDiscoveryAdapter.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { CancellationToken, Uri } from 'vscode';
+import { CancellationToken, Disposable, Uri } from 'vscode';
 import { ChildProcess } from 'child_process';
 import { IConfigurationService } from '../../../common/types';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
@@ -62,6 +62,11 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         // Setup process handlers deferred (used by both execution paths)
         const deferredTillExecClose = createTestingDeferred();
 
+        // Collect all disposables for cleanup in finally block
+        const disposables: Disposable[] = [];
+        if (tokenDisposable) {
+            disposables.push(tokenDisposable);
+        }
         try {
             // Build unittest command and arguments
             const settings = this.configSettings.getSettings(uri);
@@ -100,6 +105,9 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 const envExtCancellationHandler = token?.onCancellationRequested(() => {
                     cleanupOnCancellation('unittest', proc, deferredTillExecClose, discoveryPipeCancellation, uri);
                 });
+                if (envExtCancellationHandler) {
+                    disposables.push(envExtCancellationHandler);
+                }
                 proc.stdout.on('data', handlers.onStdout);
                 proc.stderr.on('data', handlers.onStderr);
                 proc.onExit((code, signal) => {
@@ -108,7 +116,6 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 });
 
                 await deferredTillExecClose.promise;
-                envExtCancellationHandler?.dispose();
                 traceInfo(`Unittest discovery completed for workspace ${uri.fsPath}`);
                 return;
             }
@@ -146,13 +153,15 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             };
 
             let resultProc: ChildProcess | undefined;
-            let cancellationHandler: import('vscode').Disposable | undefined;
 
             // Set up cancellation handler after all early return checks
-            cancellationHandler = token?.onCancellationRequested(() => {
+            const cancellationHandler = token?.onCancellationRequested(() => {
                 traceInfo(`Cancellation requested during unittest discovery for workspace ${uri.fsPath}`);
                 cleanupOnCancellation('unittest', resultProc, deferredTillExecClose, discoveryPipeCancellation, uri);
             });
+            if (cancellationHandler) {
+                disposables.push(cancellationHandler);
+            }
 
             try {
                 const result = execService.execObservable(execArgs, spawnOptions);
@@ -161,14 +170,12 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
                 if (!resultProc) {
                     traceError(`Failed to spawn unittest discovery subprocess for workspace ${uri.fsPath}`);
                     deferredTillExecClose.resolve();
-                    cancellationHandler?.dispose();
                     return;
                 }
                 traceInfo(`Started unittest discovery subprocess (execution factory) for workspace ${uri.fsPath}`);
             } catch (error) {
                 traceError(`Error spawning unittest discovery subprocess for workspace ${uri.fsPath}: ${error}`);
                 deferredTillExecClose.resolve();
-                cancellationHandler?.dispose();
                 throw error;
             }
             resultProc.stdout?.on('data', handlers.onStdout);
@@ -176,14 +183,6 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             resultProc.on('exit', handlers.onExit);
             resultProc.on('close', handlers.onClose);
 
-            cancellationHandler?.dispose();
-
-            // Check for early cancellation before awaiting
-            if (token?.isCancellationRequested) {
-                traceInfo(`Unittest discovery was cancelled before process completion for workspace ${uri.fsPath}`);
-                deferredTillExecClose.resolve();
-            }
-
             traceVerbose(`Waiting for unittest discovery subprocess to complete for workspace ${uri.fsPath}`);
             await deferredTillExecClose.promise;
             traceInfo(`Unittest discovery completed for workspace ${uri.fsPath}`);
@@ -193,7 +192,9 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
             throw error;
         } finally {
             traceVerbose(`Cleaning up unittest discovery resources for workspace ${uri.fsPath}`);
-            tokenDisposable?.dispose();
+            // Dispose all cancellation handlers and event subscriptions
+            disposables.forEach((d) => d.dispose());
+            // Dispose the discovery pipe cancellation token
             discoveryPipeCancellation.dispose();
         }
     }