run benchmark at cursor (#972)

package.json

@@ -90,6 +90,11 @@
  "title": "Go: Test Function At Cursor",
  "description": "Runs a unit test at the cursor."
  },
+ {
+ "command": "go.benchmark.cursor",
+ "title": "Go: Benchmark Function At Cursor",
+ "description": "Runs a benchmark at the cursor."
+ },
  {
  "command": "go.test.file",
  "title": "Go: Test File",
@@ -825,6 +830,11 @@
  "command": "go.test.cursor",
  "group": "Go group 1"
  },
+ {
+ "when": "editorTextFocus && config.go.editorContextMenuCommands.benchmarkAtCursor && resourceLangId == go && !config.editor.codeLens",
+ "command": "go.benchmark.cursor",
+ "group": "Go group 1"
+ },
  {
  "when": "editorTextFocus && config.go.editorContextMenuCommands.testFile && resourceLangId == go",
  "command": "go.test.file",

src/goMain.ts

@@ -172,7 +172,14 @@ export function activate(ctx: vscode.ExtensionContext): void {
 
  ctx.subscriptions.push(vscode.commands.registerCommand('go.test.cursor', (args) => {
  let goConfig = vscode.workspace.getConfiguration('go', vscode.window.activeTextEditor ? vscode.window.activeTextEditor.document.uri : null);
- testAtCursor(goConfig, args);
+ let isBenchmark = false;
+ testAtCursor(goConfig, isBenchmark, args);
+ }));
+
+ ctx.subscriptions.push(vscode.commands.registerCommand('go.benchmark.cursor', (args) => {
+ let goConfig = vscode.workspace.getConfiguration('go', vscode.window.activeTextEditor ? vscode.window.activeTextEditor.document.uri : null);
+ let isBenchmark = true;
+ testAtCursor(goConfig, isBenchmark, args);
  }));
 
  ctx.subscriptions.push(vscode.commands.registerCommand('go.test.package', (args) => {

src/goRunTestCodelens.ts

@@ -8,7 +8,7 @@
 import vscode = require('vscode');
 import path = require('path');
 import { CodeLensProvider, TextDocument, CancellationToken, CodeLens, Command } from 'vscode';
-import { getTestFunctions, getTestEnvVars, getTestFlags } from './testUtils';
+import { getTestFunctions, hasTestFunctionPrefix, hasBenchmarkFunctionPrefix, getTestEnvVars, getTestFlags } from './testUtils';
 import { GoDocumentSymbolProvider } from './goOutline';
 import { getCurrentGoPath } from './util';
 import { GoBaseCodeLensProvider } from './goBaseCodelens';
@@ -65,14 +65,14 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
  }
 
  private getCodeLensForFunctions(vsConfig: vscode.WorkspaceConfiguration, document: TextDocument): Thenable<CodeLens[]> {
- return getTestFunctions(document).then(testFunctions => {
- let codelens = [];
+ let codelens: CodeLens[] = [];
 
+ getTestFunctions(document, hasTestFunctionPrefix).then(testFunctions => {
  testFunctions.forEach(func => {
  let runTestCmd: Command = {
  title: 'run test',
  command: 'go.test.cursor',
- arguments: [ { functionName: func.name} ]
+ arguments: [ { functionName: func.name } ]
  };
 
  const args = ['-test.run', func.name];
@@ -95,7 +95,21 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
  codelens.push(new CodeLens(func.location.range, runTestCmd));
  codelens.push(new CodeLens(func.location.range, debugTestCmd));
  });
- return codelens;
  });
+
+ getTestFunctions(document, hasBenchmarkFunctionPrefix).then(benchmarkFunctions => {
+ benchmarkFunctions.forEach(func => {
+ let runBenchmarkCmd: Command = {
+ title: 'run benchmark',
+ command: 'go.benchmark.cursor',
+ arguments: [ { functionName: func.name, isBenchmark: true } ]
+ };
+
+ codelens.push(new CodeLens(func.location.range, runBenchmarkCmd));
+ });
+
+ });
+
+ return Promise.resolve(codelens);
  }
 }

src/goTest.ts

@@ -8,7 +8,7 @@
 import path = require('path');
 import vscode = require('vscode');
 import os = require('os');
-import { goTest, TestConfig, getTestEnvVars, getTestFlags, getTestFunctions } from './testUtils';
+import { goTest, hasTestFunctionPrefix, hasBenchmarkFunctionPrefix, TestConfig, getTestEnvVars, getTestFlags, getTestFunctions } from './testUtils';
 import { getCoverage } from './goCover';
 
 // lastTestConfig holds a reference to the last executed TestConfig which allows
@@ -21,7 +21,7 @@ let lastTestConfig: TestConfig;
 *
 * @param goConfig Configuration for the Go extension.
 */
-export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, args: any) {
+export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, isBenchmark: boolean, args: any) {
  let editor = vscode.window.activeTextEditor;
  if (!editor) {
  vscode.window.showInformationMessage('No editor is active.');
@@ -31,8 +31,14 @@ export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, args: any)
  vscode.window.showInformationMessage('No tests found. Current file is not a test file.');
  return;
  }
+
+ let checker = hasTestFunctionPrefix;
+ if (isBenchmark) {
+ checker = hasBenchmarkFunctionPrefix;
+ }
+
  editor.document.save().then(() => {
- return getTestFunctions(editor.document).then(testFunctions => {
+ return getTestFunctions(editor.document, checker).then(testFunctions => {
  let testFunctionName: string;
 
  // We use functionName if it was provided as argument
@@ -58,8 +64,10 @@ export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, args: any)
  goConfig: goConfig,
  dir: path.dirname(editor.document.fileName),
  flags: getTestFlags(goConfig, args),
- functions: [testFunctionName]
+ functions: [testFunctionName],
+ isBenchmark: isBenchmark
  };
+
  // Remember this config as the last executed test.
  lastTestConfig = testConfig;
 
@@ -152,7 +160,7 @@ export function testCurrentFile(goConfig: vscode.WorkspaceConfiguration, args: s
  }
 
  return editor.document.save().then(() => {
- return getTestFunctions(editor.document).then(testFunctions => {
+ return getTestFunctions(editor.document, hasTestFunctionPrefix).then(testFunctions => {
  const testConfig = {
  goConfig: goConfig,
  dir: path.dirname(editor.document.fileName),

src/testUtils.ts

@@ -43,6 +43,10 @@ export interface TestConfig {
  * Run all tests from all sub directories under `dir`
  */
  includeSubDirectories?: boolean;
+ /**
+ * Whether this is a benchmark.
+ */
+ isBenchmark?: boolean;
 }
 
 export function getTestEnvVars(config: vscode.WorkspaceConfiguration): any {
@@ -78,28 +82,48 @@ export function getTestFlags(goConfig: vscode.WorkspaceConfiguration, args: any)
  * @param the URI of a Go source file.
  * @return test function symbols for the source file.
  */
-export function getTestFunctions(doc: vscode.TextDocument): Thenable<vscode.SymbolInformation[]> {
+export function getTestFunctions(doc: vscode.TextDocument, checker: prefixChecker): Thenable<vscode.SymbolInformation[]> {
  let documentSymbolProvider = new GoDocumentSymbolProvider();
  return documentSymbolProvider
  .provideDocumentSymbols(doc, null)
  .then(symbols =>
  symbols.filter(sym =>
  sym.kind === vscode.SymbolKind.Function
- && hasTestFunctionPrefix(sym.name))
+ && checker(sym.name))
  );
 }
 
+/**
+ * Function type for function that given a function name has
+ * returns whether it is of a certain type of prefix.
+ *
+ * @param the function name.
+ * @return whether the name has a function prefix.
+ */
+type prefixChecker = (name: string) => boolean;
+
 /**
  * Returns whether a given function name has a test prefix.
  * Test functions have "Test" or "Example" as a prefix.
  *
  * @param the function name.
  * @return whether the name has a test function prefix.
  */
-function hasTestFunctionPrefix(name: string): boolean {
+export function hasTestFunctionPrefix(name: string): boolean {
  return name.startsWith('Test') || name.startsWith('Example');
 }
 
+/**
+ * Returns whether a given function name has a benchmark prefix.
+ * Benchmark functions have "Benchmark" as a prefix.
+ *
+ * @param the function name.
+ * @return whether the name has a benchmark function prefix.
+ */
+export function hasBenchmarkFunctionPrefix(name: string): boolean {
+ return name.startsWith('Benchmark');
+}
+
 /**
  * Runs go test and presents the output in the 'Go' channel.
  *
@@ -114,7 +138,20 @@ export function goTest(testconfig: TestConfig): Thenable<boolean> {
  }
 
  let buildTags: string = testconfig.goConfig['buildTags'];
- let args = ['test', ...testconfig.flags, '-timeout', testconfig.goConfig['testTimeout']];
+
+ let args: Array<string>;
+ let handleFunc: argsHandleFunc;
+ let testType: string;
+
+ if (testconfig.isBenchmark) {
+ args = ['test', ...testconfig.flags, '-benchmem', '-run=^$'];
+ handleFunc = benchmarkTargetArgs;
+ testType = 'Benchmarks';
+ } else {
+ args = ['test', ...testconfig.flags, '-timeout', testconfig.goConfig['testTimeout']];
+ handleFunc = testTargetArgs;
+ testType = 'Tests';
+ }
  if (buildTags && testconfig.flags.indexOf('-tags') === -1) {
  args.push('-tags');
  args.push(buildTags);
@@ -133,7 +170,7 @@ export function goTest(testconfig: TestConfig): Thenable<boolean> {
  args.push(testconfig.dir.substr(currentGoWorkspace.length + 1));
  }
 
- targetArgs(testconfig).then(targets => {
+ handleFunc(testconfig).then(targets => {
  let outTargets = args.slice(0);
  if (targets.length > 2) {
  outTargets.push('<long arguments omitted>');
@@ -163,14 +200,14 @@ export function goTest(testconfig: TestConfig): Thenable<boolean> {
  errBuf.done();
 
  if (code) {
- outputChannel.appendLine('Error: Tests failed.');
+ outputChannel.appendLine(`Error: ${testType} failed.`);
  } else {
- outputChannel.appendLine('Success: Tests passed.');
+ outputChannel.appendLine(`Success: ${testType} passed.`);
  }
  resolve(code === 0);
  });
  }, err => {
- outputChannel.appendLine('Error: Tests failed.');
+ outputChannel.appendLine(`Error: ${testType} failed.`);
  outputChannel.appendLine(err);
  resolve(false);
  });
@@ -195,12 +232,19 @@ function expandFilePathInOutput(output: string, cwd: string): string {
  return lines.join('\n');
 }
 
+/**
+ * Function type for getting the target arguments.
+ *
+ * @param testconfig Configuration for the Go extension.
+ */
+type argsHandleFunc = (testconfig: TestConfig) => Thenable<Array<string>>;
+
 /**
  * Get the test target arguments.
  *
  * @param testconfig Configuration for the Go extension.
  */
-function targetArgs(testconfig: TestConfig): Thenable<Array<string>> {
+function testTargetArgs(testconfig: TestConfig): Thenable<Array<string>> {
  if (testconfig.functions) {
  return Promise.resolve(['-run', util.format('^%s$', testconfig.functions.join('|'))]);
  } else if (testconfig.includeSubDirectories) {
@@ -213,3 +257,12 @@ function targetArgs(testconfig: TestConfig): Thenable<Array<string>> {
  }
  return Promise.resolve([]);
 }
+
+/**
+ * Get the benchmark target arguments.
+ *
+ * @param testconfig Configuration for the Go extension.
+ */
+function benchmarkTargetArgs(testconfig: TestConfig): Thenable<Array<string>> {
+ return Promise.resolve(['-bench', util.format('^%s$', testconfig.functions.join('|'))]);
+}