Merge pull request #198 from askuy/feature/bufnetclient

support bufnet client
gotomicro · Oct 15, 2021 · facb4ba · facb4ba
2 parents d9870ee + 14e6d43
commit facb4ba
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/client/egrpc/option.go b/client/egrpc/option.go
@@ -1,9 +1,12 @@
 package egrpc
 
 import (
+	"context"
+	"net"
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/test/bufconn"
 )
 
 // WithAddr setting grpc server address
@@ -78,3 +81,10 @@ func WithEnableAccessInterceptorRes(enableAccessInterceptorRes bool) Option {
 		c.config.EnableAccessInterceptorRes = enableAccessInterceptorRes
 	}
 }
+
+// WithBufnetServerListener 写入bufnet listener
+func WithBufnetServerListener(svc net.Listener) Option {
+	return WithDialOption(grpc.WithContextDialer(func(ctx context.Context, s string) (net.Conn, error) {
+		return svc.(*bufconn.Listener).Dial()
+	}))
+}
diff --git a/cmd/protoc-gen-go-test/template.go b/cmd/protoc-gen-go-test/template.go
@@ -19,18 +19,12 @@ var svc *egrpc.Component
 func init() {
 	// initial your configuration、DB instance、Redis instance、svc instance here.
 }
-
-// bufDialer generated by protoc-gen-go-test, you should not edit it.  
-// @Override=true
-func bufDialer(context.Context, string) (net.Conn, error) {
-	return svc.Listener().(*bufconn.Listener).Dial()
-}
 `
 	stream2streamTpl = `
 // Test{{.Name}} generated by protoc-gen-go-test, you can fill test logic by yourself
 // @Override=true
 func Test{{.Name}}(t *testing.T) {
-	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithDialOption(grpc.WithContextDialer(bufDialer))).ClientConn)
+	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithBufnetServerListener(svc.Listener())).ClientConn)
     ctx := context.Background()
 	tests := []struct {
 		name string
@@ -62,7 +56,7 @@ func Test{{.Name}}(t *testing.T) {
 // Test{{.Name}} generated by protoc-gen-go-test, you can fill test logic by yourself.
 // @Override=true
 func Test{{.Name}}(t *testing.T) {
-	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithDialOption(grpc.WithContextDialer(bufDialer))).ClientConn)
+	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithBufnetServerListener(svc.Listener())).ClientConn)
     ctx := context.Background()
 	tests := []struct {
 		name string
@@ -90,7 +84,7 @@ func Test{{.Name}}(t *testing.T) {
 // Test{{.Name}} generated by protoc-gen-go-test, you can fill test logic by yourself.
 // @Override=true
 func Test{{.Name}}(t *testing.T) {
-	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithDialOption(grpc.WithContextDialer(bufDialer))).ClientConn)
+	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithBufnetServerListener(svc.Listener())).ClientConn)
     ctx := context.Background()
 
     stream, err := cli.{{.Name}}(ctx)
@@ -123,7 +117,7 @@ func Test{{.Name}}(t *testing.T) {
 // Test{{.Name}} generated by protoc-gen-go-test, you can fill test logic by yourself.
 // @Override=true
 func Test{{.Name}}(t *testing.T) {
-	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithDialOption(grpc.WithContextDialer(bufDialer))).ClientConn)
+	cli := {{.Package}}.New{{.Service.Name}}Client(cegrpc.DefaultContainer().Build(cegrpc.WithBufnetServerListener(svc.Listener())).ClientConn)
     ctx := context.Background()
 	tests := []struct {
 		name string