From 96bfd877fbc5869c0a4d87de4daeb1f76aaca79d Mon Sep 17 00:00:00 2001
From: Cody Oss <6331106+codyoss@users.noreply.github.com>
Date: Mon, 11 Jan 2021 10:12:02 -0700
Subject: [PATCH] feat(transport/bytestream): Add Close method for shutdown
 (#787)

This pacakge is still labeled expirmental. At some point we should
decide if we want to keep supporting it or deprecate and remove.

Fixes: #775
---
 transport/bytestream/client.go      | 8 ++++++++
 transport/bytestream/client_test.go | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/transport/bytestream/client.go b/transport/bytestream/client.go
index 0c4d7c023b0..eabd268179f 100644
--- a/transport/bytestream/client.go
+++ b/transport/bytestream/client.go
@@ -33,6 +33,7 @@ const (
 type Client struct {
 	client  pb.ByteStreamClient
 	options []grpc.CallOption
+	conn    *grpc.ClientConn
 }
 
 // NewClient creates a new bytestream.Client.
@@ -40,6 +41,7 @@ func NewClient(cc *grpc.ClientConn, options ...grpc.CallOption) *Client {
 	return &Client{
 		client:  pb.NewByteStreamClient(cc),
 		options: options,
+		conn:    cc,
 	}
 }
 
@@ -228,3 +230,9 @@ func (c *Client) NewWriter(ctx context.Context, resourceName string) (*Writer, e
 		resourceName: resourceName,
 	}, nil
 }
+
+// Close closes the connection to the API service. The user should invoke this when
+// the client is no longer required.
+func (c *Client) Close() {
+	c.conn.Close()
+}
diff --git a/transport/bytestream/client_test.go b/transport/bytestream/client_test.go
index ff0640d2909..2a733e15c34 100644
--- a/transport/bytestream/client_test.go
+++ b/transport/bytestream/client_test.go
@@ -307,7 +307,6 @@ func TestClientWrite_WriteFails(t *testing.T) {
 }
 
 func TestClientWrite_CloseAndRecvFails(t *testing.T) {
-	t.Skip("https://github.com/googleapis/google-api-go-client/issues/775")
 	setup := newTestSetup("")
 	w, err := setup.client.NewWriter(setup.ctx, "CloseAndRecvFails")
 	if err != nil {
@@ -425,5 +424,6 @@ func newTestSetup(input string) *TestSetup {
 }
 
 func (testSetup *TestSetup) Close() {
+	testSetup.client.Close()
 	testSetup.rpcTest.Close()
 }