Don Jayamanne (DonJayamanne) · Jan 7, 2019 · Nov 11, 2018 · Nov 11, 2018 · Nov 11, 2018 · Nov 11, 2018
diff --git a/.github/test_plan.md b/.github/test_plan.md
@@ -288,6 +288,10 @@ def test_failure():
 - [ ] Tests are discovered (as shown by code lenses on each test)
   - [ ] `Run Test` works
   - [ ] `Debug Test` works
+- [ ] A `Diagnostic` is shown in the problems pane for each failed/skipped test
+  - [ ] The `Diagnostic`s are organized according to the file the test was executed from (not neccesarily the file it was defined in)
+  - [ ] The appropriate `DiagnosticRelatedInformation` is shown for each `Diagnostic`
+  - [ ] The `DiagnosticRelatedInformation` reflects the traceback for the test
 
 #### [`nose`](https://code.visualstudio.com/docs/python/unit-testing#_nose-configuration-settings)
 ```python

diff --git a/news/1 Enhancements/120.md b/news/1 Enhancements/120.md
@@ -0,0 +1,2 @@
+Create diagnostics for failed/skipped tests that were run with pytest.
+(thanks [Chris NeJame](https://github.com/SalmonMode/))
diff --git a/package.nls.json b/package.nls.json
diff --git a/src/client/common/utils/localize.ts b/src/client/common/utils/localize.ts
@@ -144,6 +144,12 @@ export namespace DebugConfigurationPrompts {
     export const attachRemoteHostValidationError = localize('debug.attachRemoteHostValidationError', 'Enter a valid host name or IP address');
 }
 
+export namespace UnitTests {
+    export const testErrorDiagnosticMessage = localize('UnitTests.testErrorDiagnosticMessage', 'Error');
+    export const testFailDiagnosticMessage = localize('UnitTests.testFailDiagnosticMessage', 'Fail');
+    export const testSkippedDiagnosticMessage = localize('UnitTests.testSkippedDiagnosticMessage', 'Skipped');
+}
+
 // Skip using vscode-nls and instead just compute our strings based on key values. Key values
 // can be loaded out of the nls.<locale>.json files
 let loadedCollection: { [index: string]: string } | undefined;

diff --git a/src/client/unittests/common/managers/baseTestManager.ts b/src/client/unittests/common/managers/baseTestManager.ts
@@ -1,12 +1,14 @@
-import { CancellationToken, CancellationTokenSource, Disposable, OutputChannel, Uri } from 'vscode';
+import { CancellationToken, CancellationTokenSource, Diagnostic, DiagnosticCollection, DiagnosticRelatedInformation, Disposable, languages, OutputChannel, Uri } from 'vscode';
 import { IWorkspaceService } from '../../../common/application/types';
 import { isNotInstalledError } from '../../../common/helpers';
+import { IFileSystem } from '../../../common/platform/types';
 import { IConfigurationService, IDisposableRegistry, IInstaller, IOutputChannel, IPythonSettings, Product } from '../../../common/types';
 import { getNamesAndValues } from '../../../common/utils/enum';
 import { IServiceContainer } from '../../../ioc/types';
 import { UNITTEST_DISCOVER, UNITTEST_RUN } from '../../../telemetry/constants';
 import { sendTelemetryEvent } from '../../../telemetry/index';
 import { TestDiscoverytTelemetry, TestRunTelemetry } from '../../../telemetry/types';
+import { IPythonUnitTestMessage, IUnitTestDiagnosticService } from '../../types';
 import { CANCELLATION_REASON, CommandSource, TEST_OUTPUT_CHANNEL } from './../constants';
 import { ITestCollectionStorageService, ITestDiscoveryService, ITestManager, ITestResultsService, ITestsHelper, TestDiscoveryOptions, TestProvider, Tests, TestStatus, TestsToRun } from './../types';
 
@@ -16,7 +18,9 @@ enum CancellationTokenType {
 }
 
 export abstract class BaseTestManager implements ITestManager {
+    public diagnosticCollection: DiagnosticCollection;
     protected readonly settings: IPythonSettings;
+    private readonly unitTestDiagnosticService: IUnitTestDiagnosticService;
     public abstract get enabled(): boolean;
     protected get outputChannel() {
         return this._outputChannel;
@@ -50,6 +54,8 @@ export abstract class BaseTestManager implements ITestManager {
         this.testCollectionStorage = this.serviceContainer.get<ITestCollectionStorageService>(ITestCollectionStorageService);
         this._testResultsService = this.serviceContainer.get<ITestResultsService>(ITestResultsService);
         this.workspaceService = this.serviceContainer.get<IWorkspaceService>(IWorkspaceService);
+        this.diagnosticCollection = languages.createDiagnosticCollection(this.testProvider);
+        this.unitTestDiagnosticService = serviceContainer.get<IUnitTestDiagnosticService>(IUnitTestDiagnosticService);
         disposables.push(this);
     }
     protected get testDiscoveryCancellationToken(): CancellationToken | undefined {
@@ -246,6 +252,42 @@ export abstract class BaseTestManager implements ITestManager {
                 return Promise.reject<Tests>(reason);
             });
     }
+    public async updateDiagnostics(tests: Tests, messages: IPythonUnitTestMessage[]): Promise<void> {
+        await this.stripStaleDiagnostics(tests, messages);
+
+        // Update relevant file diagnostics for tests that have problems.
+        const uniqueMsgFiles = messages.reduce((filtered, msg) => {
+            if (filtered.indexOf(msg.testFilePath) === -1 && msg.testFilePath !== undefined) {
+                filtered.push(msg.testFilePath);
+            }
+            return filtered;
+        }, []);
+        const fs = this.serviceContainer.get<IFileSystem>(IFileSystem);
+        for (const msgFile of uniqueMsgFiles) {
+            // Check all messages against each test file.
+            const fileUri = Uri.file(msgFile);
+            if (!this.diagnosticCollection.has(fileUri)) {
+                // Create empty diagnostic for file URI so the rest of the logic can assume one already exists.
+                const diagnostics: Diagnostic[] = [];
+                this.diagnosticCollection.set(fileUri, diagnostics);
+            }
+            // Get the diagnostics for this file's URI before updating it so old tests that weren't run can still show problems.
+            const oldDiagnostics = this.diagnosticCollection.get(fileUri);
+            const newDiagnostics: Diagnostic[] = [];
+            for (const diagnostic of oldDiagnostics) {
+                newDiagnostics.push(diagnostic);
+            }
+            for (const msg of messages) {
+                if (fs.arePathsSame(fileUri.fsPath, Uri.file(msg.testFilePath).fsPath) && msg.status !== TestStatus.Pass) {
+                    const diagnostic = this.createDiagnostics(msg);
+                    newDiagnostics.push(diagnostic);
+                }
+            }
+
+            // Set the diagnostics for the file.
+            this.diagnosticCollection.set(fileUri, newDiagnostics);
+        }
+    }
     // tslint:disable-next-line:no-any
     protected abstract runTestImpl(tests: Tests, testsToRun?: TestsToRun, runFailedTests?: boolean, debug?: boolean): Promise<any>;
     protected abstract getDiscoveryOptions(ignoreCache: boolean): TestDiscoveryOptions;
@@ -270,4 +312,51 @@ export abstract class BaseTestManager implements ITestManager {
             this.testRunnerCancellationTokenSource = undefined;
         }
     }
+    /**
+     * Whenever a test is run, any previous problems it had should be removed. This runs through
+     * every already existing set of diagnostics for any that match the tests that were just run
+     * so they can be stripped out (as they are now no longer relevant). If the tests pass, then
+     * there is no need to have a diagnostic for it. If they fail, the stale diagnostic will be
+     * replaced by an up-to-date diagnostic showing the most recent problem with that test.
+     *
+     * In order to identify diagnostics associated with the tests that were run, the `nameToRun`
+     * property of each messages is compared to the `code` property of each diagnostic.
+     *
+     * @param messages Details about the tests that were just run.
+     */
+    private async stripStaleDiagnostics(tests: Tests, messages: IPythonUnitTestMessage[]): Promise<void> {
+        this.diagnosticCollection.forEach((diagnosticUri, oldDiagnostics, collection) => {
+            const newDiagnostics: Diagnostic[] = [];
+            for (const diagnostic of oldDiagnostics) {
+                const matchingMsg = messages.find((msg) => msg.code === diagnostic.code);
+                if (matchingMsg === undefined) {
+                    // No matching message was found, so this test was not included in the test run.
+                    const matchingTest = tests.testFunctions.find((tf) => tf.testFunction.nameToRun === diagnostic.code);
+                    if (matchingTest !== undefined) {
+                        // Matching test was found, so the diagnostic is still relevant.
+                        newDiagnostics.push(diagnostic);
+                    }
+                }
+            }
+            // Set the diagnostics for the file.
+            collection.set(diagnosticUri, newDiagnostics);
+        });
+    }
+
+    private createDiagnostics(message: IPythonUnitTestMessage): Diagnostic {
+        const stackStart = message.locationStack[0];
+        const diagPrefix = this.unitTestDiagnosticService.getMessagePrefix(message.status);
+        const severity = this.unitTestDiagnosticService.getSeverity(message.severity)!;
+        const diagMsg = message.message.split('\n')[0];
+        const diagnostic = new Diagnostic(stackStart.location.range, `${diagPrefix ? `${diagPrefix}: ` : ''}${diagMsg}`, severity);
+        diagnostic.code = message.code;
+        diagnostic.source = message.provider;
+        const relatedInfoArr: DiagnosticRelatedInformation[] = [];
+        for (const frameDetails of message.locationStack) {
+            const relatedInfo = new DiagnosticRelatedInformation(frameDetails.location, frameDetails.lineText);
+            relatedInfoArr.push(relatedInfo);
+        }
+        diagnostic.relatedInformation = relatedInfoArr;
+        return diagnostic;
+    }
 }
diff --git a/src/client/unittests/common/services/unitTestDiagnosticService.ts b/src/client/unittests/common/services/unitTestDiagnosticService.ts
@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { injectable } from 'inversify';
+import { DiagnosticSeverity } from 'vscode';
+import * as localize from '../../../common/utils/localize';
+import { DiagnosticMessageType, IUnitTestDiagnosticService, PythonUnitTestMessageSeverity } from '../../types';
+import { TestStatus } from '../types';
+
+@injectable()
+export class UnitTestDiagnosticService implements IUnitTestDiagnosticService {
+    private MessageTypes = new Map<TestStatus, DiagnosticMessageType>();
+    private MessageSeverities = new Map<PythonUnitTestMessageSeverity, DiagnosticSeverity>();
+    private MessagePrefixes = new Map<DiagnosticMessageType, string>();
+
+    constructor() {
+        this.MessageTypes.set(TestStatus.Error, DiagnosticMessageType.Error);
+        this.MessageTypes.set(TestStatus.Fail, DiagnosticMessageType.Fail);
+        this.MessageTypes.set(TestStatus.Skipped, DiagnosticMessageType.Skipped);
+        this.MessageTypes.set(TestStatus.Pass, DiagnosticMessageType.Pass);
+        this.MessageSeverities.set(PythonUnitTestMessageSeverity.Error, DiagnosticSeverity.Error);
+        this.MessageSeverities.set(PythonUnitTestMessageSeverity.Failure, DiagnosticSeverity.Error);
+        this.MessageSeverities.set(PythonUnitTestMessageSeverity.Skip, DiagnosticSeverity.Information);
+        this.MessageSeverities.set(PythonUnitTestMessageSeverity.Pass, null);
+        this.MessagePrefixes.set(DiagnosticMessageType.Error, localize.UnitTests.testErrorDiagnosticMessage());
+        this.MessagePrefixes.set(DiagnosticMessageType.Fail, localize.UnitTests.testFailDiagnosticMessage());
+        this.MessagePrefixes.set(DiagnosticMessageType.Skipped, localize.UnitTests.testSkippedDiagnosticMessage());
+        this.MessagePrefixes.set(DiagnosticMessageType.Pass, '');
+    }
+    public getMessagePrefix(status: TestStatus): string {
+        return this.MessagePrefixes.get(this.MessageTypes.get(status));
+    }
+    public getSeverity(unitTestSeverity: PythonUnitTestMessageSeverity): DiagnosticSeverity {
+        return this.MessageSeverities.get(unitTestSeverity);
+    }
+}
diff --git a/src/client/unittests/common/types.ts b/src/client/unittests/common/types.ts
@@ -1,5 +1,6 @@
-import { CancellationToken, Disposable, OutputChannel, Uri } from 'vscode';
+import { CancellationToken, DiagnosticCollection, Disposable, OutputChannel, Uri } from 'vscode';
 import { IUnitTestSettings, Product } from '../../common/types';
+import { IPythonUnitTestMessage } from '../types';
 import { CommandSource } from './constants';
 
 export type TestProvider = 'nosetest' | 'pytest' | 'unittest';
@@ -66,6 +67,7 @@ export type TestResult = Node & {
     passed?: boolean;
     time: number;
     line?: number;
+    file?: string;
     message?: string;
     traceback?: string;
     functionsPassed?: number;
@@ -223,6 +225,7 @@ export interface ITestManager extends Disposable {
     readonly enabled: boolean;
     readonly workingDirectory: string;
     readonly workspaceFolder: Uri;
+    diagnosticCollection: DiagnosticCollection;
     stop(): void;
     resetTestResults(): void;
     discoverTests(cmdSource: CommandSource, ignoreCache?: boolean, quietMode?: boolean, userInitiated?: boolean): Promise<Tests>;
@@ -278,3 +281,8 @@ export type PythonVersionInformation = {
     major: number;
     minor: number;
 };
+
+export const ITestMessageService = Symbol('ITestMessageService');
+export interface ITestMessageService {
+    getFilteredTestMessages(rootDirectory: string, testResults: Tests): Promise<IPythonUnitTestMessage[]>;
+}
diff --git a/src/client/unittests/common/xUnitParser.ts b/src/client/unittests/common/xUnitParser.ts
@@ -110,6 +110,7 @@ export function updateResultsFromXmlLogFile(tests: Tests, outputXmlFile: string,
                     }
 
                     result.testFunction.line = getSafeInt(testcase.$.line, null);
+                    result.testFunction.file = testcase.$.file;
                     result.testFunction.time = parseFloat(testcase.$.time);
                     result.testFunction.passed = true;
                     result.testFunction.status = TestStatus.Pass;

diff --git a/src/client/unittests/pytest/main.ts b/src/client/unittests/pytest/main.ts
@@ -5,13 +5,14 @@ import { Product } from '../../common/types';
 import { IServiceContainer } from '../../ioc/types';
 import { PYTEST_PROVIDER } from '../common/constants';
 import { BaseTestManager } from '../common/managers/baseTestManager';
-import { ITestsHelper, TestDiscoveryOptions, TestRunOptions, Tests, TestsToRun } from '../common/types';
-import { IArgumentsService, ITestManagerRunner, TestFilter } from '../types';
+import { ITestMessageService, ITestsHelper, TestDiscoveryOptions, TestRunOptions, Tests, TestsToRun } from '../common/types';
+import { IArgumentsService, IPythonUnitTestMessage, ITestManagerRunner, TestFilter } from '../types';
 
 export class TestManager extends BaseTestManager {
     private readonly argsService: IArgumentsService;
     private readonly helper: ITestsHelper;
     private readonly runner: ITestManagerRunner;
+    private readonly testMessageService: ITestMessageService;
     public get enabled() {
         return this.settings.unitTest.pyTestEnabled;
     }
@@ -21,6 +22,7 @@ export class TestManager extends BaseTestManager {
         this.argsService = this.serviceContainer.get<IArgumentsService>(IArgumentsService, this.testProvider);
         this.helper = this.serviceContainer.get<ITestsHelper>(ITestsHelper);
         this.runner = this.serviceContainer.get<ITestManagerRunner>(ITestManagerRunner, this.testProvider);
+        this.testMessageService = this.serviceContainer.get<ITestMessageService>(ITestMessageService, this.testProvider);
     }
     public getDiscoveryOptions(ignoreCache: boolean): TestDiscoveryOptions {
         const args = this.settings.unitTest.pyTestArgs.slice(0);
@@ -51,6 +53,9 @@ export class TestManager extends BaseTestManager {
             token: this.testRunnerCancellationToken!,
             outChannel: this.outputChannel
         };
-        return this.runner.runTest(this.testResultsService, options, this);
+        const testResults = await this.runner.runTest(this.testResultsService, options, this);
+        const messages: IPythonUnitTestMessage[] = await this.testMessageService.getFilteredTestMessages(this.rootDirectory, testResults);
+        await this.updateDiagnostics(tests, messages);
+        return testResults;
     }
 }
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Create diagnostics for failed/skipped tests that were run with pytest.
		(thanks [Chris NeJame](https://github.com/SalmonMode/))