go/analysis/passes/testinggoroutine: also inspect defined/anonymous functions

Currently, testinggoroutine only inspects functions literal invoked as
"go func(){ ... }()".

For defined or anonymous functions like "go f(t)" it didn't traverse the
function body. To fix this, on encountering those kinds of functions,
retrieve the definition node then inspect it.

Fixes golang/go#47470

Change-Id: I83b6eb3bf2689c66aee32f13d34002aa3cd175b2
Reviewed-on: https://go-review.googlesource.com/c/tools/+/338529
Trust: Cuong Manh Le <cuong.manhle.vn@gmail.com>
Run-TryBot: Cuong Manh Le <cuong.manhle.vn@gmail.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Tim King <taking@google.com>
Reviewed-by: Emmanuel Odeke <emmanuel@orijtech.com>
diff --git a/go/analysis/passes/testinggoroutine/testdata/src/a/a.go b/go/analysis/passes/testinggoroutine/testdata/src/a/a.go
index 5fe9041..c211ec3 100644
--- a/go/analysis/passes/testinggoroutine/testdata/src/a/a.go
+++ b/go/analysis/passes/testinggoroutine/testdata/src/a/a.go
@@ -36,19 +36,6 @@
 	}
 }
 
-func BenchmarkBadFatalf(b *testing.B) {
-	var wg sync.WaitGroup
-	defer wg.Wait()
-
-	for i := 0; i < b.N; i++ {
-		wg.Add(1)
-		go func(id int) {
-			defer wg.Done()
-			b.Fatalf("TestFailed: id = %v\n", id) // want "call to .+B.+Fatalf from a non-test goroutine"
-		}(i)
-	}
-}
-
 func TestBadFatal(t *testing.T) {
 	var wg sync.WaitGroup
 	defer wg.Wait()
@@ -62,6 +49,32 @@
 	}
 }
 
+func f(t *testing.T, _ string) {
+	t.Fatal("TestFailed")
+}
+
+func TestBadFatalIssue47470(t *testing.T) {
+	go f(t, "failed test 1") // want "call to .+T.+Fatal from a non-test goroutine"
+
+	g := func(t *testing.T, _ string) {
+		t.Fatal("TestFailed")
+	}
+	go g(t, "failed test 2") // want "call to .+T.+Fatal from a non-test goroutine"
+}
+
+func BenchmarkBadFatalf(b *testing.B) {
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	for i := 0; i < b.N; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			b.Fatalf("TestFailed: id = %v\n", id) // want "call to .+B.+Fatalf from a non-test goroutine"
+		}(i)
+	}
+}
+
 func BenchmarkBadFatal(b *testing.B) {
 	var wg sync.WaitGroup
 	defer wg.Wait()
diff --git a/go/analysis/passes/testinggoroutine/testinggoroutine.go b/go/analysis/passes/testinggoroutine/testinggoroutine.go
index d2b9a56..800bef5 100644
--- a/go/analysis/passes/testinggoroutine/testinggoroutine.go
+++ b/go/analysis/passes/testinggoroutine/testinggoroutine.go
@@ -119,11 +119,29 @@
 	return varTypeName, ok
 }
 
+// goStmtFunc returns the ast.Node of a call expression
+// that was invoked as a go statement. Currently, only
+// function literals declared in the same function, and
+// static calls within the same package are supported.
+func goStmtFun(goStmt *ast.GoStmt) ast.Node {
+	switch goStmt.Call.Fun.(type) {
+	case *ast.Ident:
+		id := goStmt.Call.Fun.(*ast.Ident)
+		if funDecl, ok := id.Obj.Decl.(ast.Node); ok {
+			return funDecl
+		}
+	case *ast.FuncLit:
+		return goStmt.Call.Fun
+	}
+	return goStmt.Call
+}
+
 // checkGoStmt traverses the goroutine and checks for the
 // use of the forbidden *testing.(B, T) methods.
 func checkGoStmt(pass *analysis.Pass, goStmt *ast.GoStmt) {
+	fn := goStmtFun(goStmt)
 	// Otherwise examine the goroutine to check for the forbidden methods.
-	ast.Inspect(goStmt, func(n ast.Node) bool {
+	ast.Inspect(fn, func(n ast.Node) bool {
 		selExpr, ok := n.(*ast.SelectorExpr)
 		if !ok {
 			return true
@@ -147,7 +165,11 @@
 			return true
 		}
 		if typeName, ok := typeIsTestingDotTOrB(field.Type); ok {
-			pass.ReportRangef(selExpr, "call to (*%s).%s from a non-test goroutine", typeName, selExpr.Sel)
+			var fnRange analysis.Range = goStmt
+			if _, ok := fn.(*ast.FuncLit); ok {
+				fnRange = selExpr
+			}
+			pass.ReportRangef(fnRange, "call to (*%s).%s from a non-test goroutine", typeName, selExpr.Sel)
 		}
 		return true
 	})