Code Lenses for Tests

When the `swift.showTestCodeLenses` setting is `true`, show code
lenses inline in the editor above suites and tests. This setting
defaults to `false` and must be enabled first.

The three lenses are Run, Debug, and Run w/ Coverage.

Author: plemarquand

package.json

@@ -306,6 +306,27 @@
         "category": "Test",
         "icon": "$(debug-coverage)"
       },
+      {
+        "command": "swift.runTest",
+        "title": "Run Test",
+        "category": "Test",
+        "icon": "$(testing-run-icon)",
+        "enablement": "false"
+      },
+      {
+        "command": "swift.debugTest",
+        "title": "Debug Test",
+        "category": "Test",
+        "icon": "$(testing-debug-icon)",
+        "enablement": "false"
+      },
+      {
+        "command": "swift.runTestWithCoverage",
+        "title": "Run Test with Coverage",
+        "category": "Test",
+        "icon": "$(debug-coverage)",
+        "enablement": "false"
+      },
       {
         "command": "swift.openDocumentation",
         "title": "Open Documentation",
@@ -552,7 +573,7 @@
         }
       },
       {
-        "title": "Code Coverage",
+        "title": "Testing",
         "properties": {
           "swift.excludeFromCodeCoverage": {
             "description": "A list of paths to exclude from code coverage reports. Paths can be absolute or relative to the workspace root.",
@@ -561,6 +582,11 @@
               "type": "string"
             },
             "default": []
+          },
+          "swift.showTestCodeLenses": {
+            "type": "boolean",
+            "default": false,
+            "markdownDescription": "Controls whether or not to show inline code lenses for running and debugging tests inline, above test and suite declarations."
           }
         }
       },
@@ -1017,6 +1043,18 @@
           "command": "swift.coverAllTests",
           "when": "swift.isActivated"
         },
+        {
+          "command": "swift.runTest",
+          "when": "false"
+        },
+        {
+          "command": "swift.debugTest",
+          "when": "false"
+        },
+        {
+          "command": "swift.runTestWithCoverage",
+          "when": "false"
+        },
         {
           "command": "swift.openEducationalNote",
           "when": "false"
@@ -1748,4 +1786,4 @@
     "vscode-languageclient": "^9.0.1",
     "xml2js": "^0.6.2"
   }
-}
+}

src/TestExplorer/TestCodeLensProvider.ts

@@ -0,0 +1,59 @@
+import * as vscode from "vscode";
+import { TestExplorer } from "./TestExplorer";
+import { flattenTestItemCollection } from "./TestUtils";
+import configuration from "../configuration";
+
+export class TestCodeLensProvider implements vscode.CodeLensProvider, vscode.Disposable {
+    private onDidChangeCodeLensesEmitter = new vscode.EventEmitter<void>();
+    public onDidChangeCodeLenses = this.onDidChangeCodeLensesEmitter.event;
+    private disposables: vscode.Disposable[] = [];
+
+    constructor(private testExplorer: TestExplorer) {
+        this.disposables = [
+            testExplorer.onTestItemsDidChange(() => this.onDidChangeCodeLensesEmitter.fire()),
+            vscode.languages.registerCodeLensProvider({ language: "swift", scheme: "file" }, this),
+        ];
+    }
+
+    dispose() {
+        this.disposables.forEach(disposable => disposable.dispose());
+    }
+
+    public provideCodeLenses(
+        document: vscode.TextDocument,
+        _token: vscode.CancellationToken
+    ): vscode.ProviderResult<vscode.CodeLens[]> {
+        if (configuration.showTestCodeLenses === false) {
+            return [];
+        }
+
+        const items = flattenTestItemCollection(this.testExplorer.controller.items);
+        return items
+            .filter(item => item.uri?.fsPath === document.uri.fsPath)
+            .flatMap(item => this.codeLensesForTestItem(item));
+    }
+
+    private codeLensesForTestItem(item: vscode.TestItem): vscode.CodeLens[] {
+        if (!item.range) {
+            return [];
+        }
+
+        return [
+            new vscode.CodeLens(item.range, {
+                title: "$(play) Run",
+                command: "swift.runTest",
+                arguments: [item],
+            }),
+            new vscode.CodeLens(item.range, {
+                title: "$(debug) Debug",
+                command: "swift.debugTest",
+                arguments: [item],
+            }),
+            new vscode.CodeLens(item.range, {
+                title: "$(debug-coverage) Run w/ Coverage",
+                command: "swift.runTestWithCoverage",
+                arguments: [item],
+            }),
+        ];
+    }
+}

src/TestExplorer/TestExplorer.ts

@@ -27,15 +27,16 @@ import { TargetType } from "../SwiftPackage";
 import { parseTestsFromSwiftTestListOutput } from "./SPMTestDiscovery";
 import { parseTestsFromDocumentSymbols } from "./DocumentSymbolTestDiscovery";
 import { flattenTestItemCollection } from "./TestUtils";
+import { TestCodeLensProvider } from "./TestCodeLensProvider";
 
 /** Build test explorer UI */
 export class TestExplorer {
     static errorTestItemId = "#Error#";
     public controller: vscode.TestController;
     public testRunProfiles: vscode.TestRunProfile[];
+
     private lspTestDiscovery: LSPTestDiscovery;
     private subscriptions: { dispose(): unknown }[];
-    private testFileEdited = true;
     private tokenSource = new vscode.CancellationTokenSource();
 
     // Emits after the `vscode.TestController` has been updated.
@@ -45,9 +46,13 @@ export class TestExplorer {
     public onDidCreateTestRunEmitter = new vscode.EventEmitter<TestRunProxy>();
     public onCreateTestRun: vscode.Event<TestRunProxy>;
 
+    private codeLensProvider?: TestCodeLensProvider;
+
     constructor(public folderContext: FolderContext) {
         this.onTestItemsDidChange = this.onTestItemsDidChangeEmitter.event;
         this.onCreateTestRun = this.onDidCreateTestRunEmitter.event;
+        this.lspTestDiscovery = this.configureLSPTestDiscovery(folderContext);
+        this.codeLensProvider = new TestCodeLensProvider(this);
 
         this.controller = vscode.tests.createTestController(
             folderContext.name,
@@ -66,55 +71,109 @@ export class TestExplorer {
             this.onDidCreateTestRunEmitter
         );
 
+        this.subscriptions = [
+            this.tokenSource,
+            this.controller,
+            this.onTestItemsDidChangeEmitter,
+            this.onDidCreateTestRunEmitter,
+            ...this.testRunProfiles,
+            this.onTestItemsDidChange(() => this.updateSwiftTestContext()),
+            this.discoverUpdatedTestsAfterBuild(folderContext),
+        ];
+    }
+
+    /**
+     * Query the LSP for tests in the document. If the LSP is not available
+     * this method will fallback to the legacy method of parsing document symbols,
+     * but only for XCTests.
+     * @param folder The folder context.
+     * @param uri The document URI. If the document is not part of a test target, this method will do nothing.
+     * @param symbols The document symbols.
+     * @returns A promise that resolves when the tests have been retrieved.
+     */
+    public async getDocumentTests(
+        folder: FolderContext,
+        uri: vscode.Uri,
+        symbols: vscode.DocumentSymbol[]
+    ): Promise<void> {
+        const target = await folder.swiftPackage.getTarget(uri.fsPath);
+        if (!target || target.type !== "test") {
+            return;
+        }
+
+        try {
+            const tests = await this.lspTestDiscovery.getDocumentTests(folder.swiftPackage, uri);
+            TestDiscovery.updateTestsForTarget(
+                this.controller,
+                { id: target.c99name, label: target.name },
+                tests,
+                uri
+            );
+            this.onTestItemsDidChangeEmitter.fire(this.controller);
+        } catch {
+            // Fallback to parsing document symbols for XCTests only
+            const tests = parseTestsFromDocumentSymbols(target.name, symbols, uri);
+            this.updateTests(this.controller, tests, uri);
+        }
+    }
+
+    /**
+     * Creates an LSPTestDiscovery client for the given folder context.
+     */
+    private configureLSPTestDiscovery(folderContext: FolderContext): LSPTestDiscovery {
         const workspaceContext = folderContext.workspaceContext;
         const languageClientManager = workspaceContext.languageClientManager.get(folderContext);
-        this.lspTestDiscovery = new LSPTestDiscovery(languageClientManager);
+        return new LSPTestDiscovery(languageClientManager);
+    }
 
-        // add end of task handler to be called whenever a build task has finished. If
-        // it is the build task for this folder then update the tests
-        const onDidEndTask = folderContext.workspaceContext.tasks.onDidEndTaskProcess(event => {
-            const task = event.execution.task;
-            const execution = task.execution as vscode.ProcessExecution;
-            if (
-                task.scope === this.folderContext.workspaceFolder &&
-                task.group === vscode.TaskGroup.Build &&
-                execution?.options?.cwd === this.folderContext.folder.fsPath &&
-                event.exitCode === 0 &&
-                task.definition.dontTriggerTestDiscovery !== true &&
-                this.testFileEdited
-            ) {
-                this.testFileEdited = false;
+    /**
+     * Configure test discovery for updated tests after a build task has completed.
+     */
+    private discoverUpdatedTestsAfterBuild(folderContext: FolderContext): vscode.Disposable {
+        let testFileEdited = true;
+        const endProcessDisposable = folderContext.workspaceContext.tasks.onDidEndTaskProcess(
+            event => {
+                const task = event.execution.task;
+                const execution = task.execution as vscode.ProcessExecution;
+                if (
+                    task.scope === folderContext.workspaceFolder &&
+                    task.group === vscode.TaskGroup.Build &&
+                    execution?.options?.cwd === folderContext.folder.fsPath &&
+                    event.exitCode === 0 &&
+                    task.definition.dontTriggerTestDiscovery !== true &&
+                    testFileEdited
+                ) {
+                    testFileEdited = false;
 
-                // only run discover tests if the library has tests
-                void this.folderContext.swiftPackage.getTargets(TargetType.test).then(targets => {
-                    if (targets.length > 0) {
-                        void this.discoverTestsInWorkspace(this.tokenSource.token);
-                    }
-                });
+                    // only run discover tests if the library has tests
+                    void folderContext.swiftPackage.getTargets(TargetType.test).then(targets => {
+                        if (targets.length > 0) {
+                            void this.discoverTestsInWorkspace(this.tokenSource.token);
+                        }
+                    });
+                }
             }
-        });
+        );
 
         // add file watcher to catch changes to swift test files
-        const fileWatcher = this.folderContext.workspaceContext.onDidChangeSwiftFiles(({ uri }) => {
-            if (this.testFileEdited === false) {
-                void this.folderContext.getTestTarget(uri).then(target => {
-                    if (target) {
-                        this.testFileEdited = true;
-                    }
-                });
+        const didChangeSwiftFileDisposable = folderContext.workspaceContext.onDidChangeSwiftFiles(
+            ({ uri }) => {
+                if (testFileEdited === false) {
+                    void folderContext.getTestTarget(uri).then(target => {
+                        if (target) {
+                            testFileEdited = true;
+                        }
+                    });
+                }
             }
-        });
+        );
 
-        this.subscriptions = [
-            this.tokenSource,
-            fileWatcher,
-            onDidEndTask,
-            this.controller,
-            this.onTestItemsDidChangeEmitter,
-            this.onDidCreateTestRunEmitter,
-            ...this.testRunProfiles,
-            this.onTestItemsDidChange(() => this.updateSwiftTestContext()),
-        ];
+        return {
+            dispose: () => {
+                endProcessDisposable.dispose();
+                didChangeSwiftFileDisposable.dispose();
+            },
+        };
     }
 
     dispose() {
@@ -193,32 +252,6 @@ export class TestExplorer {
         });
     }
 
-    async getDocumentTests(
-        folder: FolderContext,
-        uri: vscode.Uri,
-        symbols: vscode.DocumentSymbol[]
-    ): Promise<void> {
-        const target = await folder.swiftPackage.getTarget(uri.fsPath);
-        if (!target || target.type !== "test") {
-            return;
-        }
-
-        try {
-            const tests = await this.lspTestDiscovery.getDocumentTests(folder.swiftPackage, uri);
-            TestDiscovery.updateTestsForTarget(
-                this.controller,
-                { id: target.c99name, label: target.name },
-                tests,
-                uri
-            );
-            this.onTestItemsDidChangeEmitter.fire(this.controller);
-        } catch {
-            // Fallback to parsing document symbols for XCTests only
-            const tests = parseTestsFromDocumentSymbols(target.name, symbols, uri);
-            this.updateTests(this.controller, tests, uri);
-        }
-    }
-
     private updateTests(
         controller: vscode.TestController,
         tests: TestDiscovery.TestClass[],
@@ -231,7 +264,7 @@ export class TestExplorer {
     /**
      * Discover tests
      */
-    async discoverTestsInWorkspace(token: vscode.CancellationToken) {
+    private async discoverTestsInWorkspace(token: vscode.CancellationToken) {
         try {
             // If the LSP cannot produce a list of tests it throws and
             // we fall back to discovering tests with SPM.
@@ -249,7 +282,7 @@ export class TestExplorer {
      * Discover tests
      * Uses `swift test --list-tests` to get the list of tests
      */
-    async discoverTestsInWorkspaceSPM(token: vscode.CancellationToken) {
+    private async discoverTestsInWorkspaceSPM(token: vscode.CancellationToken) {
         async function runDiscover(explorer: TestExplorer, firstTry: boolean) {
             try {
                 // we depend on sourcekit-lsp to detect swift-testing tests so let the user know
@@ -379,7 +412,7 @@ export class TestExplorer {
     /**
      * Discover tests
      */
-    async discoverTestsInWorkspaceLSP(token: vscode.CancellationToken) {
+    private async discoverTestsInWorkspaceLSP(token: vscode.CancellationToken) {
         const tests = await this.lspTestDiscovery.getWorkspaceTests(
             this.folderContext.swiftPackage
         );