Run: ensure that all transport stop on exit

server.go

@@ -40,6 +40,8 @@ func NewServer(t ...Transport) *Server {
 
 // Run starts the server.
 func (s *Server) Run(ctx context.Context) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 	stop := make(chan error)
 	exit := make(chan error)
 	ctx = s.addLoggerToContext(ctx, nil)

server_test.go

@@ -132,11 +132,28 @@ func (*failingTransport) RegisterEndpoints(m endpoint.Middleware) error { return
 func (*failingTransport) Start(ctx context.Context) error               { return errors.New("unable to start") }
 func (*failingTransport) Shutdown(ctx context.Context) error            { return nil }
 
+type workingTransport struct {
+	running chan struct{}
+	Transport
+}
+
+func (*workingTransport) RegisterEndpoints(m endpoint.Middleware) error { return nil }
+func (t *workingTransport) Start(ctx context.Context) error {
+	t.running = make(chan struct{})
+	defer func() {
+		close(t.running)
+	}()
+	<-ctx.Done()
+	return nil
+}
+func (*workingTransport) Shutdown(ctx context.Context) error { return nil }
+
 func TestStartError(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.TODO(), time.Second)
+	ctx, cancel := context.WithTimeout(context.TODO(), 5*time.Second)
 	defer cancel()
 	exitError := make(chan error)
-	srv := NewServer(&failingTransport{})
+	wt := &workingTransport{}
+	srv := NewServer(&failingTransport{}, wt)
 	go func() {
 		exitError <- srv.Run(ctx)
 	}()
@@ -149,4 +166,9 @@ func TestStartError(t *testing.T) {
 			t.Errorf("Server.Run returned an invalid error : %v", err)
 		}
 	}
+	select {
+	case <-ctx.Done():
+		t.Error("Alternate transport has not stopped after 1sec")
+	case <-wt.running:
+	}
 }