github.com/apache/beam/sdks/v2@v2.48.2/go/cmd/starcgen/starcgen_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package main
    17  
    18  import (
    19  	"bytes"
    20  	"go/ast"
    21  	"go/parser"
    22  	"go/token"
    23  	"strings"
    24  	"testing"
    25  )
    26  
    27  func TestGenerate(t *testing.T) {
    28  	tests := []struct {
    29  		name     string
    30  		pkg      string
    31  		files    []string
    32  		ids      []string
    33  		expected []string
    34  		excluded []string
    35  	}{
    36  		{name: "genAllSingleFile", files: []string{hello1}, pkg: "hello", ids: []string{},
    37  			expected: []string{"runtime.RegisterFunction(MyTitle)", "runtime.RegisterFunction(MyOtherDoFn)", "runtime.RegisterType(reflect.TypeOf((*foo)(nil)).Elem())", "funcMakerContext۰ContextStringГString", "funcMakerFooГString"},
    38  		},
    39  		{name: "genSpecificSingleFile", files: []string{hello1}, pkg: "hello", ids: []string{"MyTitle"},
    40  			expected: []string{"runtime.RegisterFunction(MyTitle)", "funcMakerContext۰ContextStringГString"},
    41  			excluded: []string{"MyOtherDoFn", "runtime.RegisterType(reflect.TypeOf((*foo)(nil)).Elem())", "funcMakerFooГString"},
    42  		},
    43  		{name: "genAllMultiFile", files: []string{hello1, hello2}, pkg: "hello", ids: []string{},
    44  			expected: []string{"runtime.RegisterFunction(MyTitle)", "runtime.RegisterFunction(MyOtherDoFn)", "runtime.RegisterFunction(anotherFn)", "runtime.RegisterType(reflect.TypeOf((*foo)(nil)).Elem())", "funcMakerContext۰ContextStringГString", "funcMakerFooГString", "funcMakerShimx۰EmitterГString", "funcMakerShimx۰EmitterГFoo"},
    45  		},
    46  		{name: "genSpecificMultiFile1", files: []string{hello1, hello2}, pkg: "hello", ids: []string{"MyTitle"},
    47  			expected: []string{"runtime.RegisterFunction(MyTitle)", "funcMakerContext۰ContextStringГString"},
    48  			excluded: []string{"MyOtherDoFn", "anotherFn", "runtime.RegisterType(reflect.TypeOf((*foo)(nil)).Elem())", "funcMakerFooГString", "funcMakerShimx۰EmitterГString", "funcMakerShimx۰EmitterГFoo"},
    49  		},
    50  		{name: "genSpecificMultiFile2", files: []string{hello1, hello2}, pkg: "hello", ids: []string{"anotherFn"},
    51  			expected: []string{"funcMakerShimx۰EmitterГString", "funcMakerShimx۰EmitterГString"},
    52  			excluded: []string{"MyOtherDoFn", "MyTitle", "runtime.RegisterType(reflect.TypeOf((*foo)(nil)).Elem())", "funcMakerFooГString"},
    53  		},
    54  	}
    55  	for _, test := range tests {
    56  		test := test
    57  		t.Run(test.name, func(t *testing.T) {
    58  			fset := token.NewFileSet()
    59  			var fs []*ast.File
    60  			for i, f := range test.files {
    61  				n, err := parser.ParseFile(fset, "", f, 0)
    62  				if err != nil {
    63  					t.Fatalf("couldn't parse test.files[%d]: %v", i, err)
    64  				}
    65  				fs = append(fs, n)
    66  			}
    67  			var b bytes.Buffer
    68  			if err := Generate(&b, test.name+".go", test.pkg, test.ids, fset, fs); err != nil {
    69  				t.Fatal(err)
    70  			}
    71  			s := string(b.Bytes())
    72  			for _, i := range test.expected {
    73  				if !strings.Contains(s, i) {
    74  					t.Errorf("expected %q in generated file", i)
    75  				}
    76  			}
    77  			for _, i := range test.excluded {
    78  				if strings.Contains(s, i) {
    79  					t.Errorf("found %q in generated file", i)
    80  				}
    81  			}
    82  			t.Log(s)
    83  		})
    84  	}
    85  }
    86  
    87  const hello1 = `
    88  package hello
    89  
    90  import (
    91  	"context"
    92  	"strings"
    93  )
    94  
    95  func MyTitle(ctx context.Context, v string) string {
    96  	return strings.Title(v)
    97  }
    98  
    99  type foo struct{}
   100  
   101  func MyOtherDoFn(v foo) string {
   102  	return "constant"
   103  }
   104  `
   105  
   106  const hello2 = `
   107  package hello
   108  
   109  import (
   110  	"context"
   111  	"strings"
   112  
   113  	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/shimx"
   114  )
   115  
   116  func anotherFn(v shimx.Emitter) string {
   117  	return v.Name
   118  }
   119  
   120  func fooFn(v shimx.Emitter) foo {
   121  	return foo{}
   122  }
   123  `