From 3fcbd3737ec6baff505795417e48f162a7a3183c Mon Sep 17 00:00:00 2001
From: Michal Witkowski <michal@improbable.io>
Date: Mon, 27 Mar 2017 20:19:13 +0100
Subject: [PATCH] fixup closing conns

---
 proxy/handler_test.go | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/proxy/handler_test.go b/proxy/handler_test.go
index 4b3b047..1c8e3cf 100644
--- a/proxy/handler_test.go
+++ b/proxy/handler_test.go
@@ -99,10 +99,11 @@ func (s *assertingService) PingStream(stream pb.TestService_PingStreamServer) er
 type ProxyHappySuite struct {
 	suite.Suite
 
-	serverListener net.Listener
-	server         *grpc.Server
-	proxyListener  net.Listener
-	proxy          *grpc.Server
+	serverListener   net.Listener
+	server           *grpc.Server
+	proxyListener    net.Listener
+	proxy            *grpc.Server
+	serverClientConn *grpc.ClientConn
 
 	client     *grpc.ClientConn
 	testClient pb.TestServiceClient
@@ -201,7 +202,7 @@ func (s *ProxyHappySuite) SetupSuite() {
 	pb.RegisterTestServiceServer(s.server, &assertingService{t: s.T()})
 
 	// Setup of the proxy's Director.
-	proxyClientConn, err := grpc.Dial(s.serverListener.Addr().String(), grpc.WithInsecure(), grpc.WithCodec(proxy.Codec()))
+	s.serverClientConn, err = grpc.Dial(s.serverListener.Addr().String(), grpc.WithInsecure(), grpc.WithCodec(proxy.Codec()))
 	require.NoError(s.T(), err, "must not error on deferred client Dial")
 	director := func(ctx context.Context, fullName string) (*grpc.ClientConn, error) {
 		md, ok := metadata.FromContext(ctx)
@@ -210,7 +211,7 @@ func (s *ProxyHappySuite) SetupSuite() {
 				return nil, grpc.Errorf(codes.PermissionDenied, "testing rejection")
 			}
 		}
-		return proxyClientConn, nil
+		return s.serverClientConn, nil
 	}
 	s.proxy = grpc.NewServer(
 		grpc.CustomCodec(proxy.Codec()),
@@ -237,6 +238,14 @@ func (s *ProxyHappySuite) SetupSuite() {
 }
 
 func (s *ProxyHappySuite) TearDownSuite() {
+	if s.client != nil {
+		s.client.Close()
+	}
+	if s.serverClientConn != nil {
+		s.serverClientConn.Close()
+	}
+	// Close all transports so the logs don't get spammy.
+	time.Sleep(10 * time.Millisecond)
 	if s.proxy != nil {
 		s.proxy.Stop()
 		s.proxyListener.Close()
@@ -245,9 +254,6 @@ func (s *ProxyHappySuite) TearDownSuite() {
 		s.server.Stop()
 		s.serverListener.Close()
 	}
-	if s.client != nil {
-		s.client.Close()
-	}
 }
 
 func TestProxyHappySuite(t *testing.T) {