src/goTest: fix multifile suite test fails to debug

Author: nirhaas

extension/src/goTest.ts

@@ -18,10 +18,12 @@ import {
 	getBenchmarkFunctions,
 	getTestFlags,
 	getTestFunctionDebugArgs,
+	getSuiteToTestMap,
 	getTestFunctions,
 	getTestTags,
 	goTest,
-	TestConfig
+	TestConfig,
+	SuiteToTestMap
 } from './testUtils';
 
 // lastTestConfig holds a reference to the last executed TestConfig which allows
@@ -54,6 +56,7 @@ async function _testAtCursor(
 
 	const getFunctions = cmd === 'benchmark' ? getBenchmarkFunctions : getTestFunctions;
 	const testFunctions = (await getFunctions(goCtx, editor.document)) ?? [];
+	const suiteToTest = await getSuiteToTestMap(goCtx, editor.document);
 	// We use functionName if it was provided as argument
 	// Otherwise find any test function containing the cursor.
 	const testFunctionName =
@@ -67,9 +70,9 @@ async function _testAtCursor(
 	await editor.document.save();
 
 	if (cmd === 'debug') {
-		return debugTestAtCursor(editor, testFunctionName, testFunctions, goConfig);
+		return debugTestAtCursor(editor, testFunctionName, testFunctions, suiteToTest, goConfig);
 	} else if (cmd === 'benchmark' || cmd === 'test') {
-		return runTestAtCursor(editor, testFunctionName, testFunctions, goConfig, cmd, args);
+		return runTestAtCursor(editor, testFunctionName, testFunctions, suiteToTest, goConfig, cmd, args);
 	} else {
 		throw new Error(`Unsupported command: ${cmd}`);
 	}
@@ -93,6 +96,7 @@ async function _subTestAtCursor(
 
 	await editor.document.save();
 	const testFunctions = (await getTestFunctions(goCtx, editor.document)) ?? [];
+	const suiteToTest = await getSuiteToTestMap(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));
@@ -142,9 +146,9 @@ async function _subTestAtCursor(
 	const escapedName = escapeSubTestName(testFunctionName, subTestName);
 
 	if (cmd === 'debug') {
-		return debugTestAtCursor(editor, escapedName, testFunctions, goConfig);
+		return debugTestAtCursor(editor, escapedName, testFunctions, suiteToTest, goConfig);
 	} else if (cmd === 'test') {
-		return runTestAtCursor(editor, escapedName, testFunctions, goConfig, cmd, args);
+		return runTestAtCursor(editor, escapedName, testFunctions, suiteToTest, goConfig, cmd, args);
 	} else {
 		throw new Error(`Unsupported command: ${cmd}`);
 	}
@@ -202,13 +206,14 @@ async function runTestAtCursor(
 	editor: vscode.TextEditor,
 	testFunctionName: string,
 	testFunctions: vscode.DocumentSymbol[],
+	suiteToTest: SuiteToTestMap,
 	goConfig: vscode.WorkspaceConfiguration,
 	cmd: TestAtCursorCmd,
 	args: any
 ) {
 	const testConfigFns = [testFunctionName];
 	if (cmd !== 'benchmark' && extractInstanceTestName(testFunctionName)) {
-		testConfigFns.push(...findAllTestSuiteRuns(editor.document, testFunctions).map((t) => t.name));
+		testConfigFns.push(...findAllTestSuiteRuns(editor.document, testFunctions, suiteToTest).map((t) => t.name));
 	}
 
 	const isMod = await isModSupported(editor.document.uri);
@@ -259,11 +264,12 @@ export async function debugTestAtCursor(
 	editorOrDocument: vscode.TextEditor | vscode.TextDocument,
 	testFunctionName: string,
 	testFunctions: vscode.DocumentSymbol[],
+	suiteToFunc: SuiteToTestMap,
 	goConfig: vscode.WorkspaceConfiguration,
 	sessionID?: string
 ) {
 	const doc = 'document' in editorOrDocument ? editorOrDocument.document : editorOrDocument;
-	const args = getTestFunctionDebugArgs(doc, testFunctionName, testFunctions);
+	const args = getTestFunctionDebugArgs(doc, testFunctionName, testFunctions, suiteToFunc);
 	const tags = getTestTags(goConfig);
 	const buildFlags = tags ? ['-tags', tags] : [];
 	const flagsFromConfig = getTestFlags(goConfig);

extension/src/goTest/run.ts

@@ -21,7 +21,14 @@ import vscode = require('vscode');
 import { outputChannel } from '../goStatus';
 import { isModSupported } from '../goModules';
 import { getGoConfig } from '../config';
-import { getBenchmarkFunctions, getTestFlags, getTestFunctions, goTest, GoTestOutput } from '../testUtils';
+import {
+	getBenchmarkFunctions,
+	getTestFlags,
+	getSuiteToTestMap,
+	getTestFunctions,
+	goTest,
+	GoTestOutput
+} from '../testUtils';
 import { GoTestResolver } from './resolve';
 import { dispose, forEachAsync, GoTest, Workspace } from './utils';
 import { GoTestProfiler, ProfilingOptions } from './profile';
@@ -163,6 +170,7 @@ export class GoTestRunner {
 		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);
 
 		// TODO Can we get output from the debug session, in order to check for
 		// run/pass/fail events?
@@ -191,7 +199,8 @@ export class GoTestRunner {
 
 		const run = this.ctrl.createTestRun(request, `Debug ${name}`);
 		if (!testFunctions) return;
-		const started = await debugTestAtCursor(doc, escapeSubTestName(name), testFunctions, goConfig, id);
+		const started = await debugTestAtCursor(doc, escapeSubTestName(name), testFunctions, suiteToTest, goConfig, id);
+
 		if (!started) {
 			subs.forEach((s) => s.dispose());
 			run.end();

extension/src/testUtils.ts

@@ -11,6 +11,7 @@ import cp = require('child_process');
 import path = require('path');
 import util = require('util');
 import vscode = require('vscode');
+import { promises as fs } from 'fs';
 
 import { applyCodeCoverageToAllEditors } from './goCover';
 import { toolExecutionEnvironment } from './goEnv';
@@ -50,6 +51,7 @@ const testMethodRegex = /^\(([^)]+)\)\.(Test|Test\P{Ll}.*)$/u;
 const benchmarkRegex = /^Benchmark$|^Benchmark\P{Ll}.*/u;
 const fuzzFuncRegx = /^Fuzz$|^Fuzz\P{Ll}.*/u;
 const testMainRegex = /TestMain\(.*\*testing.M\)/;
+const runTestSuiteRegex = /^\s*suite\.Run\(\w+,\s*(?:&?(?<type1>\w+)\{|new\((?<type2>\w+)\))/mu;
 
 /**
  * Input to goTest.
@@ -164,7 +166,7 @@ export async function getTestFunctions(
 	}
 	const children = symbol.children;
 
-	// With gopls dymbol provider symbols, the symbols have the imports of all
+	// 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(
@@ -199,14 +201,15 @@ export function extractInstanceTestName(symbolName: string): string {
 export function getTestFunctionDebugArgs(
 	document: vscode.TextDocument,
 	testFunctionName: string,
-	testFunctions: vscode.DocumentSymbol[]
+	testFunctions: vscode.DocumentSymbol[],
+	suiteToFunc: SuiteToTestMap
 ): string[] {
 	if (benchmarkRegex.test(testFunctionName)) {
 		return ['-test.bench', '^' + testFunctionName + '$', '-test.run', 'a^'];
 	}
 	const instanceMethod = extractInstanceTestName(testFunctionName);
 	if (instanceMethod) {
-		const testFns = findAllTestSuiteRuns(document, testFunctions);
+		const testFns = findAllTestSuiteRuns(document, testFunctions, suiteToFunc);
 		const testSuiteRuns = ['-test.run', `^${testFns.map((t) => t.name).join('|')}$`];
 		const testSuiteTests = ['-testify.m', `^${instanceMethod}$`];
 		return [...testSuiteRuns, ...testSuiteTests];
@@ -222,12 +225,22 @@ export function getTestFunctionDebugArgs(
  */
 export function findAllTestSuiteRuns(
 	doc: vscode.TextDocument,
-	allTests: vscode.DocumentSymbol[]
+	allTests: vscode.DocumentSymbol[],
+	suiteToFunc: SuiteToTestMap
 ): vscode.DocumentSymbol[] {
-	// get non-instance test functions
-	const testFunctions = allTests?.filter((t) => !testMethodRegex.test(t.name));
-	// filter further to ones containing suite.Run()
-	return testFunctions?.filter((t) => doc.getText(t.range).includes('suite.Run(')) ?? [];
+	const suites = allTests
+		// Find all tests with receivers.
+		?.map((e) => e.name.match(testMethodRegex))
+		.filter((e) => e?.length === 3)
+		// Take out receiever, strip leading *.
+		.map((e) => e && e[1].replace(/^\*/g, ''))
+		// Map receiver name to test that runs "suite.Run".
+		.map((e) => e && suiteToFunc[e])
+		// Filter out empty results.
+		.filter((e): e is vscode.DocumentSymbol => !!e);
+
+	// Dedup.
+	return [...new Set(suites)];
 }
 
 /**
@@ -254,6 +267,59 @@ export async function getBenchmarkFunctions(
 	return children.filter((sym) => sym.kind === vscode.SymbolKind.Function && benchmarkRegex.test(sym.name));
 }
 
+export type SuiteToTestMap = Record<string, vscode.DocumentSymbol>;
+
+/**
+ * Returns a mapping between a package's function receivers to
+ * the test method that initiated them with "suite.Run".
+ *
+ * @param the URI of a Go source file.
+ * @return function symbols from all source files of the package, mapped by target suite names.
+ */
+export async function getSuiteToTestMap(
+	goCtx: GoExtensionContext,
+	doc: vscode.TextDocument,
+	token?: vscode.CancellationToken
+) {
+	// Get all the package documents.
+	const packageDir = path.parse(doc.fileName).dir;
+	const packageContent = await fs.readdir(packageDir, { withFileTypes: true });
+	const packageFilenames = packageContent
+		// Only go files.
+		.filter((dirent) => dirent.isFile())
+		.map((dirent) => dirent.name)
+		.filter((name) => name.endsWith('.go'));
+	const packageDocs = await Promise.all(
+		packageFilenames.map((e) => path.join(packageDir, e)).map(vscode.workspace.openTextDocument)
+	);
+
+	const suiteToTest: SuiteToTestMap = {};
+	for (const packageDoc of packageDocs) {
+		const funcs = await getTestFunctions(goCtx, packageDoc, token);
+		if (!funcs) {
+			continue;
+		}
+
+		for (const func of funcs) {
+			const funcText = packageDoc.getText(func.range);
+
+			// Matches run suites of the types:
+			// type1: suite.Run(t, MySuite{
+			// type1: suite.Run(t, &MySuite{
+			// type2: suite.Run(t, new(MySuite)
+			const matchRunSuite = funcText.match(runTestSuiteRegex);
+			if (!matchRunSuite) {
+				continue;
+			}
+
+			const g = matchRunSuite.groups;
+			suiteToTest[g?.type1 || g?.type2 || ''] = func;
+		}
+	}
+
+	return suiteToTest;
+}
+
 /**
  * go test -json output format.
  * which is a subset of https://golang.org/cmd/test2json/#hdr-Output_Format