Fix pytest discovery using the workspace-root interpreter after missed startup environment events

Subscribe to Python Environments extension project/environment changes
before the initial test project discovery, and queue re-discovery for
workspaces whose initial project registration is still in flight when an
environment is assigned. Previously, environment assignments raised
while initial discovery was running were missed entirely, leaving the
workspace stuck on the fallback default project that discovers tests
with the workspace-root (e.g. system) interpreter instead of the
project's environment.

Fixes #25718

Author: wasim-builds

src/client/testing/testController/controller.ts

@@ -230,6 +230,16 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
         if (useEnvExtension()) {
             traceInfo('[test-by-project] Activating project-based testing mode');
 
+            // Subscribe to project and environment changes BEFORE the initial project discovery.
+            // Acquiring the environments API during discovery activates the environments extension
+            // and kicks off its initial refresh; project registrations and environment assignments
+            // raised before these subscriptions are attached would be missed entirely. A missed
+            // environment assignment leaves the workspace stuck on the fallback default project,
+            // which discovers tests with the workspace-root interpreter instead of the
+            // project-specific environment (e.g. a monorepo sub-project venv).
+            await this.subscribeToProjectChanges();
+            await this.subscribeToEnvironmentChanges();
+
             // Discover projects in parallel across all workspaces
             // Promise.allSettled ensures one workspace failure doesn't block others
             const results = await Promise.allSettled(
@@ -254,11 +264,6 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
                     this.activateLegacyWorkspace(workspace);
                 }
             });
-            // Subscribe to project changes to update test tree when projects are added/removed
-            await this.subscribeToProjectChanges();
-            // Subscribe to environment changes so projects that had no resolved environment at
-            // activation get discovered once the environments extension assigns one.
-            await this.subscribeToEnvironmentChanges();
             return;
         }
 
@@ -340,8 +345,13 @@ export class PythonTestController implements ITestController, IExtensionSingleAc
 
         let queued = false;
         for (const workspace of affected) {
-            // Only workspaces already in project-based mode can be re-discovered this way.
-            if (this.projectRegistry.hasProjects(workspace.uri)) {
+            // Workspaces in legacy mode (a legacy adapter was registered for them) are not
+            // re-discovered this way. All other workspaces are: both workspaces with registered
+            // projects and workspaces whose initial project registration is still in flight.
+            // The latter ensures environment assignments raised during startup discovery (before
+            // the workspace had any projects registered) still trigger re-discovery once the
+            // environments extension resolves them.
+            if (this.projectRegistry.hasProjects(workspace.uri) || !this.testAdapters.has(workspace.uri)) {
                 this.pendingEnvChangeWorkspaces.set(workspace.uri.toString(), workspace);
                 queued = true;
             }

src/test/testing/testController/controller.unit.test.ts

@@ -409,8 +409,11 @@ suite('PythonTestController', () => {
             assert.strictEqual(triggerStub.notCalled, true);
         });
 
-        test('ignores workspaces that are not in project-based mode', () => {
+        test('ignores workspaces in legacy mode (legacy adapter registered)', () => {
             const { controller, triggerStub } = setupController(false);
+            // A registered legacy adapter marks the workspace as legacy mode, so environment
+            // changes must not queue a project-based re-discovery for it.
+            (controller as any).testAdapters.set(workspaceUri, {});
 
             (controller as any).handleEnvironmentChange({ uri: workspaceUri, old: undefined, new: newEnv });
 
@@ -419,6 +422,21 @@ suite('PythonTestController', () => {
             assert.strictEqual(triggerStub.notCalled, true);
         });
 
+        test('queues workspaces whose initial project registration is still in flight', () => {
+            // Startup race: the environments extension assigns an environment while the initial
+            // project discovery for the workspace is still running. No projects are registered yet
+            // (hasProjects === false) and no legacy adapter exists. The assignment must still queue
+            // a re-discovery, otherwise the workspace stays on the fallback default project (using
+            // the workspace-root interpreter) until a manual refresh.
+            const { controller, triggerStub } = setupController(false);
+
+            (controller as any).handleEnvironmentChange({ uri: workspaceUri, old: undefined, new: newEnv });
+
+            const pending = (controller as any).pendingEnvChangeWorkspaces as Map<string, vscode.WorkspaceFolder>;
+            assert.strictEqual(pending.has(workspaceUri.toString()), true);
+            assert.strictEqual(triggerStub.calledOnce, true);
+        });
+
         test('re-discovers each pending workspace and clears the queue', async () => {
             const { controller, rediscoverStub } = setupController(true);
 
@@ -430,5 +448,35 @@ suite('PythonTestController', () => {
             assert.strictEqual(rediscoverStub.calledOnceWithExactly(workspaceUri), true);
             assert.strictEqual(pending.size, 0);
         });
+
+        test('activate() subscribes to project and environment changes before initial project discovery', async () => {
+            // Subscribing after the initial discovery misses project registrations and environment
+            // assignments raised by the environments extension while discovery is running (its
+            // initial refresh is kicked off by acquiring its API during discovery).
+            sandbox.stub(envExtApiInternal, 'useEnvExtension').returns(true);
+            sandbox.stub(envExtApiInternal, 'getEnvExtApi').resolves({
+                getPythonProjects: () => [],
+                getEnvironment: sandbox.stub().resolves(undefined),
+                onDidChangePythonProjects: sandbox.stub().returns({ dispose: () => undefined }),
+                onDidChangeEnvironment: sandbox.stub().returns({ dispose: () => undefined }),
+            } as any);
+
+            const controller = createController({
+                workspaceService: ({
+                    workspaceFolders: [workspaceFolder],
+                    getWorkspaceFolder: () => workspaceFolder,
+                } as unknown) as any,
+            });
+
+            const subscribeProjectsSpy = sandbox.spy(controller as any, 'subscribeToProjectChanges');
+            const subscribeEnvSpy = sandbox.spy(controller as any, 'subscribeToEnvironmentChanges');
+            const discoverSpy = sandbox.spy((controller as any).projectRegistry, 'discoverAndRegisterProjects');
+
+            await controller.activate();
+
+            assert.strictEqual(discoverSpy.calledOnceWithExactly(workspaceUri), true);
+            assert.strictEqual(subscribeProjectsSpy.calledBefore(discoverSpy), true);
+            assert.strictEqual(subscribeEnvSpy.calledBefore(discoverSpy), true);
+        });
     });
 });