github.com/cockroachdb/tools@v0.0.0-20230222021103-a6d27438930d/go/callgraph/cha/cha_test.go (about)

     1  // Copyright 2014 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // No testdata on Android.
     6  
     7  //go:build !android
     8  // +build !android
     9  
    10  package cha_test
    11  
    12  import (
    13  	"bytes"
    14  	"fmt"
    15  	"go/ast"
    16  	"go/parser"
    17  	"go/token"
    18  	"go/types"
    19  	"io/ioutil"
    20  	"sort"
    21  	"strings"
    22  	"testing"
    23  
    24  	"golang.org/x/tools/go/callgraph"
    25  	"golang.org/x/tools/go/callgraph/cha"
    26  	"golang.org/x/tools/go/loader"
    27  	"golang.org/x/tools/go/ssa"
    28  	"golang.org/x/tools/go/ssa/ssautil"
    29  	"golang.org/x/tools/internal/typeparams"
    30  )
    31  
    32  var inputs = []string{
    33  	"testdata/func.go",
    34  	"testdata/iface.go",
    35  	"testdata/recv.go",
    36  	"testdata/issue23925.go",
    37  }
    38  
    39  func expectation(f *ast.File) (string, token.Pos) {
    40  	for _, c := range f.Comments {
    41  		text := strings.TrimSpace(c.Text())
    42  		if t := strings.TrimPrefix(text, "WANT:\n"); t != text {
    43  			return t, c.Pos()
    44  		}
    45  	}
    46  	return "", token.NoPos
    47  }
    48  
    49  // TestCHA runs CHA on each file in inputs, prints the dynamic edges of
    50  // the call graph, and compares it with the golden results embedded in
    51  // the WANT comment at the end of the file.
    52  func TestCHA(t *testing.T) {
    53  	for _, filename := range inputs {
    54  		prog, f, mainPkg, err := loadProgInfo(filename, ssa.InstantiateGenerics)
    55  		if err != nil {
    56  			t.Error(err)
    57  			continue
    58  		}
    59  
    60  		want, pos := expectation(f)
    61  		if pos == token.NoPos {
    62  			t.Error(fmt.Errorf("No WANT: comment in %s", filename))
    63  			continue
    64  		}
    65  
    66  		cg := cha.CallGraph(prog)
    67  
    68  		if got := printGraph(cg, mainPkg.Pkg, "dynamic", "Dynamic calls"); got != want {
    69  			t.Errorf("%s: got:\n%s\nwant:\n%s",
    70  				prog.Fset.Position(pos), got, want)
    71  		}
    72  	}
    73  }
    74  
    75  // TestCHAGenerics is TestCHA tailored for testing generics,
    76  func TestCHAGenerics(t *testing.T) {
    77  	if !typeparams.Enabled {
    78  		t.Skip("TestCHAGenerics requires type parameters")
    79  	}
    80  
    81  	filename := "testdata/generics.go"
    82  	prog, f, mainPkg, err := loadProgInfo(filename, ssa.InstantiateGenerics)
    83  	if err != nil {
    84  		t.Fatal(err)
    85  	}
    86  
    87  	want, pos := expectation(f)
    88  	if pos == token.NoPos {
    89  		t.Fatal(fmt.Errorf("No WANT: comment in %s", filename))
    90  	}
    91  
    92  	cg := cha.CallGraph(prog)
    93  
    94  	if got := printGraph(cg, mainPkg.Pkg, "", "All calls"); got != want {
    95  		t.Errorf("%s: got:\n%s\nwant:\n%s",
    96  			prog.Fset.Position(pos), got, want)
    97  	}
    98  }
    99  
   100  func loadProgInfo(filename string, mode ssa.BuilderMode) (*ssa.Program, *ast.File, *ssa.Package, error) {
   101  	content, err := ioutil.ReadFile(filename)
   102  	if err != nil {
   103  		return nil, nil, nil, fmt.Errorf("couldn't read file '%s': %s", filename, err)
   104  	}
   105  
   106  	conf := loader.Config{
   107  		ParserMode: parser.ParseComments,
   108  	}
   109  	f, err := conf.ParseFile(filename, content)
   110  	if err != nil {
   111  		return nil, nil, nil, err
   112  	}
   113  
   114  	conf.CreateFromFiles("main", f)
   115  	iprog, err := conf.Load()
   116  	if err != nil {
   117  		return nil, nil, nil, err
   118  	}
   119  
   120  	prog := ssautil.CreateProgram(iprog, mode)
   121  	prog.Build()
   122  
   123  	return prog, f, prog.Package(iprog.Created[0].Pkg), nil
   124  }
   125  
   126  // printGraph returns a string representation of cg involving only edges
   127  // whose description contains edgeMatch. The string representation is
   128  // prefixed with a desc line.
   129  func printGraph(cg *callgraph.Graph, from *types.Package, edgeMatch string, desc string) string {
   130  	var edges []string
   131  	callgraph.GraphVisitEdges(cg, func(e *callgraph.Edge) error {
   132  		if strings.Contains(e.Description(), edgeMatch) {
   133  			edges = append(edges, fmt.Sprintf("%s --> %s",
   134  				e.Caller.Func.RelString(from),
   135  				e.Callee.Func.RelString(from)))
   136  		}
   137  		return nil
   138  	})
   139  	sort.Strings(edges)
   140  
   141  	var buf bytes.Buffer
   142  	buf.WriteString(desc + "\n")
   143  	for _, edge := range edges {
   144  		fmt.Fprintf(&buf, "  %s\n", edge)
   145  	}
   146  	return strings.TrimSpace(buf.String())
   147  }