better way to find the actual expression

In gomega expression, the current code knows how to skip the helper method like `WithOffset()` and get to the actual method (e.g. `Expect`). The problem is that this logic is done by comparing to a static list of method names. If a new method will be added, we'll have to add it to the list each time. This commit fixes the issue by skiping any non-actual method. Also, move things around for better oerdering.
nunnatsa · Nov 4, 2024 · de271ea · de271ea
1 parent 1309578
commit de271ea
Show file tree

Hide file tree

Showing 5 changed files with 248 additions and 232 deletions.
diff --git a/internal/expression/actual/asyncfuncarg.go b/internal/expression/actual/asyncfuncarg.go
@@ -1,9 +1,10 @@
 package actual
 
 import (
-	"github.com/nunnatsa/ginkgolinter/internal/gomegahandler"
-	"github.com/nunnatsa/ginkgolinter/internal/interfaces"
 	gotypes "go/types"
+
+	"github.com/nunnatsa/ginkgolinter/internal/gomegainfo"
+	"github.com/nunnatsa/ginkgolinter/internal/interfaces"
 )
 
 func getAsyncFuncArg(sig *gotypes.Signature) ArgPayload {
@@ -16,7 +17,7 @@ func getAsyncFuncArg(sig *gotypes.Signature) ArgPayload {
 
 	if sig.Params().Len() > 0 {
 		arg := sig.Params().At(0).Type()
-		if gomegahandler.IsGomegaType(arg) && sig.Results().Len() == 0 {
+		if gomegainfo.IsGomegaType(arg) && sig.Results().Len() == 0 {
 			argType |= FuncSigArgType | GomegaParamArgType
 		}
 	}

diff --git a/internal/gomegahandler/dothandler.go b/internal/gomegahandler/dothandler.go
@@ -0,0 +1,99 @@
+package gomegahandler
+
+import (
+	"go/ast"
+
+	"golang.org/x/tools/go/analysis"
+
+	"github.com/nunnatsa/ginkgolinter/internal/gomegainfo"
+)
+
+// dotHandler is used when importing gomega with dot; i.e.
+// import . "github.com/onsi/gomega"
+type dotHandler struct {
+	pass *analysis.Pass
+}
+
+// GetActualFuncName returns the name of the gomega function, e.g. `Expect`
+func (h dotHandler) GetActualFuncName(expr *ast.CallExpr) (string, bool) {
+	switch actualFunc := expr.Fun.(type) {
+	case *ast.Ident:
+		return actualFunc.Name, true
+	case *ast.SelectorExpr:
+		if h.isGomegaVar(actualFunc.X) {
+			return actualFunc.Sel.Name, true
+		}
+
+		if x, ok := actualFunc.X.(*ast.CallExpr); ok {
+			return h.GetActualFuncName(x)
+		}
+
+	case *ast.CallExpr:
+		return h.GetActualFuncName(actualFunc)
+	}
+	return "", false
+}
+
+// ReplaceFunction replaces the function with another one, for fix suggestions
+func (dotHandler) ReplaceFunction(caller *ast.CallExpr, newExpr *ast.Ident) {
+	switch f := caller.Fun.(type) {
+	case *ast.Ident:
+		caller.Fun = newExpr
+	case *ast.SelectorExpr:
+		f.Sel = newExpr
+	}
+}
+
+func (dotHandler) GetNewWrapperMatcher(name string, existing *ast.CallExpr) *ast.CallExpr {
+	return &ast.CallExpr{
+		Fun:  ast.NewIdent(name),
+		Args: []ast.Expr{existing},
+	}
+}
+
+func (h dotHandler) GetActualExpr(assertionFunc *ast.SelectorExpr) *ast.CallExpr {
+	actualExpr, ok := assertionFunc.X.(*ast.CallExpr)
+	if !ok {
+		return nil
+	}
+
+	switch fun := actualExpr.Fun.(type) {
+	case *ast.Ident:
+		return actualExpr
+	case *ast.SelectorExpr:
+		if gomegainfo.IsActualMethod(fun.Sel.Name) {
+			if h.isGomegaVar(fun.X) {
+				return actualExpr
+			}
+		} else {
+			return h.GetActualExpr(fun)
+		}
+	}
+	return nil
+}
+
+func (h dotHandler) GetActualExprClone(origFunc, funcClone *ast.SelectorExpr) *ast.CallExpr {
+	actualExpr, ok := funcClone.X.(*ast.CallExpr)
+	if !ok {
+		return nil
+	}
+
+	switch funClone := actualExpr.Fun.(type) {
+	case *ast.Ident:
+		return actualExpr
+	case *ast.SelectorExpr:
+		origFun := origFunc.X.(*ast.CallExpr).Fun.(*ast.SelectorExpr)
+		if gomegainfo.IsActualMethod(funClone.Sel.Name) {
+			if h.isGomegaVar(origFun.X) {
+				return actualExpr
+			}
+		} else {
+			return h.GetActualExprClone(origFun, funClone)
+		}
+	}
+	return nil
+}
+
+func (h dotHandler) isGomegaVar(x ast.Expr) bool {
+	return gomegainfo.IsGomegaVar(x, h.pass)
+}
diff --git a/internal/gomegahandler/handler.go b/internal/gomegahandler/handler.go
@@ -2,9 +2,6 @@ package gomegahandler
 
 import (
 	"go/ast"
-	gotypes "go/types"
-	"regexp"
-
 	"golang.org/x/tools/go/analysis"
 )
 
@@ -48,229 +45,3 @@ func GetGomegaHandler(file *ast.File, pass *analysis.Pass) Handler {
 
 	return nil // no gomega import; this file does not use gomega
 }
-
-// dotHandler is used when importing gomega with dot; i.e.
-// import . "github.com/onsi/gomega"
-type dotHandler struct {
-	pass *analysis.Pass
-}
-
-// GetActualFuncName returns the name of the gomega function, e.g. `Expect`
-func (h dotHandler) GetActualFuncName(expr *ast.CallExpr) (string, bool) {
-	switch actualFunc := expr.Fun.(type) {
-	case *ast.Ident:
-		return actualFunc.Name, true
-	case *ast.SelectorExpr:
-		if h.isGomegaVar(actualFunc.X) {
-			return actualFunc.Sel.Name, true
-		}
-
-		if x, ok := actualFunc.X.(*ast.CallExpr); ok {
-			return h.GetActualFuncName(x)
-		}
-
-	case *ast.CallExpr:
-		return h.GetActualFuncName(actualFunc)
-	}
-	return "", false
-}
-
-// ReplaceFunction replaces the function with another one, for fix suggestions
-func (dotHandler) ReplaceFunction(caller *ast.CallExpr, newExpr *ast.Ident) {
-	switch f := caller.Fun.(type) {
-	case *ast.Ident:
-		caller.Fun = newExpr
-	case *ast.SelectorExpr:
-		f.Sel = newExpr
-	}
-}
-
-func (dotHandler) GetNewWrapperMatcher(name string, existing *ast.CallExpr) *ast.CallExpr {
-	return &ast.CallExpr{
-		Fun:  ast.NewIdent(name),
-		Args: []ast.Expr{existing},
-	}
-}
-
-// nameHandler is used when importing gomega without name; i.e.
-// import "github.com/onsi/gomega"
-//
-// or with a custom name; e.g.
-// import customname "github.com/onsi/gomega"
-type nameHandler struct {
-	name string
-	pass *analysis.Pass
-}
-
-// GetActualFuncName returns the name of the gomega function, e.g. `Expect`
-func (g nameHandler) GetActualFuncName(expr *ast.CallExpr) (string, bool) {
-	selector, ok := expr.Fun.(*ast.SelectorExpr)
-	if !ok {
-		return "", false
-	}
-
-	switch x := selector.X.(type) {
-	case *ast.Ident:
-		if x.Name != g.name {
-			if !g.isGomegaVar(x) {
-				return "", false
-			}
-		}
-
-		return selector.Sel.Name, true
-
-	case *ast.CallExpr:
-		return g.GetActualFuncName(x)
-	}
-
-	return "", false
-}
-
-// ReplaceFunction replaces the function with another one, for fix suggestions
-func (nameHandler) ReplaceFunction(caller *ast.CallExpr, newExpr *ast.Ident) {
-	caller.Fun.(*ast.SelectorExpr).Sel = newExpr
-}
-
-func (g nameHandler) isGomegaVar(x ast.Expr) bool {
-	return isGomegaVar(x, g.pass)
-}
-
-var gomegaTypeRegex = regexp.MustCompile(`github\.com/onsi/gomega/(?:internal|types)\.Gomega`)
-
-func isGomegaVar(x ast.Expr, pass *analysis.Pass) bool {
-	if tx, ok := pass.TypesInfo.Types[x]; ok {
-		return IsGomegaType(tx.Type)
-	}
-
-	return false
-}
-
-func IsGomegaType(t gotypes.Type) bool {
-	var typeStr string
-	switch ttx := t.(type) {
-	case *gotypes.Pointer:
-		tp := ttx.Elem()
-		typeStr = tp.String()
-
-	case *gotypes.Named:
-		typeStr = ttx.String()
-
-	default:
-		return false
-	}
-
-	return gomegaTypeRegex.MatchString(typeStr)
-}
-
-func (h dotHandler) GetActualExpr(assertionFunc *ast.SelectorExpr) *ast.CallExpr {
-	actualExpr, ok := assertionFunc.X.(*ast.CallExpr)
-	if !ok {
-		return nil
-	}
-
-	switch fun := actualExpr.Fun.(type) {
-	case *ast.Ident:
-		return actualExpr
-	case *ast.SelectorExpr:
-		if isHelperMethods(fun.Sel.Name) {
-			return h.GetActualExpr(fun)
-		}
-		if h.isGomegaVar(fun.X) {
-			return actualExpr
-		}
-	}
-	return nil
-}
-
-func (h dotHandler) GetActualExprClone(origFunc, funcClone *ast.SelectorExpr) *ast.CallExpr {
-	actualExpr, ok := funcClone.X.(*ast.CallExpr)
-	if !ok {
-		return nil
-	}
-
-	switch funClone := actualExpr.Fun.(type) {
-	case *ast.Ident:
-		return actualExpr
-	case *ast.SelectorExpr:
-		origFun := origFunc.X.(*ast.CallExpr).Fun.(*ast.SelectorExpr)
-		if isHelperMethods(funClone.Sel.Name) {
-			return h.GetActualExprClone(origFun, funClone)
-		}
-		if h.isGomegaVar(origFun.X) {
-			return actualExpr
-		}
-	}
-	return nil
-}
-
-func (h dotHandler) isGomegaVar(x ast.Expr) bool {
-	return isGomegaVar(x, h.pass)
-}
-
-func (g nameHandler) GetActualExpr(assertionFunc *ast.SelectorExpr) *ast.CallExpr {
-	actualExpr, ok := assertionFunc.X.(*ast.CallExpr)
-	if !ok {
-		return nil
-	}
-
-	switch fun := actualExpr.Fun.(type) {
-	case *ast.Ident:
-		return actualExpr
-	case *ast.SelectorExpr:
-		if x, ok := fun.X.(*ast.Ident); ok && x.Name == g.name {
-			return actualExpr
-		}
-		if isHelperMethods(fun.Sel.Name) {
-			return g.GetActualExpr(fun)
-		}
-
-		if g.isGomegaVar(fun.X) {
-			return actualExpr
-		}
-	}
-	return nil
-}
-
-func (g nameHandler) GetActualExprClone(origFunc, funcClone *ast.SelectorExpr) *ast.CallExpr {
-	actualExpr, ok := funcClone.X.(*ast.CallExpr)
-	if !ok {
-		return nil
-	}
-
-	switch funClone := actualExpr.Fun.(type) {
-	case *ast.Ident:
-		return actualExpr
-	case *ast.SelectorExpr:
-		if x, ok := funClone.X.(*ast.Ident); ok && x.Name == g.name {
-			return actualExpr
-		}
-		origFun := origFunc.X.(*ast.CallExpr).Fun.(*ast.SelectorExpr)
-		if isHelperMethods(funClone.Sel.Name) {
-			return g.GetActualExprClone(origFun, funClone)
-		}
-
-		if g.isGomegaVar(origFun.X) {
-			return actualExpr
-		}
-	}
-	return nil
-}
-
-func (g nameHandler) GetNewWrapperMatcher(name string, existing *ast.CallExpr) *ast.CallExpr {
-	return &ast.CallExpr{
-		Fun: &ast.SelectorExpr{
-			X:   ast.NewIdent(g.name),
-			Sel: ast.NewIdent(name),
-		},
-		Args: []ast.Expr{existing},
-	}
-}
-
-func isHelperMethods(funcName string) bool {
-	switch funcName {
-	case "WithOffset", "WithTimeout", "WithPolling", "Within", "ProbeEvery", "WithContext", "WithArguments", "MustPassRepeatedly":
-		return true
-	}
-
-	return false
-}