wip

Author: Cr4zySheep

extension/src/goTest.ts

@@ -18,12 +18,12 @@ import {
 	getBenchmarkFunctions,
 	getTestFlags,
 	getTestFunctionDebugArgs,
-	getSuiteToTestMap,
-	getTestFunctions,
+	getTestFunctionsAndTestSuite,
 	getTestTags,
 	goTest,
 	TestConfig,
-	SuiteToTestMap
+	SuiteToTestMap,
+	getTestFunctions
 } from './testUtils';
 
 // lastTestConfig holds a reference to the last executed TestConfig which allows
@@ -54,9 +54,11 @@ async function _testAtCursor(
 		throw new NotFoundError('No tests found. Current file is not a test file.');
 	}
 
-	const getFunctions = cmd === 'benchmark' ? getBenchmarkFunctions : getTestFunctions;
-	const testFunctions = (await getFunctions(goCtx, editor.document)) ?? [];
-	const suiteToTest = await getSuiteToTestMap(goCtx, editor.document);
+	const { testFunctions, suiteToTest } = await getTestFunctionsAndTestSuite(
+		cmd === 'benchmark',
+		goCtx,
+		editor.document
+	);
 	// We use functionName if it was provided as argument
 	// Otherwise find any test function containing the cursor.
 	const testFunctionName =
@@ -95,8 +97,7 @@ async function _subTestAtCursor(
 	}
 
 	await editor.document.save();
-	const testFunctions = (await getTestFunctions(goCtx, editor.document)) ?? [];
-	const suiteToTest = await getSuiteToTestMap(goCtx, editor.document);
+	const { testFunctions, suiteToTest } = await getTestFunctionsAndTestSuite(false, goCtx, editor.document);
 	// We use functionName if it was provided as argument
 	// Otherwise find any test function containing the cursor.
 	const currentTestFunctions = testFunctions.filter((func) => func.range.contains(editor.selection.start));

extension/src/goTest/run.ts

@@ -21,14 +21,7 @@ import vscode = require('vscode');
 import { outputChannel } from '../goStatus';
 import { isModSupported } from '../goModules';
 import { getGoConfig } from '../config';
-import {
-	getBenchmarkFunctions,
-	getTestFlags,
-	getSuiteToTestMap,
-	getTestFunctions,
-	goTest,
-	GoTestOutput
-} from '../testUtils';
+import { getTestFlags, getTestFunctionsAndTestSuite, goTest, GoTestOutput } from '../testUtils';
 import { GoTestResolver } from './resolve';
 import { dispose, forEachAsync, GoTest, Workspace } from './utils';
 import { GoTestProfiler, ProfilingOptions } from './profile';
@@ -168,9 +161,11 @@ export class GoTestRunner {
 		await doc.save();
 
 		const goConfig = getGoConfig(test.uri);
-		const getFunctions = kind === 'benchmark' ? getBenchmarkFunctions : getTestFunctions;
-		const testFunctions = await getFunctions(this.goCtx, doc, token);
-		const suiteToTest = await getSuiteToTestMap(this.goCtx, doc, token);
+		const { testFunctions, suiteToTest } = await getTestFunctionsAndTestSuite(
+			kind === 'benchmark',
+			this.goCtx,
+			doc
+		);
 
 		// TODO Can we get output from the debug session, in order to check for
 		// run/pass/fail events?

extension/src/testUtils.ts

@@ -155,27 +155,46 @@ export async function getTestFunctions(
 	doc: vscode.TextDocument,
 	token?: vscode.CancellationToken
 ): Promise<vscode.DocumentSymbol[] | undefined> {
+	const result = await getTestFunctionsAndTestifyHint(goCtx, doc, token);
+	return result.testFunctions;
+}
+
+export async function getTestFunctionsAndTestifyHint(
+	goCtx: GoExtensionContext,
+	doc: vscode.TextDocument,
+	token?: vscode.CancellationToken
+): Promise<{ testFunctions?: vscode.DocumentSymbol[]; foundTestifyTestFunction?: boolean }> {
 	const documentSymbolProvider = GoDocumentSymbolProvider(goCtx, true);
 	const symbols = await documentSymbolProvider.provideDocumentSymbols(doc);
 	if (!symbols || symbols.length === 0) {
-		return;
+		return {};
 	}
 	const symbol = symbols[0];
 	if (!symbol) {
-		return;
+		return {};
 	}
 	const children = symbol.children;
 
 	// With gopls symbol provider, the symbols have the imports of all
 	// the package, so suite tests from all files will be found.
 	const testify = importsTestify(symbols);
-	return children.filter(
+
+	const allTestFunctions = children.filter(
 		(sym) =>
-			(sym.kind === vscode.SymbolKind.Function || sym.kind === vscode.SymbolKind.Method) &&
+			sym.kind === vscode.SymbolKind.Function &&
 			// Skip TestMain(*testing.M) - see https://github.com/golang/vscode-go/issues/482
 			!testMainRegex.test(doc.lineAt(sym.range.start.line).text) &&
-			(testFuncRegex.test(sym.name) || fuzzFuncRegx.test(sym.name) || (testify && testMethodRegex.test(sym.name)))
+			(testFuncRegex.test(sym.name) || fuzzFuncRegx.test(sym.name))
 	);
+
+	const allTestMethods = testify
+		? children.filter((sym) => sym.kind === vscode.SymbolKind.Method && testMethodRegex.test(sym.name))
+		: [];
+
+	return {
+		testFunctions: allTestFunctions.concat(allTestMethods),
+		foundTestifyTestFunction: allTestMethods.length > 0
+	};
 }
 
 /**
@@ -725,3 +744,23 @@ export function importsTestify(syms: vscode.DocumentSymbol[]): boolean {
 			(sym.name === '"github.com/stretchr/testify/suite"' || sym.name === 'github.com/stretchr/testify/suite')
 	);
 }
+
+export async function getTestFunctionsAndTestSuite(
+	isBenchmark: boolean,
+	goCtx: GoExtensionContext,
+	document: vscode.TextDocument
+): Promise<{ testFunctions: vscode.DocumentSymbol[]; suiteToTest: SuiteToTestMap }> {
+	if (isBenchmark) {
+		return {
+			testFunctions: (await getBenchmarkFunctions(goCtx, document)) ?? [],
+			suiteToTest: {}
+		};
+	}
+
+	const { testFunctions, foundTestifyTestFunction } = await getTestFunctionsAndTestifyHint(goCtx, document);
+
+	return {
+		testFunctions: testFunctions ?? [],
+		suiteToTest: foundTestifyTestFunction ? await getSuiteToTestMap(goCtx, document) : {}
+	};
+}