call.go

Documentation: gotest.tools/v3/assert/cmd/gty-migrate-from-testify

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"go/ast"
     7  	"go/format"
     8  	"go/token"
     9  )
    10  
    11  // call wraps a testify/assert ast.CallExpr and exposes properties of the
    12  // expression to facilitate migrating the expression to a gotest.tools/v3/assert
    13  type call struct {
    14  	fileset *token.FileSet
    15  	expr    *ast.CallExpr
    16  	xIdent  *ast.Ident
    17  	selExpr *ast.SelectorExpr
    18  	// assert is Assert (if the testify package was require), or Check (if the
    19  	// testify package was assert).
    20  	assert string
    21  }
    22  
    23  func (c call) String() string {
    24  	buf := new(bytes.Buffer)
    25  	_ = format.Node(buf, token.NewFileSet(), c.expr)
    26  	return buf.String()
    27  }
    28  
    29  func (c call) StringWithFileInfo() string {
    30  	if c.fileset.File(c.expr.Pos()) == nil {
    31  		return fmt.Sprintf("%s at unknown file", c)
    32  	}
    33  	return fmt.Sprintf("%s at %s:%d", c,
    34  		relativePath(c.fileset.File(c.expr.Pos()).Name()),
    35  		c.fileset.Position(c.expr.Pos()).Line)
    36  }
    37  
    38  // testingT returns the first argument of the call, which is assumed to be a
    39  // *ast.Ident of *testing.T or a compatible interface.
    40  func (c call) testingT() ast.Expr {
    41  	if len(c.expr.Args) == 0 {
    42  		return nil
    43  	}
    44  	return c.expr.Args[0]
    45  }
    46  
    47  // extraArgs returns the arguments of the expression starting from index
    48  func (c call) extraArgs(index int) []ast.Expr {
    49  	if len(c.expr.Args) <= index {
    50  		return nil
    51  	}
    52  	return c.expr.Args[index:]
    53  }
    54  
    55  // args returns a range of arguments from the expression
    56  func (c call) args(from, to int) []ast.Expr {
    57  	return c.expr.Args[from:to]
    58  }
    59  
    60  // arg returns a single argument from the expression
    61  func (c call) arg(index int) ast.Expr {
    62  	return c.expr.Args[index]
    63  }
    64  
    65  // newCallFromCallExpr returns a new call build from a ast.CallExpr. Returns false
    66  // if the call expression does not have the expected ast nodes.
    67  func newCallFromCallExpr(callExpr *ast.CallExpr, migration migration) (call, bool) {
    68  	c := call{}
    69  	selector, ok := callExpr.Fun.(*ast.SelectorExpr)
    70  	if !ok {
    71  		return c, false
    72  	}
    73  	ident, ok := selector.X.(*ast.Ident)
    74  	if !ok {
    75  		return c, false
    76  	}
    77  
    78  	return call{
    79  		fileset: migration.fileset,
    80  		xIdent:  ident,
    81  		selExpr: selector,
    82  		expr:    callExpr,
    83  	}, true
    84  }
    85  
    86  // newTestifyCallFromNode returns a call that wraps a valid testify assertion.
    87  // Returns false if the call expression is not a testify assertion.
    88  func newTestifyCallFromNode(callExpr *ast.CallExpr, migration migration) (call, bool) {
    89  	tcall, ok := newCallFromCallExpr(callExpr, migration)
    90  	if !ok {
    91  		return tcall, false
    92  	}
    93  
    94  	testifyNewAssignStmt := testifyAssertionsAssignment(tcall, migration)
    95  	switch {
    96  	case testifyNewAssignStmt != nil:
    97  		return updateCallForTestifyNew(tcall, testifyNewAssignStmt, migration)
    98  	case isTestifyPkgCall(tcall, migration):
    99  		tcall.assert = migration.importNames.funcNameFromTestifyName(tcall.xIdent.Name)
   100  		return tcall, true
   101  	}
   102  	return tcall, false
   103  }
   104  
   105  // isTestifyPkgCall returns true if the call is a testify package-level assertion
   106  // (as apposed to an assertion method on the Assertions type)
   107  //
   108  // TODO: check if the xIdent.Obj.Decl is an import declaration instead of
   109  // assuming that a name matching the import name is always an import. Some code
   110  // may shadow import names, which could lead to incorrect results.
   111  func isTestifyPkgCall(tcall call, migration migration) bool {
   112  	return migration.importNames.matchesTestify(tcall.xIdent)
   113  }
   114  
   115  // testifyAssertionsAssignment returns an ast.AssignStmt if the call is a testify
   116  // call from an Assertions object returned from assert.New(t) (not a package level
   117  // assert). Otherwise returns nil.
   118  func testifyAssertionsAssignment(tcall call, migration migration) *ast.AssignStmt {
   119  	if tcall.xIdent.Obj == nil {
   120  		return nil
   121  	}
   122  
   123  	assignStmt, ok := tcall.xIdent.Obj.Decl.(*ast.AssignStmt)
   124  	if !ok {
   125  		return nil
   126  	}
   127  
   128  	if isAssignmentFromAssertNew(assignStmt, migration) {
   129  		return assignStmt
   130  	}
   131  	return nil
   132  }
   133  
   134  func updateCallForTestifyNew(
   135  	tcall call,
   136  	testifyNewAssignStmt *ast.AssignStmt,
   137  	migration migration,
   138  ) (call, bool) {
   139  	testifyNewCallExpr := callExprFromAssignment(testifyNewAssignStmt)
   140  	if testifyNewCallExpr == nil {
   141  		return tcall, false
   142  	}
   143  	testifyNewCall, ok := newCallFromCallExpr(testifyNewCallExpr, migration)
   144  	if !ok {
   145  		return tcall, false
   146  	}
   147  
   148  	tcall.assert = migration.importNames.funcNameFromTestifyName(testifyNewCall.xIdent.Name)
   149  	tcall.expr = addMissingTestingTArgToCallExpr(tcall.expr, testifyNewCall.testingT())
   150  	return tcall, true
   151  }
   152  
   153  // addMissingTestingTArgToCallExpr adds a testingT arg as the first arg of the
   154  // ast.CallExpr and returns a copy of the ast.CallExpr
   155  func addMissingTestingTArgToCallExpr(callExpr *ast.CallExpr, testingT ast.Expr) *ast.CallExpr {
   156  	return &ast.CallExpr{
   157  		Fun:  callExpr.Fun,
   158  		Args: append([]ast.Expr{removePos(testingT)}, callExpr.Args...),
   159  	}
   160  }
   161  
   162  func removePos(node ast.Expr) ast.Expr {
   163  	switch typed := node.(type) {
   164  	case *ast.Ident:
   165  		return &ast.Ident{Name: typed.Name}
   166  	}
   167  	return node
   168  }
   169  
   170  // TODO: use pkgInfo and walkForType instead?
   171  func isAssignmentFromAssertNew(assign *ast.AssignStmt, migration migration) bool {
   172  	callExpr := callExprFromAssignment(assign)
   173  	if callExpr == nil {
   174  		return false
   175  	}
   176  	tcall, ok := newCallFromCallExpr(callExpr, migration)
   177  	if !ok {
   178  		return false
   179  	}
   180  	if !migration.importNames.matchesTestify(tcall.xIdent) {
   181  		return false
   182  	}
   183  
   184  	if len(tcall.expr.Args) != 1 {
   185  		return false
   186  	}
   187  	return tcall.selExpr.Sel.Name == "New"
   188  }
   189  
   190  func callExprFromAssignment(assign *ast.AssignStmt) *ast.CallExpr {
   191  	if len(assign.Rhs) != 1 {
   192  		return nil
   193  	}
   194  
   195  	callExpr, ok := assign.Rhs[0].(*ast.CallExpr)
   196  	if !ok {
   197  		return nil
   198  	}
   199  	return callExpr
   200  }
   201
View as plain text