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

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);
+        });
     });
 });