diff --git a/credentials/alts/alts_test.go b/credentials/alts/alts_test.go index 6c083a2fe..48f871a00 100644 --- a/credentials/alts/alts_test.go +++ b/credentials/alts/alts_test.go @@ -36,6 +36,7 @@ import ( altspb "google.golang.org/grpc/credentials/alts/internal/proto/grpc_gcp" "google.golang.org/grpc/credentials/alts/internal/testutil" "google.golang.org/grpc/internal/grpctest" + "google.golang.org/grpc/internal/stubserver" "google.golang.org/grpc/internal/testutils" testgrpc "google.golang.org/grpc/interop/grpc_testing" testpb "google.golang.org/grpc/interop/grpc_testing" @@ -323,7 +324,7 @@ func (s) TestFullHandshake(t *testing.T) { defer wait.Wait() stopHandshaker, handshakerAddress := startFakeHandshakerService(t, &wait) defer stopHandshaker() - stopServer, serverAddress := startServer(t, handshakerAddress, &wait) + stopServer, serverAddress := startServer(t, handshakerAddress) defer stopServer() // Ping the server, authenticating with ALTS. @@ -349,7 +350,7 @@ func (s) TestConcurrentHandshakes(t *testing.T) { defer wait.Wait() stopHandshaker, handshakerAddress := startFakeHandshakerService(t, &wait) defer stopHandshaker() - stopServer, serverAddress := startServer(t, handshakerAddress, &wait) + stopServer, serverAddress := startServer(t, handshakerAddress) defer stopServer() // Ping the server, authenticating with ALTS. @@ -444,31 +445,22 @@ func startFakeHandshakerService(t *testing.T, wait *sync.WaitGroup) (stop func() return func() { s.Stop() }, listener.Addr().String() } -func startServer(t *testing.T, handshakerServiceAddress string, wait *sync.WaitGroup) (stop func(), address string) { +func startServer(t *testing.T, handshakerServiceAddress string) (stop func(), address string) { listener, err := testutils.LocalTCPListener() if err != nil { t.Fatalf("LocalTCPListener() failed: %v", err) } serverOpts := &ServerOptions{HandshakerServiceAddress: handshakerServiceAddress} creds := NewServerCreds(serverOpts) - s := grpc.NewServer(grpc.Creds(creds)) - testgrpc.RegisterTestServiceServer(s, &testServer{}) - wait.Add(1) - go func() { - defer wait.Done() - if err := s.Serve(listener); err != nil { - t.Errorf("s.Serve(%v) failed: %v", listener, err) - } - }() - return func() { s.Stop() }, listener.Addr().String() -} - -type testServer struct { - testgrpc.UnimplementedTestServiceServer -} - -func (s *testServer) UnaryCall(_ context.Context, _ *testpb.SimpleRequest) (*testpb.SimpleResponse, error) { - return &testpb.SimpleResponse{ - Payload: &testpb.Payload{}, - }, nil + stub := &stubserver.StubServer{ + Listener: listener, + UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) { + return &testpb.SimpleResponse{ + Payload: &testpb.Payload{}, + }, nil + }, + S: grpc.NewServer(grpc.Creds(creds)), + } + stubserver.StartTestService(t, stub) + return func() { stub.S.Stop() }, listener.Addr().String() } diff --git a/test/stats_test.go b/test/stats_test.go index 0c36b6f15..5ee00bb37 100644 --- a/test/stats_test.go +++ b/test/stats_test.go @@ -27,8 +27,10 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/internal/stubserver" "google.golang.org/grpc/interop" testgrpc "google.golang.org/grpc/interop/grpc_testing" + testpb "google.golang.org/grpc/interop/grpc_testing" "google.golang.org/grpc/peer" "google.golang.org/grpc/stats" ) @@ -59,18 +61,15 @@ func (s) TestPeerForClientStatsHandler(t *testing.T) { if err != nil { t.Fatal(err) } - s := grpc.NewServer() - testgrpc.RegisterTestServiceServer(s, interop.NewTestServer()) - errCh := make(chan error) - go func() { - errCh <- s.Serve(l) - }() - defer func() { - s.Stop() - if err := <-errCh; err != nil { - t.Error(err) - } - }() + ss := &stubserver.StubServer{ + Listener: l, + EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) { + return &testpb.Empty{}, nil + }, + S: grpc.NewServer(), + } + stubserver.StartTestService(t, ss) + defer ss.S.Stop() // Create client with stats handler and do some calls. cc, err := grpc.NewClient( diff --git a/xds/internal/server/listener_wrapper_test.go b/xds/internal/server/listener_wrapper_test.go index 981d65e46..e181a7b82 100644 --- a/xds/internal/server/listener_wrapper_test.go +++ b/xds/internal/server/listener_wrapper_test.go @@ -23,7 +23,6 @@ import ( "fmt" "net" "strconv" - "sync" "testing" "time" @@ -31,6 +30,7 @@ import ( "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/internal" + "google.golang.org/grpc/internal/stubserver" "google.golang.org/grpc/internal/testutils" "google.golang.org/grpc/internal/testutils/xds/e2e" testgrpc "google.golang.org/grpc/interop/grpc_testing" @@ -159,14 +159,6 @@ func (s) TestListenerWrapper(t *testing.T) { } } -type testService struct { - testgrpc.TestServiceServer -} - -func (*testService) EmptyCall(context.Context, *testpb.Empty) (*testpb.Empty, error) { - return &testpb.Empty{}, nil -} - // TestConnsCleanup tests that the listener wrapper clears it's connection // references when connections close. It sets up a listener wrapper and gRPC // Server, and connects to the server 100 times and makes an RPC each time, and @@ -220,14 +212,15 @@ func (s) TestConnsCleanup(t *testing.T) { } } - server := grpc.NewServer(grpc.Creds(insecure.NewCredentials())) - testgrpc.RegisterTestServiceServer(server, &testService{}) - wg := sync.WaitGroup{} - go func() { - if err := server.Serve(lw); err != nil { - t.Errorf("failed to serve: %v", err) - } - }() + ss := &stubserver.StubServer{ + Listener: lis, + EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) { + return &testpb.Empty{}, nil + }, + S: grpc.NewServer(grpc.Creds(insecure.NewCredentials())), + } + stubserver.StartTestService(t, ss) + defer ss.S.Stop() // Make 100 connections to the server, and make an RPC on each one. for i := 0; i < 100; i++ { @@ -255,6 +248,4 @@ func (s) TestConnsCleanup(t *testing.T) { t.Fatalf("timeout waiting for lis wrapper conns to clear, size: %v", lenConns) } - server.Stop() - wg.Wait() }