diff --git a/pkg/grpc_testing/stub_server.go b/pkg/grpc_testing/stub_server.go index e7e8c49d4..50b385476 100644 --- a/pkg/grpc_testing/stub_server.go +++ b/pkg/grpc_testing/stub_server.go @@ -26,10 +26,14 @@ type StubServer struct { s *grpc.Server cleanups []func() // Lambdas executed in Stop(); populated by Start(). + started chan struct{} } func New(testService testpb.TestServiceServer) *StubServer { - return &StubServer{testService: testService} + return &StubServer{ + testService: testService, + started: make(chan struct{}), + } } // Start starts the server and creates a client connected to it. @@ -50,7 +54,10 @@ func (ss *StubServer) Start(sopts []grpc.ServerOption, dopts ...grpc.DialOption) s := grpc.NewServer(sopts...) testpb.RegisterTestServiceServer(s, ss.testService) - go s.Serve(lis) + go func() { + close(ss.started) + s.Serve(lis) + }() ss.cleanups = append(ss.cleanups, s.Stop) ss.s = s @@ -59,6 +66,7 @@ func (ss *StubServer) Start(sopts []grpc.ServerOption, dopts ...grpc.DialOption) // Stop stops ss and cleans up all resources it consumed. func (ss *StubServer) Stop() { + <-ss.started for i := len(ss.cleanups) - 1; i >= 0; i-- { ss.cleanups[i]() }