Adds support for running gocheck tests

package.json

@@ -1567,4 +1567,4 @@
       ]
     }
   }
-}
+}

src/goRunTestCodelens.ts

@@ -9,7 +9,7 @@ import vscode = require('vscode');
 import { CancellationToken, CodeLens, Command, TextDocument } from 'vscode';
 import { GoBaseCodeLensProvider } from './goBaseCodelens';
 import { GoDocumentSymbolProvider } from './goOutline';
-import { getBenchmarkFunctions, getTestFunctions } from './testUtils';
+import { getBenchmarkFunctions, getDocumentSymbols, getTestFunctions, hasMethodTests } from './testUtils';
 import { getCurrentGoPath, getGoConfig } from './util';
 
 export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
@@ -94,7 +94,10 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
 	): Promise<CodeLens[]> {
 		const codelens: CodeLens[] = [];
 
-		const testPromise = getTestFunctions(document, token).then((testFunctions) => {
+		const testPromise = getDocumentSymbols(document, token, true).then((symbols) => {
+			const [methodTests, _] = hasMethodTests(symbols);
+			const testFunctions = getTestFunctions(symbols, methodTests);
+
 			testFunctions.forEach((func) => {
 				const runTestCmd: Command = {
 					title: 'run test',
@@ -114,7 +117,9 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
 			});
 		});
 
-		const benchmarkPromise = getBenchmarkFunctions(document, token).then((benchmarkFunctions) => {
+		const benchmarkPromise = getDocumentSymbols(document, token).then((symbols) => {
+			const benchmarkFunctions = getBenchmarkFunctions(symbols);
+
 			benchmarkFunctions.forEach((func) => {
 				const runBenchmarkCmd: Command = {
 					title: 'run benchmark',

src/goTest.ts

@@ -12,11 +12,13 @@ import {
 	extractInstanceTestName,
 	findAllTestSuiteRuns,
 	getBenchmarkFunctions,
+	getDocumentSymbols,
 	getTestFlags,
 	getTestFunctionDebugArgs,
 	getTestFunctions,
 	getTestTags,
 	goTest,
+	hasMethodTests,
 	TestConfig
 } from './testUtils';
 import { getTempFilePath } from './util';
@@ -45,28 +47,31 @@ export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, cmd: TestA
 		return;
 	}
 
+	const includeImports = cmd === 'benchmark' ? null : true;
 	const getFunctions = cmd === 'benchmark' ? getBenchmarkFunctions : getTestFunctions;
 
 	editor.document.save().then(async () => {
 		try {
-			const testFunctions = await getFunctions(editor.document, null);
+			const documentSymbols = await getDocumentSymbols(editor.document, null, includeImports);
+			const [methodTests, usingGoCheck] = hasMethodTests(documentSymbols);
+			const testFunctions = getFunctions(documentSymbols, methodTests);
 			// We use functionName if it was provided as argument
 			// Otherwise find any test function containing the cursor.
 			const testFunctionName =
 				args && args.functionName
 					? args.functionName
 					: testFunctions
-							.filter((func) => func.range.contains(editor.selection.start))
-							.map((el) => el.name)[0];
+						.filter((func) => func.range.contains(editor.selection.start))
+						.map((el) => el.name)[0];
 			if (!testFunctionName) {
 				vscode.window.showInformationMessage('No test function found at cursor.');
 				return;
 			}
 
 			if (cmd === 'debug') {
-				await debugTestAtCursor(editor, testFunctionName, testFunctions, goConfig);
+				await debugTestAtCursor(editor, testFunctionName, testFunctions, usingGoCheck, goConfig);
 			} else if (cmd === 'benchmark' || cmd === 'test') {
-				await runTestAtCursor(editor, testFunctionName, testFunctions, goConfig, cmd, args);
+				await runTestAtCursor(editor, testFunctionName, testFunctions, usingGoCheck, goConfig, cmd, args);
 			} else {
 				throw new Error('Unsupported command.');
 			}
@@ -83,6 +88,7 @@ async function runTestAtCursor(
 	editor: vscode.TextEditor,
 	testFunctionName: string,
 	testFunctions: vscode.DocumentSymbol[],
+	usingGoCheck: boolean,
 	goConfig: vscode.WorkspaceConfiguration,
 	cmd: TestAtCursorCmd,
 	args: any
@@ -100,7 +106,8 @@ async function runTestAtCursor(
 		functions: testConfigFns,
 		isBenchmark: cmd === 'benchmark',
 		isMod,
-		applyCodeCoverage: goConfig.get<boolean>('coverOnSingleTest')
+		applyCodeCoverage: goConfig.get<boolean>('coverOnSingleTest'),
+		usingGoCheck
 	};
 	// Remember this config as the last executed test.
 	lastTestConfig = testConfig;
@@ -114,9 +121,10 @@ async function debugTestAtCursor(
 	editor: vscode.TextEditor,
 	testFunctionName: string,
 	testFunctions: vscode.DocumentSymbol[],
+	usingGoCheck: boolean,
 	goConfig: vscode.WorkspaceConfiguration
 ) {
-	const args = getTestFunctionDebugArgs(editor.document, testFunctionName, testFunctions);
+	const args = getTestFunctionDebugArgs(editor.document, testFunctionName, testFunctions, usingGoCheck);
 	const tags = getTestTags(goConfig);
 	const buildFlags = tags ? ['-tags', tags] : [];
 	const flagsFromConfig = getTestFlags(goConfig);
@@ -229,31 +237,34 @@ export async function testCurrentFile(
 		return;
 	}
 
+	const includeImports = isBenchmark ? null : true;
 	const getFunctions = isBenchmark ? getBenchmarkFunctions : getTestFunctions;
 	const isMod = await isModSupported(editor.document.uri);
 
-	return editor.document
-		.save()
-		.then(() => {
-			return getFunctions(editor.document, null).then((testFunctions) => {
-				const testConfig: TestConfig = {
-					goConfig,
-					dir: path.dirname(editor.document.fileName),
-					flags: getTestFlags(goConfig, args),
-					functions: testFunctions.map((sym) => sym.name),
-					isBenchmark,
-					isMod,
-					applyCodeCoverage: goConfig.get<boolean>('coverOnSingleTestFile')
-				};
-				// Remember this config as the last executed test.
-				lastTestConfig = testConfig;
-				return goTest(testConfig);
-			});
-		})
-		.then(null, (err) => {
+	return editor.document.save().then(async () => {
+		try {
+			const documentSymbols = await getDocumentSymbols(editor.document, null, includeImports);
+			const [methodTests, usingGoCheck] = hasMethodTests(documentSymbols);
+			const testFunctions = getFunctions(documentSymbols, methodTests);
+
+			const testConfig: TestConfig = {
+				goConfig,
+				dir: path.dirname(editor.document.fileName),
+				flags: getTestFlags(goConfig, args),
+				functions: testFunctions.map((sym) => sym.name),
+				isBenchmark,
+				isMod,
+				applyCodeCoverage: goConfig.get<boolean>('coverOnSingleTestFile'),
+				usingGoCheck
+			};
+			// Remember this config as the last executed test.
+			lastTestConfig = testConfig;
+			return goTest(testConfig);
+		} catch (err) {
 			console.error(err);
 			return Promise.resolve(false);
-		});
+		}
+	});
 }
 
 /**

src/testUtils.ts

@@ -37,6 +37,9 @@ const testFuncRegex = /^Test.*|^Example.*/;
 const testMethodRegex = /^\(([^)]+)\)\.(Test.*)$/;
 const benchmarkRegex = /^Benchmark.*/;
 
+const goCheckPkg = '"gopkg.in/check.v1"';
+const testifyPkg = '"github.com/stretchr/testify/suite"';
+
 /**
  * Input to goTest.
  */
@@ -77,6 +80,10 @@ export interface TestConfig {
 	 * Whether code coverage should be generated and applied.
 	 */
 	applyCodeCoverage?: boolean;
+	/**
+	 * Whether `gocheck` was found in current test file.
+	 */
+	usingGoCheck?: boolean;
 }
 
 export function getTestEnvVars(config: vscode.WorkspaceConfiguration): any {
@@ -119,27 +126,55 @@ export function getTestTags(goConfig: vscode.WorkspaceConfiguration): string {
 /**
  * Returns all Go unit test functions in the given source file.
  *
- * @param the URI of a Go source file.
+ * @param symbols symbols in the Go source file.
+ * @param methodTests whether or not the file contains method tests.
  * @return test function symbols for the source file.
  */
 export function getTestFunctions(
+	symbols: vscode.DocumentSymbol[],
+	methodTests: boolean
+): vscode.DocumentSymbol[] {
+	return symbols.filter(
+		(sym) =>
+			sym.kind === vscode.SymbolKind.Function &&
+			(testFuncRegex.test(sym.name) || (methodTests && testMethodRegex.test(sym.name)))
+	);
+}
+
+/**
+ * Pull all symbols from the document.
+ */
+export function getDocumentSymbols(
 	doc: vscode.TextDocument,
-	token: vscode.CancellationToken
+	token: vscode.CancellationToken,
+	includeImports?: boolean
 ): Thenable<vscode.DocumentSymbol[]> {
-	const documentSymbolProvider = new GoDocumentSymbolProvider(true);
+	const documentSymbolProvider = new GoDocumentSymbolProvider(includeImports);
 	return documentSymbolProvider
 		.provideDocumentSymbols(doc, token)
-		.then((symbols) => symbols[0].children)
-		.then((symbols) => {
-			const testify = symbols.some(
-				(sym) => sym.kind === vscode.SymbolKind.Namespace && sym.name === '"github.com/stretchr/testify/suite"'
-			);
-			return symbols.filter(
-				(sym) =>
-					sym.kind === vscode.SymbolKind.Function &&
-					(testFuncRegex.test(sym.name) || (testify && testMethodRegex.test(sym.name)))
-			);
-		});
+		.then((symbols) => symbols[0].children);
+}
+
+/**
+ * Check symbols for the presence of test suites that use methods.
+ */
+export function hasMethodTests(
+	symbols: vscode.DocumentSymbol[]
+): boolean[] {
+	let usingGoChck = false;
+	const usingMethodTests = symbols.some(
+		(sym) => {
+			const isNamespace = sym.kind === vscode.SymbolKind.Namespace;
+			if (isNamespace && sym.name === goCheckPkg) {
+				usingGoChck = true;
+				return true;
+			} else {
+				return isNamespace && sym.name === testifyPkg;
+			}
+		}
+	);
+
+	return [usingMethodTests, usingGoChck];
 }
 
 /**
@@ -156,20 +191,46 @@ export function extractInstanceTestName(symbolName: string): string {
 	return match[2];
 }
 
+/**
+ * Extracts test method name and the method suite of a suite test function.
+ * For example a symbol with name "(*testSuite).TestMethod" will return ["testSuite", "TestMethod"].
+ *
+ * @param symbolName Symbol Name to extract method name from.
+ */
+export function extractInstanceTestNameWithSuite(symbolName: string): string[] {
+	const match = symbolName.match(testMethodRegex);
+	if (!match || match.length !== 3) {
+		return null;
+	}
+	const [suiteName, methodName] = match.slice(1, 3);
+	let trimmedSuiteName = suiteName;
+	if (trimmedSuiteName[0] === '*') {
+		trimmedSuiteName = trimmedSuiteName.substring(1);
+	}
+
+	return [trimmedSuiteName, methodName];
+}
+
 /**
  * Gets the appropriate debug arguments for a debug session on a test function.
  * @param document  The document containing the tests
  * @param testFunctionName The test function to get the debug args
  * @param testFunctions The test functions found in the document
+ * @param usingGoCheck Whether or not the file is using `gocheck`
  */
 export function getTestFunctionDebugArgs(
 	document: vscode.TextDocument,
 	testFunctionName: string,
-	testFunctions: vscode.DocumentSymbol[]
+	testFunctions: vscode.DocumentSymbol[],
+	usingGoCheck: boolean,
 ): string[] {
 	if (benchmarkRegex.test(testFunctionName)) {
 		return ['-test.bench', '^' + testFunctionName + '$', '-test.run', 'a^'];
 	}
+	if (usingGoCheck) {
+		const [suiteName, testName] = extractInstanceTestNameWithSuite(testFunctionName);
+		return ['-check.f', `${suiteName}.${testName}`];
+	}
 	const instanceMethod = extractInstanceTestName(testFunctionName);
 	if (instanceMethod) {
 		const testFns = findAllTestSuiteRuns(document, testFunctions);
@@ -199,26 +260,20 @@ export function findAllTestSuiteRuns(
 /**
  * Returns all Benchmark functions in the given source file.
  *
- * @param the URI of a Go source file.
+ * @param symbols symbols in the Go source file.
  * @return benchmark function symbols for the source file.
  */
 export function getBenchmarkFunctions(
-	doc: vscode.TextDocument,
-	token: vscode.CancellationToken
-): Thenable<vscode.DocumentSymbol[]> {
-	const documentSymbolProvider = new GoDocumentSymbolProvider();
-	return documentSymbolProvider
-		.provideDocumentSymbols(doc, token)
-		.then((symbols) => symbols[0].children)
-		.then((symbols) =>
-			symbols.filter((sym) => sym.kind === vscode.SymbolKind.Function && benchmarkRegex.test(sym.name))
-		);
+	symbols: vscode.DocumentSymbol[],
+	_?: boolean
+): vscode.DocumentSymbol[] {
+	return symbols.filter((sym) => sym.kind === vscode.SymbolKind.Function && benchmarkRegex.test(sym.name));
 }
 
 /**
  * Runs go test and presents the output in the 'Go' channel.
  *
- * @param goConfig Configuration for the Go extension.
+ * @param testconfig Configuration for running tests.
  */
 export async function goTest(testconfig: TestConfig): Promise<boolean> {
 	const tmpCoverPath = getTempFilePath('go-code-cover');
@@ -444,11 +499,12 @@ function targetArgs(testconfig: TestConfig): Array<string> {
 			params = ['-bench', util.format('^(%s)$', testconfig.functions.join('|'))];
 		} else {
 			let testFunctions = testconfig.functions;
-			let testifyMethods = testFunctions.filter((fn) => testMethodRegex.test(fn));
-			if (testifyMethods.length > 0) {
-				// filter out testify methods
+			const testMethods = testFunctions.filter((fn) => testMethodRegex.test(fn));
+			let testMethodsWithSuite: string[][] = [];
+			if (testMethods.length > 0) {
+				// filter out test methods
 				testFunctions = testFunctions.filter((fn) => !testMethodRegex.test(fn));
-				testifyMethods = testifyMethods.map(extractInstanceTestName);
+				testMethodsWithSuite = testMethods.map(extractInstanceTestNameWithSuite);
 			}
 
 			// we might skip the '-run' param when running only testify methods, which will result
@@ -457,8 +513,28 @@ function targetArgs(testconfig: TestConfig): Array<string> {
 			if (testFunctions.length > 0) {
 				params = params.concat(['-run', util.format('^(%s)$', testFunctions.join('|'))]);
 			}
-			if (testifyMethods.length > 0) {
-				params = params.concat(['-testify.m', util.format('^(%s)$', testifyMethods.join('|'))]);
+
+			const numMethods = testMethodsWithSuite.length;
+			if (numMethods > 0) {
+				if (testconfig.usingGoCheck) {
+					const testNames = testMethodsWithSuite.reduce((acc, [suiteName, instanceMethod], index) => {
+						if (index === numMethods - 1) {
+							return acc + `${suiteName}.${instanceMethod}`;
+						} else {
+							return acc + `${suiteName}.${instanceMethod}|`;
+						}
+					}, '');
+					params = params.concat(['-check.f', util.format('(%s)', testNames)]);
+				} else {
+					const testNames = testMethodsWithSuite.reduce((acc, [_, instanceMethod], index) => {
+						if (index === numMethods - 1) {
+							return acc + instanceMethod;
+						} else {
+							return acc + `${instanceMethod}|`;
+						}
+					}, '');
+					params = params.concat(['-testify.m', util.format('^(%s)$', testNames)]);
+				}
 			}
 		}
 		return params;