Enhance pytest installation flow and error handling with Environment Extension integration

src/client/testing/configuration/pytest/testConfigurationManager.ts

@@ -3,12 +3,19 @@ import { QuickPickItem, Uri } from 'vscode';
 import { IFileSystem } from '../../../common/platform/types';
 import { Product } from '../../../common/types';
 import { IServiceContainer } from '../../../ioc/types';
+import { IApplicationShell } from '../../../common/application/types';
 import { TestConfigurationManager } from '../../common/testConfigurationManager';
 import { ITestConfigSettingsService } from '../../common/types';
+import { PytestInstallationHelper } from '../pytestInstallationHelper';
+import { traceInfo } from '../../../logging';
 
 export class ConfigurationManager extends TestConfigurationManager {
+    private readonly pytestInstallationHelper: PytestInstallationHelper;
+
     constructor(workspace: Uri, serviceContainer: IServiceContainer, cfg?: ITestConfigSettingsService) {
         super(workspace, Product.pytest, serviceContainer, cfg);
+        const appShell = serviceContainer.get<IApplicationShell>(IApplicationShell);
+        this.pytestInstallationHelper = new PytestInstallationHelper(appShell);
     }
 
     public async requiresUserToConfigure(wkspace: Uri): Promise<boolean> {
@@ -42,10 +49,22 @@ export class ConfigurationManager extends TestConfigurationManager {
             args.push(testDir);
         }
         const installed = await this.installer.isInstalled(Product.pytest);
+        await this.testConfigSettingsService.updateTestArgs(wkspace.fsPath, Product.pytest, args);
         if (!installed) {
-            await this.installer.install(Product.pytest);
+            // Check if Python Environments extension is available for enhanced installation flow
+            if (this.pytestInstallationHelper.isEnvExtensionAvailable()) {
+                traceInfo('pytest not installed, prompting user with environment extension integration');
+                const installAttempted = await this.pytestInstallationHelper.promptToInstallPytest(wkspace);
+                if (!installAttempted) {
+                    // User chose to ignore or installation failed
+                    return;
+                }
+            } else {
+                // Fall back to traditional installer
+                traceInfo('pytest not installed, falling back to traditional installer');
+                await this.installer.install(Product.pytest);
+            }
         }
-        await this.testConfigSettingsService.updateTestArgs(wkspace.fsPath, Product.pytest, args);
     }
 
     private async getConfigFiles(rootDir: string): Promise<string[]> {

src/client/testing/configuration/pytestInstallationHelper.ts

@@ -0,0 +1,95 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { Uri, l10n } from 'vscode';
+import { IApplicationShell } from '../../common/application/types';
+import { traceInfo, traceError } from '../../logging';
+import { useEnvExtension, getEnvExtApi } from '../../envExt/api.internal';
+import { getEnvironment } from '../../envExt/api.internal';
+
+/**
+ * Helper class to handle pytest installation using the appropriate method
+ * based on whether the Python Environments extension is available.
+ */
+export class PytestInstallationHelper {
+    constructor(private readonly appShell: IApplicationShell) {}
+
+    /**
+     * Prompts the user to install pytest with appropriate installation method.
+     * @param workspaceUri The workspace URI where pytest should be installed
+     * @returns Promise that resolves to true if installation was attempted, false otherwise
+     */
+    async promptToInstallPytest(workspaceUri: Uri): Promise<boolean> {
+        const message = l10n.t('pytest selected but not installed. Would you like to install pytest?');
+        const installOption = l10n.t('Install pytest');
+
+        const selection = await this.appShell.showInformationMessage(message, { modal: true }, installOption);
+
+        if (selection === installOption) {
+            return this.installPytest(workspaceUri);
+        }
+
+        return false;
+    }
+
+    /**
+     * Installs pytest using the appropriate method based on available extensions.
+     * @param workspaceUri The workspace URI where pytest should be installed
+     * @returns Promise that resolves to true if installation was successful, false otherwise
+     */
+    private async installPytest(workspaceUri: Uri): Promise<boolean> {
+        try {
+            if (useEnvExtension()) {
+                return this.installPytestWithEnvExtension(workspaceUri);
+            } else {
+                // Fall back to traditional installer if environments extension is not available
+                traceInfo(
+                    'Python Environments extension not available, installation cannot proceed via environment extension',
+                );
+                return false;
+            }
+        } catch (error) {
+            traceError('Error installing pytest:', error);
+            return false;
+        }
+    }
+
+    /**
+     * Installs pytest using the Python Environments extension.
+     * @param workspaceUri The workspace URI where pytest should be installed
+     * @returns Promise that resolves to true if installation was successful, false otherwise
+     */
+    private async installPytestWithEnvExtension(workspaceUri: Uri): Promise<boolean> {
+        try {
+            const envExtApi = await getEnvExtApi();
+            const environment = await getEnvironment(workspaceUri);
+
+            if (!environment) {
+                traceError('No Python environment found for workspace:', workspaceUri.fsPath);
+                await this.appShell.showErrorMessage(
+                    l10n.t('No Python environment found. Please set up a Python environment first.'),
+                );
+                return false;
+            }
+
+            traceInfo('Installing pytest using Python Environments extension...');
+            await envExtApi.managePackages(environment, {
+                install: ['pytest'],
+            });
+
+            traceInfo('pytest installation completed successfully');
+            return true;
+        } catch (error) {
+            traceError('Failed to install pytest using Python Environments extension:', error);
+            return false;
+        }
+    }
+
+    /**
+     * Checks if the Python Environments extension is available for package management.
+     * @returns True if the extension is available, false otherwise
+     */
+    isEnvExtensionAvailable(): boolean {
+        return useEnvExtension();
+    }
+}

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

@@ -174,12 +174,34 @@ export async function startDiscoveryNamedPipe(
     return pipeName;
 }
 
+/**
+ * Detects if an error message indicates that pytest is not installed.
+ * @param message The error message to check
+ * @returns True if the error indicates pytest is not installed
+ */
+function isPytestNotInstalledError(message: string): boolean {
+    return (
+        (message.includes('ModuleNotFoundError') && message.includes('pytest')) ||
+        (message.includes('No module named') && message.includes('pytest')) ||
+        (message.includes('ImportError') && message.includes('pytest'))
+    );
+}
+
 export function buildErrorNodeOptions(uri: Uri, message: string, testType: string): ErrorTestItemOptions {
-    const labelText = testType === 'pytest' ? 'pytest Discovery Error' : 'Unittest Discovery Error';
+    let labelText = testType === 'pytest' ? 'pytest Discovery Error' : 'Unittest Discovery Error';
+    let errorMessage = message;
+
+    // Provide more specific error message if pytest is not installed
+    if (testType === 'pytest' && isPytestNotInstalledError(message)) {
+        labelText = 'pytest Not Installed';
+        errorMessage =
+            'pytest is not installed in the selected Python environment. Please install pytest to enable test discovery and execution.';
+    }
+
     return {
         id: `DiscoveryError:${uri.fsPath}`,
         label: `${labelText} [${path.basename(uri.fsPath)}]`,
-        error: message,
+        error: errorMessage,
     };
 }
 

src/test/testing/configuration/pytestInstallationHelper.unit.test.ts

@@ -0,0 +1,131 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { expect } from 'chai';
+import * as sinon from 'sinon';
+import { Uri } from 'vscode';
+import * as TypeMoq from 'typemoq';
+import { IApplicationShell } from '../../../client/common/application/types';
+import { PytestInstallationHelper } from '../../../client/testing/configuration/pytestInstallationHelper';
+import * as envExtApi from '../../../client/envExt/api.internal';
+
+suite('PytestInstallationHelper', () => {
+    let appShell: TypeMoq.IMock<IApplicationShell>;
+    let helper: PytestInstallationHelper;
+    let useEnvExtensionStub: sinon.SinonStub;
+    let getEnvExtApiStub: sinon.SinonStub;
+    let getEnvironmentStub: sinon.SinonStub;
+
+    const workspaceUri = Uri.file('/test/workspace');
+
+    setup(() => {
+        appShell = TypeMoq.Mock.ofType<IApplicationShell>();
+        helper = new PytestInstallationHelper(appShell.object);
+
+        useEnvExtensionStub = sinon.stub(envExtApi, 'useEnvExtension');
+        getEnvExtApiStub = sinon.stub(envExtApi, 'getEnvExtApi');
+        getEnvironmentStub = sinon.stub(envExtApi, 'getEnvironment');
+    });
+
+    teardown(() => {
+        sinon.restore();
+    });
+
+    test('promptToInstallPytest should return false if user selects ignore', async () => {
+        appShell
+            .setup((a) =>
+                a.showInformationMessage(
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                ),
+            )
+            .returns(() => Promise.resolve('Ignore'))
+            .verifiable(TypeMoq.Times.once());
+
+        const result = await helper.promptToInstallPytest(workspaceUri);
+
+        expect(result).to.be.false;
+        appShell.verifyAll();
+    });
+
+    test('promptToInstallPytest should return false if user cancels', async () => {
+        appShell
+            .setup((a) =>
+                a.showInformationMessage(
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                ),
+            )
+            .returns(() => Promise.resolve(undefined))
+            .verifiable(TypeMoq.Times.once());
+
+        const result = await helper.promptToInstallPytest(workspaceUri);
+
+        expect(result).to.be.false;
+        appShell.verifyAll();
+    });
+
+    test('isEnvExtensionAvailable should return result from useEnvExtension', () => {
+        useEnvExtensionStub.returns(true);
+
+        const result = helper.isEnvExtensionAvailable();
+
+        expect(result).to.be.true;
+        expect(useEnvExtensionStub.calledOnce).to.be.true;
+    });
+
+    test('promptToInstallPytest should return false if env extension not available', async () => {
+        useEnvExtensionStub.returns(false);
+
+        appShell
+            .setup((a) =>
+                a.showInformationMessage(
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                ),
+            )
+            .returns(() => Promise.resolve('Install pytest'))
+            .verifiable(TypeMoq.Times.once());
+
+        const result = await helper.promptToInstallPytest(workspaceUri);
+
+        expect(result).to.be.false;
+        appShell.verifyAll();
+    });
+
+    test('promptToInstallPytest should attempt installation when env extension is available', async () => {
+        useEnvExtensionStub.returns(true);
+
+        const mockEnvironment = { envId: { id: 'test-env', managerId: 'test-manager' } };
+        const mockEnvExtApi = {
+            managePackages: sinon.stub().resolves(),
+        };
+
+        getEnvExtApiStub.resolves(mockEnvExtApi);
+        getEnvironmentStub.resolves(mockEnvironment);
+
+        appShell
+            .setup((a) =>
+                a.showInformationMessage(
+                    TypeMoq.It.is((msg: string) => msg.includes('pytest selected but not installed')),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                    TypeMoq.It.isAny(),
+                ),
+            )
+            .returns(() => Promise.resolve('Install pytest'))
+            .verifiable(TypeMoq.Times.once());
+
+        const result = await helper.promptToInstallPytest(workspaceUri);
+
+        expect(result).to.be.true;
+        expect(mockEnvExtApi.managePackages.calledOnceWithExactly(mockEnvironment, { install: ['pytest'] })).to.be.true;
+        appShell.verifyAll();
+    });
+});

src/test/testing/testController/common/buildErrorNodeOptions.unit.test.ts

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { expect } from 'chai';
+import { Uri } from 'vscode';
+import { buildErrorNodeOptions } from '../../../../client/testing/testController/common/utils';
+
+suite('buildErrorNodeOptions - pytest not installed detection', () => {
+    const workspaceUri = Uri.file('/test/workspace');
+
+    test('Should detect pytest ModuleNotFoundError and provide specific message', () => {
+        const errorMessage =
+            'Traceback (most recent call last):\n  File "<string>", line 1, in <module>\n    import pytest\nModuleNotFoundError: No module named \'pytest\'';
+
+        const result = buildErrorNodeOptions(workspaceUri, errorMessage, 'pytest');
+
+        expect(result.label).to.equal('pytest Not Installed [workspace]');
+        expect(result.error).to.equal(
+            'pytest is not installed in the selected Python environment. Please install pytest to enable test discovery and execution.',
+        );
+    });
+
+    test('Should detect pytest ImportError and provide specific message', () => {
+        const errorMessage = 'ImportError: No module named pytest';
+
+        const result = buildErrorNodeOptions(workspaceUri, errorMessage, 'pytest');
+
+        expect(result.label).to.equal('pytest Not Installed [workspace]');
+        expect(result.error).to.equal(
+            'pytest is not installed in the selected Python environment. Please install pytest to enable test discovery and execution.',
+        );
+    });
+
+    test('Should use generic error for non-pytest-related errors', () => {
+        const errorMessage = 'Some other error occurred';
+
+        const result = buildErrorNodeOptions(workspaceUri, errorMessage, 'pytest');
+
+        expect(result.label).to.equal('pytest Discovery Error [workspace]');
+        expect(result.error).to.equal('Some other error occurred');
+    });
+
+    test('Should use generic error for unittest errors', () => {
+        const errorMessage = "ModuleNotFoundError: No module named 'pytest'";
+
+        const result = buildErrorNodeOptions(workspaceUri, errorMessage, 'unittest');
+
+        expect(result.label).to.equal('Unittest Discovery Error [workspace]');
+        expect(result.error).to.equal("ModuleNotFoundError: No module named 'pytest'");
+    });
+});