client: Allow creating a GoBGPClient using a user specific GRPC connection and client

This allows wrapping the dialer, implementing your own wrapped
GobgpApiClient or dialing with a context for cancelation on unreliable
endpoints.

2 files changed, 20 insertions, 3 deletions

diff --git a/client/client.go b/client/client.go
index b85360c2..7fa9d0ec 100644
--- a/client/client.go
+++ b/client/client.go
@@ -39,14 +39,23 @@ func defaultGRPCOptions() []grpc.DialOption {
 	return []grpc.DialOption{grpc.WithTimeout(time.Second), grpc.WithBlock(), grpc.WithInsecure()}
 }
 
-func NewGoBGPClient(target string, opts ...grpc.DialOption) (*GoBGPClient, error) {
+// New returns a new GoBGPClient using the given target and options for dialing
+// to the grpc server. If an error occurs during dialing it will be returned and
+// GoBGPClient will be nil.
+func New(target string, opts ...grpc.DialOption) (*GoBGPClient, error) {
+	return NewWith(context.Background(), target, opts...)
+}
+
+// NewWith is like New, but uses the given ctx to cancel or expire the current
+// attempt to connect if it becomes Done before the connection succeeds.
+func NewWith(ctx context.Context, target string, opts ...grpc.DialOption) (*GoBGPClient, error) {
 	if target == "" {
 		target = ":50051"
 	}
 	if len(opts) == 0 {
 		opts = defaultGRPCOptions()
 	}
-	conn, err := grpc.Dial(target, opts...)
+	conn, err := grpc.DialContext(ctx, target, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -54,6 +63,14 @@ func NewGoBGPClient(target string, opts ...grpc.DialOption) (*GoBGPClient, error
 	return &GoBGPClient{conn: conn, cli: cli}, nil
 }
 
+// NewFrom returns a new GoBGPClient, using the given conn and cli for the
+// underlying connection. The given grpc.ClientConn connection is expected to be
+// initialized and paired with the api client. See New to have the connection
+// dialed for you.
+func NewFrom(conn *grpc.ClientConn, cli api.GobgpApiClient) *GoBGPClient {
+	return &GoBGPClient{conn: conn, cli: cli}
+}
+
 func (cli *GoBGPClient) Close() error {
 	return cli.conn.Close()
 }
diff --git a/client/client_test.go b/client/client_test.go
index 9825f91b..3f07a6ad 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -30,7 +30,7 @@ func TestGetNeighbor(test *testing.T) {
 	go s.Serve()
 	g := api.NewGrpcServer(s, ":50051")
 	go g.Serve()
-	cli, err := NewGoBGPClient("")
+	cli, err := New("")
 	err = cli.StartServer(&config.Global{
 		Config: config.GlobalConfig{
 			As:       1,