From cdd3331e3e7cbd08ce6d2645ac2ee6ae1c05825b Mon Sep 17 00:00:00 2001
From: Abhilash Gnan <abhilashgnan@gmail.com>
Date: Thu, 14 Jan 2021 23:05:11 +0100
Subject: [PATCH] Add TestClose

Signed-off-by: Abhilash Gnan <abhilashgnan@gmail.com>
---
 cmux_test.go    | 33 +++++++++++++++++++++++++++++----
 example_test.go |  5 ++++-
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/cmux_test.go b/cmux_test.go
index 6ef66dc..7b3cfb7 100644
--- a/cmux_test.go
+++ b/cmux_test.go
@@ -128,7 +128,7 @@ func runTestHTTPServer(errCh chan<- error, l net.Listener) {
 			mu.Unlock()
 		},
 	}
-	if err := s.Serve(l); err != ErrListenerClosed {
+	if err := s.Serve(l); err != ErrListenerClosed && err != ErrServerClosed {
 		errCh <- err
 	}
 }
@@ -218,7 +218,7 @@ func runTestRPCServer(errCh chan<- error, l net.Listener) {
 	for {
 		c, err := l.Accept()
 		if err != nil {
-			if err != ErrListenerClosed {
+			if err != ErrListenerClosed && err != ErrServerClosed {
 				errCh <- err
 			}
 			return
@@ -651,7 +651,7 @@ func TestMultipleMatchers(t *testing.T) {
 	runTestHTTP1Client(t, l.Addr())
 }
 
-func TestClose(t *testing.T) {
+func TestListenerClose(t *testing.T) {
 	defer leakCheck(t)()
 	errCh := make(chan error)
 	defer func() {
@@ -685,7 +685,7 @@ func TestClose(t *testing.T) {
 
 	// Second connection either goes through or it is closed.
 	if _, err := anyl.Accept(); err != nil {
-		if err != ErrListenerClosed {
+		if err != ErrListenerClosed && err != ErrServerClosed {
 			t.Fatal(err)
 		}
 		// The error is either io.ErrClosedPipe or net.OpError wrapping
@@ -696,6 +696,31 @@ func TestClose(t *testing.T) {
 	}
 }
 
+func TestClose(t *testing.T) {
+	defer leakCheck(t)()
+	errCh := make(chan error)
+	defer func() {
+		select {
+		case err := <-errCh:
+			t.Fatal(err)
+		default:
+		}
+	}()
+	l, cleanup := testListener(t)
+	defer cleanup()
+
+	muxl := New(l)
+	anyl := muxl.Match(Any())
+
+	go safeServe(errCh, muxl)
+
+	muxl.Close()
+
+	if _, err := anyl.Accept(); err != ErrServerClosed {
+		t.Fatal(err)
+	}
+}
+
 // Cribbed from google.golang.org/grpc/test/end2end_test.go.
 
 // interestingGoroutines returns all goroutines we care about for the purpose
diff --git a/example_test.go b/example_test.go
index 7144f5e..c275aa7 100644
--- a/example_test.go
+++ b/example_test.go
@@ -29,6 +29,7 @@ import (
 	"golang.org/x/net/websocket"
 
 	"github.com/soheilhy/cmux"
+	"google.golang.org/grpc/examples/helloworld/helloworld"
 	grpchello "google.golang.org/grpc/examples/helloworld/helloworld"
 )
 
@@ -86,7 +87,9 @@ func serveRPC(l net.Listener) {
 	}
 }
 
-type grpcServer struct{}
+type grpcServer struct {
+	helloworld.UnimplementedGreeterServer
+}
 
 func (s *grpcServer) SayHello(ctx context.Context, in *grpchello.HelloRequest) (
 	*grpchello.HelloReply, error) {