test/gracefulstop: use stubserver instead of testservice implementation

test/gracefulstop_test.go

@@ -22,7 +22,6 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"sync"
 	"testing"
 	"time"
 
@@ -110,36 +109,29 @@ func (s) TestGracefulStop(t *testing.T) {
 		closeCalled:  make(chan struct{}),
 		allowCloseCh: make(chan struct{}),
 	}
-	d := func(ctx context.Context, _ string) (net.Conn, error) { return dlis.Dial(ctx) }
 
 	ss := &stubserver.StubServer{
+		Listener: dlis,
 		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
 			_, err := stream.Recv()
 			if err != nil {
 				return err
 			}
 			return stream.Send(&testpb.StreamingOutputCallResponse{})
 		},
+		S: grpc.NewServer(),
 	}
-	s := grpc.NewServer()
-	testgrpc.RegisterTestServiceServer(s, ss)
+	stubserver.StartTestService(t, ss)
 
 	// 1. Start Server
-	wg := sync.WaitGroup{}
-	wg.Add(1)
+	gracefulStopDone := make(chan struct{})
+	<-dlis.acceptCalled
 	go func() {
-		s.Serve(dlis)
-		wg.Done()
+		ss.S.GracefulStop()
+		close(gracefulStopDone)
 	}()
-
 	// 2. GracefulStop() Server after listener's Accept is called, but don't
 	//    allow Accept() to exit when Close() is called on it.
-	<-dlis.acceptCalled
-	wg.Add(1)
-	go func() {
-		s.GracefulStop()
-		wg.Done()
-	}()
 
 	// 3. Create a new connection to the server after listener.Close() is called.
 	//    Server should close this connection immediately, before handshaking.
@@ -150,7 +142,8 @@ func (s) TestGracefulStop(t *testing.T) {
 	// even though GracefulStop has closed the listener.
 	ctx, dialCancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer dialCancel()
-	cc, err := grpc.DialContext(ctx, "", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(d))
+	dialer := func(ctx context.Context, _ string) (net.Conn, error) { return dlis.Dial(ctx) }
+	cc, err := grpc.DialContext(ctx, "", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(dialer))
 	if err != nil {
 		t.Fatalf("grpc.DialContext(_, %q, _) = %v", lis.Addr().String(), err)
 	}
@@ -165,7 +158,7 @@ func (s) TestGracefulStop(t *testing.T) {
 		t.Fatalf("FullDuplexCall= _, %v; want _, <status code Unavailable>", err)
 	}
 	cancel()
-	wg.Wait()
+	<-gracefulStopDone
 }
 
 // TestGracefulStopClosesConnAfterLastStream ensures that a server closes the