Lint when t.Cleanup closes over t (microsoft#3237)

Author: jakebailey

_tools/customlint/cleanup.go

@@ -0,0 +1,107 @@
+package customlint
+
+import (
+	"fmt"
+	"go/ast"
+	"go/types"
+
+	"golang.org/x/tools/go/analysis"
+	"golang.org/x/tools/go/analysis/passes/inspect"
+	"golang.org/x/tools/go/ast/inspector"
+)
+
+var cleanupAnalyzer = &analysis.Analyzer{
+	Name: "cleanup",
+	Doc:  "finds t.Cleanup calls where the closure captures a testing variable from the enclosing function",
+	Requires: []*analysis.Analyzer{
+		inspect.Analyzer,
+	},
+	Run: func(pass *analysis.Pass) (any, error) {
+		return (&cleanupPass{pass: pass}).run()
+	},
+}
+
+type cleanupPass struct {
+	pass *analysis.Pass
+}
+
+func (c *cleanupPass) run() (any, error) {
+	in := c.pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
+
+	for cursor := range in.Root().Preorder((*ast.CallExpr)(nil)) {
+		call := cursor.Node().(*ast.CallExpr)
+		c.checkCall(call)
+	}
+
+	return nil, nil
+}
+
+func (c *cleanupPass) checkCall(call *ast.CallExpr) {
+	sel, ok := call.Fun.(*ast.SelectorExpr)
+	if !ok || sel.Sel.Name != "Cleanup" {
+		return
+	}
+
+	recvType := c.pass.TypesInfo.TypeOf(sel.X)
+	if recvType == nil || !isTestingType(recvType) {
+		return
+	}
+
+	if len(call.Args) != 1 {
+		return
+	}
+	funcLit, ok := call.Args[0].(*ast.FuncLit)
+	if !ok {
+		return
+	}
+
+	ast.Inspect(funcLit.Body, func(n ast.Node) bool {
+		c.checkCleanupIdent(n, funcLit)
+		return true
+	})
+}
+
+func (c *cleanupPass) checkCleanupIdent(n ast.Node, funcLit *ast.FuncLit) {
+	ident, ok := n.(*ast.Ident)
+	if !ok {
+		return
+	}
+	obj := c.pass.TypesInfo.Uses[ident]
+	if obj == nil {
+		return
+	}
+	v, ok := obj.(*types.Var)
+	if !ok {
+		return
+	}
+	if !isTestingType(v.Type()) {
+		return
+	}
+	// Flag if the variable is defined outside the closure (captured).
+	if v.Pos() < funcLit.Pos() || v.Pos() >= funcLit.End() {
+		c.pass.Report(analysis.Diagnostic{
+			Pos:     ident.Pos(),
+			End:     ident.End(),
+			Message: fmt.Sprintf("cleanup closure captures %s; the test will have ended when cleanup runs", ident.Name),
+		})
+	}
+}
+
+func isTestingType(t types.Type) bool {
+	if ptr, ok := t.(*types.Pointer); ok {
+		t = ptr.Elem()
+	}
+	named, ok := t.(*types.Named)
+	if !ok {
+		return false
+	}
+	obj := named.Obj()
+	if obj.Pkg() == nil || obj.Pkg().Path() != "testing" {
+		return false
+	}
+	switch obj.Name() {
+	case "T", "B", "F", "TB":
+		return true
+	}
+	return false
+}

_tools/customlint/plugin.go

@@ -16,6 +16,7 @@ type plugin struct{}
 func (f *plugin) BuildAnalyzers() ([]*analysis.Analyzer, error) {
 	return []*analysis.Analyzer{
 		bitclearAnalyzer,
+		cleanupAnalyzer,
 		emptyCaseAnalyzer,
 		forbidParentAccessAnalyzer,
 		shadowAnalyzer,

_tools/customlint/plugin_test.go

@@ -33,7 +33,7 @@ func TestPlugin(t *testing.T) {
 	var plugin plugin
 
 	config := &packages.Config{
-		Mode: packages.LoadSyntax,
+		Mode: packages.LoadAllSyntax,
 		Dir:  testdataDir,
 		Env:  append(os.Environ(), "GO111MODULE=on", "GOPROXY=off", "GOWORK=off"),
 	}

_tools/customlint/testdata/cleanup/cleanup.go

@@ -0,0 +1,40 @@
+package cleanup
+
+import "testing"
+
+func badT(t *testing.T) {
+	t.Cleanup(func() {
+		t.Log("done")
+	})
+}
+
+func badB(b *testing.B) {
+	b.Cleanup(func() {
+		b.Log("done")
+	})
+}
+
+func badTB(tb testing.TB) {
+	tb.Cleanup(func() {
+		tb.Log("done")
+	})
+}
+
+func badMultipleRefs(t *testing.T) {
+	t.Cleanup(func() {
+		t.Log("first")
+		t.Log("second")
+	})
+}
+
+func good(t *testing.T) {
+	t.Cleanup(func() {
+		println("no capture")
+	})
+}
+
+func goodNamedFunc(t *testing.T) {
+	t.Cleanup(namedFunc)
+}
+
+func namedFunc() {}

_tools/customlint/testdata/cleanup/cleanup.go.golden

@@ -0,0 +1,51 @@
+	package cleanup
+	
+	import "testing"
+	
+	func badT(t *testing.T) {
+		t.Cleanup(func() {
+			t.Log("done")
+			~
+!!! cleanup: cleanup closure captures t; the test will have ended when cleanup runs
+		})
+	}
+	
+	func badB(b *testing.B) {
+		b.Cleanup(func() {
+			b.Log("done")
+			~
+!!! cleanup: cleanup closure captures b; the test will have ended when cleanup runs
+		})
+	}
+	
+	func badTB(tb testing.TB) {
+		tb.Cleanup(func() {
+			tb.Log("done")
+			~~
+!!! cleanup: cleanup closure captures tb; the test will have ended when cleanup runs
+		})
+	}
+	
+	func badMultipleRefs(t *testing.T) {
+		t.Cleanup(func() {
+			t.Log("first")
+			~
+!!! cleanup: cleanup closure captures t; the test will have ended when cleanup runs
+			t.Log("second")
+			~
+!!! cleanup: cleanup closure captures t; the test will have ended when cleanup runs
+		})
+	}
+	
+	func good(t *testing.T) {
+		t.Cleanup(func() {
+			println("no capture")
+		})
+	}
+	
+	func goodNamedFunc(t *testing.T) {
+		t.Cleanup(namedFunc)
+	}
+	
+	func namedFunc() {}
+	

internal/lsp/server_completion_test.go

@@ -46,7 +46,7 @@ func initCompletionClient(t *testing.T, files map[string]string, prefs *lsutil.U
 		FS:                 fs,
 		DefaultLibraryPath: bundled.LibPath(),
 	}, onServerRequest)
-	t.Cleanup(func() { assert.NilError(t, closeClient()) })
+	t.Cleanup(func() { _ = closeClient() })
 
 	initMsg, _, ok := lsptestutil.SendRequest(t, client, lsproto.InitializeInfo, &lsproto.InitializeParams{
 		Capabilities: &lsproto.ClientCapabilities{},

internal/lsp/server_projectinfo_test.go

@@ -37,7 +37,7 @@ func initProjectInfoClient(t *testing.T, files map[string]string) *lsptestutil.L
 		FS:                 fs,
 		DefaultLibraryPath: bundled.LibPath(),
 	}, onServerRequest)
-	t.Cleanup(func() { assert.NilError(t, closeClient()) })
+	t.Cleanup(func() { _ = closeClient() })
 
 	initMsg, _, ok := lsptestutil.SendRequest(t, client, lsproto.InitializeInfo, &lsproto.InitializeParams{
 		Capabilities: &lsproto.ClientCapabilities{},