Refactor pytest and unittest test discovery

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

@@ -0,0 +1,108 @@
+// 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) {
+        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();
+}