From 01c7665e37f4674a6a574d38f4f7af9344ec0485 Mon Sep 17 00:00:00 2001
From: Eleanor Boyd <eleanorboyd@microsoft.com>
Date: Tue, 17 Oct 2023 16:45:44 -0700
Subject: [PATCH] add correct retrieval of workspace adapter for test discovery
 in multiroot context (#22246)

fixes: https://github.com/microsoft/vscode-python/issues/22218
---
 .../testing/testController/controller.ts      | 43 +++++++++++++------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/src/client/testing/testController/controller.ts b/src/client/testing/testController/controller.ts
index a87017a26a51..329326d84af9 100644
--- a/src/client/testing/testController/controller.ts
+++ b/src/client/testing/testController/controller.ts
@@ -269,13 +269,20 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
             if (settings.testing.pytestEnabled) {
                 if (pythonTestAdapterRewriteEnabled(this.serviceContainer)) {
                     traceInfo(`Running discovery for pytest using the new test adapter.`);
-                    const testAdapter =
-                        this.testAdapters.get(uri) || (this.testAdapters.values().next().value as WorkspaceTestAdapter);
-                    testAdapter.discoverTests(
-                        this.testController,
-                        this.refreshCancellation.token,
-                        this.pythonExecFactory,
-                    );
+                    if (workspace && workspace.uri) {
+                        const testAdapter = this.testAdapters.get(workspace.uri);
+                        if (testAdapter) {
+                            testAdapter.discoverTests(
+                                this.testController,
+                                this.refreshCancellation.token,
+                                this.pythonExecFactory,
+                            );
+                        } else {
+                            traceError('Unable to find test adapter for workspace.');
+                        }
+                    } else {
+                        traceError('Unable to find workspace for given file');
+                    }
                 } else {
                     // else use OLD test discovery mechanism
                     await this.pytest.refreshTestData(this.testController, uri, this.refreshCancellation.token);
@@ -283,13 +290,21 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
             } else if (settings.testing.unittestEnabled) {
                 if (pythonTestAdapterRewriteEnabled(this.serviceContainer)) {
                     traceInfo(`Running discovery for unittest using the new test adapter.`);
-                    const testAdapter =
-                        this.testAdapters.get(uri) || (this.testAdapters.values().next().value as WorkspaceTestAdapter);
-                    testAdapter.discoverTests(
-                        this.testController,
-                        this.refreshCancellation.token,
-                        this.pythonExecFactory,
-                    );
+                    traceInfo(`Running discovery for pytest using the new test adapter.`);
+                    if (workspace && workspace.uri) {
+                        const testAdapter = this.testAdapters.get(workspace.uri);
+                        if (testAdapter) {
+                            testAdapter.discoverTests(
+                                this.testController,
+                                this.refreshCancellation.token,
+                                this.pythonExecFactory,
+                            );
+                        } else {
+                            traceError('Unable to find test adapter for workspace.');
+                        }
+                    } else {
+                        traceError('Unable to find workspace for given file');
+                    }
                 } else {
                     // else use OLD test discovery mechanism
                     await this.unittest.refreshTestData(this.testController, uri, this.refreshCancellation.token);