github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/other/cmd/rewrite/rewrite_test.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"go/ast"
     6  	"go/format"
     7  	"go/parser"
     8  	"go/token"
     9  	"go/types"
    10  	"testing"
    11  
    12  	"golang.org/x/tools/go/loader"
    13  )
    14  
    15  func TestGos(t *testing.T) {
    16  	src := `
    17  package main
    18  
    19  import "golang.org/x/net/context"
    20  
    21  func one(c context.Context) {}
    22  
    23  func two(c context.Context, a int) {}
    24  
    25  func three(c context.Context, a int, b int) {}
    26  
    27  func variadic(c context.Context, a ...int) {}
    28  
    29  func blank(_ context.Context) {}
    30  
    31  func nested(c context.Context) {
    32  	func(c context.Context) {}(c)
    33  }
    34  
    35  func main() {
    36  	go one(nil)
    37  
    38  	go two(nil, 0)
    39  
    40  	go three(nil, 0, 1)
    41  
    42  	go variadic(nil)
    43  
    44  	go variadic(nil, 0)
    45  
    46  	go variadic(nil, 0, 1)
    47  
    48  	go variadic(nil, []int{1, 2, 3}...)
    49  
    50  	go blank(nil)
    51  
    52  	go nested(nil)
    53  }
    54  `
    55  
    56  	expect := `package main
    57  
    58  import (
    59  	"golang.org/x/net/context"
    60  	__runtime "runtime"
    61  )
    62  
    63  func one(c context.Context) {}
    64  
    65  func two(c context.Context, a int) {}
    66  
    67  func three(c context.Context, a int, b int) {}
    68  
    69  func variadic(c context.Context, a ...int) {}
    70  
    71  func blank(_ context.Context) {}
    72  
    73  func nested(c context.Context) {
    74  	func(c context.Context) {}(c)
    75  }
    76  
    77  func main() {
    78  	go func(__f func(c context.
    79  		Context), arg0 context.Context) {
    80  		__runtime.SetLocal(arg0)
    81  		__f(arg0)
    82  	}(one, nil)
    83  	go func(__f func(c context.
    84  		Context, a int), arg0 context.
    85  		Context, arg1 int) {
    86  		__runtime.SetLocal(
    87  			arg0)
    88  		__f(arg0, arg1)
    89  	}(two, nil, 0)
    90  	go func(__f func(c context.
    91  		Context, a int, b int), arg0 context.
    92  		Context, arg1 int, arg2 int) {
    93  		__runtime.
    94  			SetLocal(arg0)
    95  		__f(arg0, arg1, arg2)
    96  	}(three,
    97  
    98  		nil, 0, 1)
    99  	go func(__f func(c context.
   100  		Context, a ...int), arg0 context.
   101  		Context, arg1 ...[]int) {
   102  		__runtime.
   103  			SetLocal(arg0)
   104  		__f(arg0, arg1...)
   105  	}(variadic, nil)
   106  	go func(__f func(c context.
   107  		Context, a ...int), arg0 context.
   108  		Context, arg1 ...[]int) {
   109  		__runtime.
   110  			SetLocal(arg0)
   111  		__f(arg0, arg1...)
   112  	}(variadic, nil,
   113  
   114  		0)
   115  	go func(__f func(c context.
   116  		Context, a ...int), arg0 context.
   117  		Context, arg1 ...[]int) {
   118  		__runtime.
   119  			SetLocal(arg0)
   120  		__f(arg0, arg1...)
   121  	}(variadic, nil,
   122  
   123  		0, 1)
   124  	go func(__f func(c context.
   125  		Context, a ...int), arg0 context.
   126  		Context, arg1 ...[]int) {
   127  		__runtime.
   128  			SetLocal(arg0)
   129  		__f(arg0, arg1...)
   130  	}(variadic, nil,
   131  
   132  		[]int{1, 2, 3})
   133  	go func(__f func(_ context.
   134  		Context), arg0 context.Context) {
   135  		__runtime.SetLocal(arg0)
   136  		__f(arg0)
   137  	}(blank, nil)
   138  	go func(__f func(c context.
   139  		Context), arg0 context.Context) {
   140  		__runtime.SetLocal(arg0)
   141  		__f(arg0)
   142  	}(nested, nil)
   143  
   144  }
   145  `
   146  
   147  	new := testHelper(t, src, rewriteGos)
   148  	if new != expect {
   149  		var idx int
   150  		for i, c := range []byte(new) {
   151  			if c != expect[i] {
   152  				idx = i
   153  				break
   154  			}
   155  		}
   156  		t.Errorf("unexpected output (see source for expected output) at character %v:\n%v", idx, new)
   157  	}
   158  }
   159  
   160  type rewriter func(fset *token.FileSet, info types.Info, qual types.Qualifier, f *ast.File) (bool, error)
   161  
   162  func testHelper(t *testing.T, src string, r rewriter) string {
   163  	c := loader.Config{
   164  		Fset:        token.NewFileSet(),
   165  		ParserMode:  parser.ParseComments | parser.DeclarationErrors,
   166  		AllowErrors: true,
   167  	}
   168  
   169  	f, err := c.ParseFile("test.go", src)
   170  	if err != nil {
   171  		t.Fatalf("unexpected error: %v", err)
   172  	}
   173  	c.CreateFromFiles("main", f)
   174  	p, err := c.Load()
   175  	if err != nil {
   176  		t.Fatalf("unexpected error: %v", err)
   177  	}
   178  	pi := p.Package("main")
   179  
   180  	_, err = r(c.Fset, pi.Info, qualifierForFile(pi.Pkg, f), f)
   181  	if err != nil {
   182  		t.Fatalf("unexpected error: %v", err)
   183  	}
   184  
   185  	var buf bytes.Buffer
   186  	err = format.Node(&buf, c.Fset, f)
   187  	if err != nil {
   188  		t.Fatalf("unexpected error: %v", err)
   189  	}
   190  	b, err := format.Source(buf.Bytes())
   191  	if err != nil {
   192  		t.Fatalf("unexpected error: %v", err)
   193  	}
   194  	return string(b)
   195  }