tools/v2check: refactor for better testing

tools/v2check/main.go

@@ -11,5 +11,6 @@ import (
 )
 
 func main() {
-	singlechecker.Main(v2check.Analyzer)
+	c := v2check.NewChecker(nil)
+	c.Run(singlechecker.Main)
 }

tools/v2check/v2check/known_change.go

@@ -6,34 +6,43 @@
 package v2check
 
 import (
+	"context"
+	"fmt"
 	"go/ast"
 
 	"golang.org/x/tools/go/analysis"
 )
 
-type context map[string]any
-
-// knownChange models code expressions that must be changed to migrate to v2.
+// KnownChange models code expressions that must be changed to migrate to v2.
 // It is defined by a set of probes that must be true to report the analyzed expression.
 // It also contains a message function that returns a string describing the change.
 // The probes are evaluated in order, and the first one that returns false
 // will cause the expression to be ignored.
 // A predicate can store information in the context, which is available to the message function and
 // to the following probes.
 // It is possible to declare fixes that will be suggested to the user or applied automatically.
-type knownChange struct {
-	ctx     context
-	probes  []func(context, ast.Node, *analysis.Pass) bool
-	fixes   []analysis.SuggestedFix
-	message func() string
+type KnownChange interface {
+	fmt.Stringer
+
+	// Context returns the context associated with the known change.
+	Context() context.Context
+
+	// Probes returns a list of probes that must be true to report the analyzed expression.
+	Probes() []Probe
+
+	// UpdateContext updates the context with the given value.
+	UpdateContext(context.Context)
 }
 
-func (c *knownChange) eval(n ast.Node, pass *analysis.Pass) bool {
-	c.ctx = make(context)
-	for _, p := range c.probes {
-		if !p(c.ctx, n, pass) {
+func eval(k KnownChange, n ast.Node, pass *analysis.Pass) bool {
+	for _, p := range k.Probes() {
+		ctx, ok := p(k.Context(), n, pass)
+		if !ok {
 			return false
 		}
+
+		k.UpdateContext(ctx)
 	}
+
 	return true
 }

tools/v2check/v2check/probe.go

@@ -0,0 +1,48 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2023 Datadog, Inc.
+
+package v2check
+
+import (
+	"context"
+	"go/ast"
+	"go/types"
+	"strings"
+
+	"golang.org/x/tools/go/analysis"
+	"golang.org/x/tools/go/types/typeutil"
+)
+
+type Probe func(context.Context, ast.Node, *analysis.Pass) (context.Context, bool)
+
+// IsFuncCall returns true if the node is a function call.
+// The function call expression is stored in the context as "fn".
+func IsFuncCall(ctx context.Context, n ast.Node, pass *analysis.Pass) (context.Context, bool) {
+	c, ok := n.(*ast.CallExpr)
+	if !ok {
+		return ctx, false
+	}
+	callee := typeutil.Callee(pass.TypesInfo, c)
+	if callee == nil {
+		return ctx, false
+	}
+	fn, ok := callee.(*types.Func)
+	if !ok {
+		return ctx, false
+	}
+	ctx = context.WithValue(ctx, "fn", fn)
+	return ctx, true
+}
+
+// HasPackagePrefix returns true if the selector expression has a package prefix.
+func HasPackagePrefix(prefix string) Probe {
+	return func(ctx context.Context, n ast.Node, pass *analysis.Pass) (context.Context, bool) {
+		fn, ok := ctx.Value("fn").(*types.Func)
+		if !ok {
+			return ctx, false
+		}
+		return ctx, strings.HasPrefix(fn.Pkg().Path(), prefix)
+	}
+}

tools/v2check/v2check/v2check.go

@@ -14,44 +14,65 @@ import (
 	"golang.org/x/tools/go/ast/inspector"
 )
 
-var knownChanges = []*knownChange{}
-
-var Analyzer = &analysis.Analyzer{
-	Name:     "v2check",
-	Doc:      "Migration tool to assist with the dd-trace-go v2 upgrade",
-	Requires: []*analysis.Analyzer{inspect.Analyzer},
-	Run:      run,
+type Checker struct {
+	knownChanges []KnownChange
 }
 
-func run(pass *analysis.Pass) (interface{}, error) {
-	ins, ok := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
-	if !ok {
-		return nil, errors.New("analyzer is not type *inspector.Inspector")
+func (c Checker) Run(handler func(*analysis.Analyzer)) {
+	analyzer := &analysis.Analyzer{
+		Name:     "v2check",
+		Doc:      "Migration tool to assist with the dd-trace-go v2 upgrade",
+		Requires: []*analysis.Analyzer{inspect.Analyzer},
+		Run:      c.runner(),
 	}
-	filter := []ast.Node{
-		(*ast.CallExpr)(nil),
-		(*ast.ImportSpec)(nil),
+
+	if handler == nil {
+		return
 	}
-	ins.Preorder(filter, func(n ast.Node) {
-		var k *knownChange
-		for _, c := range knownChanges {
-			if c.eval(n, pass) {
-				k = c
-				break
-			}
+
+	handler(analyzer)
+}
+
+func (c Checker) runner() func(*analysis.Pass) (interface{}, error) {
+	knownChanges := c.knownChanges
+
+	return func(pass *analysis.Pass) (interface{}, error) {
+		filter := []ast.Node{
+			(*ast.CallExpr)(nil),
+			(*ast.ImportSpec)(nil),
 		}
-		if k == nil {
-			return
+
+		ins, ok := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
+		if !ok {
+			return nil, errors.New("analyzer is not type *inspector.Inspector")
 		}
-		pass.Report(analysis.Diagnostic{
-			Pos:            n.Pos(),
-			End:            n.End(),
-			Category:       "",
-			Message:        k.message(),
-			URL:            "",
-			SuggestedFixes: nil,
-			Related:        nil,
+		ins.Preorder(filter, func(n ast.Node) {
+			var k KnownChange
+			for _, c := range knownChanges {
+				if eval(c, n, pass) {
+					k = c
+					break
+				}
+			}
+			if k == nil {
+				return
+			}
+
+			pass.Report(analysis.Diagnostic{
+				Pos:            n.Pos(),
+				End:            n.End(),
+				Category:       "",
+				Message:        k.String(),
+				URL:            "",
+				SuggestedFixes: nil,
+				Related:        nil,
+			})
 		})
-	})
-	return nil, nil
+
+		return nil, nil
+	}
+}
+
+func NewChecker(knownChanges ...KnownChange) *Checker {
+	return &Checker{knownChanges: knownChanges}
 }

tools/v2check/v2check/v2check_test.go

@@ -3,49 +3,66 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2023 Datadog, Inc.
 
-package v2check
+package v2check_test
 
 import (
-	"go/ast"
+	"context"
 	"go/types"
 	"os"
 	"path"
 	"testing"
 
+	"github.com/DataDog/dd-trace-go/v2/tools/v2check/v2check"
+
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/analysistest"
 )
 
-// newV1Usage detects the usage of any v1 function.
-func newV1Usage() *knownChange {
-	var d knownChange
-	d.ctx = make(context)
-	d.probes = []func(context, ast.Node, *analysis.Pass) bool{
-		isFuncCall,
-		hasPackagePrefix("gopkg.in/DataDog/dd-trace-go.v1/"),
-	}
-	d.message = func() string {
-		fn, ok := d.ctx["fn"].(*types.Func)
-		if !ok {
-			return "unknown"
-		}
-		return fn.FullName()
+type V1Usage struct {
+	ctx context.Context
+}
+
+func (c V1Usage) Context() context.Context {
+	return c.ctx
+}
+
+func (c *V1Usage) UpdateContext(ctx context.Context) {
+	c.ctx = ctx
+}
+
+func (c V1Usage) Probes() []v2check.Probe {
+	return []v2check.Probe{
+		v2check.IsFuncCall,
+		v2check.HasPackagePrefix("gopkg.in/DataDog/dd-trace-go.v1/"),
 	}
-	return &d
 }
 
-func TestMain(m *testing.M) {
-	knownChanges = []*knownChange{
-		newV1Usage(),
+func (c V1Usage) String() string {
+	fn, ok := c.ctx.Value("fn").(*types.Func)
+	if !ok {
+		return "unknown"
 	}
 
-	os.Exit(m.Run())
+	return fn.FullName()
 }
 
 func TestSimple(t *testing.T) {
+	c := v2check.NewChecker(&V1Usage{
+		ctx: context.Background(),
+	})
+	c.Run(testRunner(t))
+}
+
+func testRunner(t *testing.T) func(*analysis.Analyzer) {
+	t.Helper()
+
 	cwd, err := os.Getwd()
 	if err != nil {
 		t.Error(err)
+		return nil
+	}
+
+	return func(a *analysis.Analyzer) {
+		analysistest.Run(t, path.Join(cwd, "..", "_stage"), a)
 	}
-	_ = analysistest.Run(t, path.Join(cwd, "../_stage"), Analyzer)
 }