RunTest with CodeLens

src/goMain.ts

@@ -16,6 +16,7 @@ import { GoReferenceProvider } from './goReferences';
 import { GoDocumentFormattingEditProvider, Formatter } from './goFormat';
 import { GoRenameProvider } from './goRename';
 import { GoDocumentSymbolProvider } from './goOutline';
+import { GoRunTestCodeLensProvider } from './goRunTestCodelens';
 import { GoSignatureHelpProvider } from './goSignature';
 import { GoWorkspaceSymbolProvider } from './goSymbol';
 import { GoCodeActionProvider } from './goCodeAction';
@@ -83,6 +84,7 @@ export function activate(ctx: vscode.ExtensionContext): void {
 	ctx.subscriptions.push(vscode.languages.registerDocumentFormattingEditProvider(GO_MODE, new GoDocumentFormattingEditProvider()));
 	ctx.subscriptions.push(vscode.languages.registerRenameProvider(GO_MODE, new GoRenameProvider()));
 	ctx.subscriptions.push(vscode.languages.registerCodeActionsProvider(GO_MODE, new GoCodeActionProvider()));
+	ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, new GoRunTestCodeLensProvider()));
 	ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, new GoCodeLensProvider()));
 
 	errorDiagnosticCollection = vscode.languages.createDiagnosticCollection('go-error');

src/goRunTestCodelens.ts

@@ -0,0 +1,60 @@
+/*---------------------------------------------------------
+ * Copyright (C) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------*/
+
+'use strict';
+
+import vscode = require('vscode');
+import { CodeLensProvider, TextDocument, CancellationToken, CodeLens, Command } from 'vscode';
+import { getTestFunctions } from './goTest';
+import { GoDocumentSymbolProvider } from './goOutline';
+
+export class GoRunTestCodeLensProvider implements CodeLensProvider {
+	public provideCodeLenses(document: TextDocument, token: CancellationToken): CodeLens[] | Thenable<CodeLens[]> {
+		if (!document.fileName.endsWith('_test.go')) {
+			return;
+		}
+
+		return Promise.all([
+			this.getCodeLensForPackage(document),
+			this.getCodeLensForFunctions(document)
+		]).then(res => {
+			return res[0].concat(res[1]);
+		});
+	}
+
+	private getCodeLensForPackage(document: TextDocument): Thenable<CodeLens[]> {
+		let documentSymbolProvider = new GoDocumentSymbolProvider();
+		return documentSymbolProvider.provideDocumentSymbols(document, null)
+				.then(symbols => symbols.find(sym => sym.kind === vscode.SymbolKind.Package && !!sym.name))
+				.then(pkg => {
+					if (pkg) {
+						const range = pkg.location.range;
+						return [
+							new CodeLens(range, {
+								title: 'run package tests',
+								command: 'go.test.package'
+							}),
+							new CodeLens(range, {
+								title: 'run file tests',
+								command: 'go.test.file'
+							})
+						];
+					}
+				});
+	}
+
+	private getCodeLensForFunctions(document: TextDocument): Thenable<CodeLens[]> {
+		return getTestFunctions(document).then(testFunctions => {
+			return testFunctions.map(func => {
+				let command: Command = {
+					title: 'run test',
+					command: 'go.test.cursor',
+					arguments: [ { functionName: func.name} ]
+				};
+				return new CodeLens(func.location.range, command);
+			});
+		});
+	}
+}

src/goTest.ts

@@ -70,24 +70,32 @@ export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, args: any)
 		return;
 	}
 	getTestFunctions(editor.document).then(testFunctions => {
-		let testFunction: vscode.SymbolInformation;
-		// Find any test function containing the cursor.
-		for (let func of testFunctions) {
-			let selection = editor.selection;
-			if (selection && func.location.range.contains(selection.start)) {
-				testFunction = func;
-				break;
-			}
-		};
-		if (!testFunction) {
+		let testFunctionName: string;
+
+		// We use functionName if it was provided as argument
+		// Otherwise find any test function containing the cursor.
+		if (args.functionName) {
+			testFunctionName = args.functionName;
+		} else {
+			for (let func of testFunctions) {
+				let selection = editor.selection;
+				if (selection && func.location.range.contains(selection.start)) {
+					testFunctionName = func.name;
+					break;
+				}
+			};
+		}
+
+		if (!testFunctionName) {
 			vscode.window.showInformationMessage('No test function found at cursor.');
 			return;
 		}
+
 		return goTest({
 			goConfig: goConfig,
 			dir: path.dirname(editor.document.fileName),
 			flags: getTestFlags(goConfig, args),
-			functions: [testFunction.name]
+			functions: [ testFunctionName ]
 		});
 	}).then(null, err => {
 		console.error(err);
@@ -237,7 +245,7 @@ export function goTest(testconfig: TestConfig): Thenable<boolean> {
  * @param the URI of a Go source file.
  * @return test function symbols for the source file.
  */
-function getTestFunctions(doc: vscode.TextDocument): Thenable<vscode.SymbolInformation[]> {
+export function getTestFunctions(doc: vscode.TextDocument): Thenable<vscode.SymbolInformation[]> {
 	let documentSymbolProvider = new GoDocumentSymbolProvider();
 	return documentSymbolProvider
 		.provideDocumentSymbols(doc, null)