internal/tdtest: make function detection more robust

Instead of detecting the package (which was already brittle),
test the first argument of the closure of the second argument of
Run. This allows other packages to wrap the tdtest.Run function,
as long as they keep the same signature.

This is necessary to handle errors in cuetest.Run.

Signed-off-by: Marcel van Lohuizen <[email protected]>
Change-Id: If8dea69244fec9111916df667b0a8c09dc85fa4d
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1167818
TryBot-Result: CUEcueckoo <[email protected]>
Unity-Result: CUE porcuepine <[email protected]>
Reviewed-by: Daniel Martí <[email protected]>

internal/tdtest/update.go

@@ -204,10 +204,7 @@ func findFileAndPackage(path string, pkgs []*packages.Package) (*ast.File, *pack
 	return nil, nil
 }
 
-const (
-	typeT       = "*cuelang.org/go/internal/tdtest.T"
-	tdtestParen = `("cuelang.org/go/internal/tdtest")`
-)
+const typeT = "*cuelang.org/go/internal/tdtest.T"
 
 // findCalls finds all call expressions within a given block for functions
 // or methods defined within the tdtest package.
@@ -229,10 +226,19 @@ func (i *info) findCalls(block *ast.BlockStmt, names ...string) []*callInfo {
 		info := i.testPkg.TypesInfo
 		for _, name := range names {
 			if sel.Sel.Name == name {
-				if info.TypeOf(sel.X).String() == typeT {
-				} else if ident, ok := sel.X.(*ast.Ident); !ok {
-					return true // Run method.
-				} else if id, ok := info.Uses[ident].(*types.PkgName); ok && strings.Contains(id.String(), tdtestParen) {
+				receiver := info.TypeOf(sel.X).String()
+				if receiver == typeT {
+					// Method.
+				} else if len(c.Args) == 3 {
+					// Run function.
+					fn := c.Args[2].(*ast.FuncLit)
+					if len(fn.Type.Params.List) != 2 {
+						return true
+					}
+					argType := info.TypeOf(fn.Type.Params.List[0].Type).String()
+					if argType != typeT {
+						return true
+					}
 				} else {
 					return true
 				}