From db2122860d22f63c25ff5efc29621503119bb5dc Mon Sep 17 00:00:00 2001
From: letian <letian0805@sina.com>
Date: Fri, 17 Sep 2021 04:19:49 +0800
Subject: [PATCH] add endpoint option in grpc server (#1488)

---
 transport/grpc/server.go      | 10 ++++++++++
 transport/grpc/server_test.go |  9 +++++++++
 2 files changed, 19 insertions(+)

diff --git a/transport/grpc/server.go b/transport/grpc/server.go
index 748302c24..31efebbf2 100644
--- a/transport/grpc/server.go
+++ b/transport/grpc/server.go
@@ -89,6 +89,13 @@ func Options(opts ...grpc.ServerOption) ServerOption {
 	}
 }
 
+// Endpoint with server endpoint.
+func Endpoint(endpoint *url.URL) ServerOption {
+	return func(o *Server) {
+		o.endpoint = endpoint
+	}
+}
+
 // Server is a gRPC server wrapper.
 type Server struct {
 	*grpc.Server
@@ -151,6 +158,9 @@ func NewServer(opts ...ServerOption) *Server {
 //   grpc://127.0.0.1:9000?isSecure=false
 func (s *Server) Endpoint() (*url.URL, error) {
 	s.once.Do(func() {
+		if s.endpoint != nil {
+			return
+		}
 		lis, err := net.Listen(s.network, s.address)
 		if err != nil {
 			s.err = err
diff --git a/transport/grpc/server_test.go b/transport/grpc/server_test.go
index 3b9f1ede1..10a3df25d 100644
--- a/transport/grpc/server_test.go
+++ b/transport/grpc/server_test.go
@@ -82,6 +82,15 @@ func TestMiddleware(t *testing.T) {
 	assert.Equal(t, v, o.middleware)
 }
 
+func TestEndpoint(t *testing.T) {
+	u, err := url.Parse("grpc://hello/world")
+	assert.NoError(t, err)
+	o := &Server{}
+	Endpoint(u)(o)
+	assert.Equal(t, "hello", o.endpoint.Host)
+	assert.Equal(t, "grpc", o.endpoint.Scheme)
+}
+
 type mockLogger struct {
 	level log.Level
 	key   string