Fix grpc handler initialization

Signed-off-by: Pavol Loffay <ploffay@redhat.com>
jaegertracing · Nov 8, 2018 · d0ac010 · d0ac010
1 parent 2d4fed8
commit d0ac010
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/cmd/agent/app/reporter/grpc/reporter_test.go b/cmd/agent/app/reporter/grpc/reporter_test.go
@@ -50,9 +50,10 @@ func (h *mockSpanHandler) PostSpans(c context.Context, r *api_v2.PostSpansReques
 }
 
 func TestReporter_EmitZipkinBatch(t *testing.T) {
-	s, addr := initializeGRPCTestServer(t)
 	handler := &mockSpanHandler{}
-	api_v2.RegisterCollectorServiceServer(s, handler)
+	s, addr := initializeGRPCTestServer(t, func(s *grpc.Server) {
+		api_v2.RegisterCollectorServiceServer(s, handler)
+	})
 	defer s.Stop()
 	conn, err := grpc.Dial(addr.String(), grpc.WithInsecure())
 	defer conn.Close()
@@ -84,9 +85,10 @@ func TestReporter_EmitZipkinBatch(t *testing.T) {
 }
 
 func TestReporter_EmitBatch(t *testing.T) {
-	s, addr := initializeGRPCTestServer(t)
 	handler := &mockSpanHandler{}
-	api_v2.RegisterCollectorServiceServer(s, handler)
+	s, addr := initializeGRPCTestServer(t, func(s *grpc.Server) {
+		api_v2.RegisterCollectorServiceServer(s, handler)
+	})
 	defer s.Stop()
 	conn, err := grpc.Dial(addr.String(), grpc.WithInsecure())
 	defer conn.Close()

diff --git a/cmd/agent/app/reporter/grpc/sampling_manager_test.go b/cmd/agent/app/reporter/grpc/sampling_manager_test.go
@@ -28,8 +28,9 @@ import (
 )
 
 func TestSamplingManager_GetSamplingStrategy(t *testing.T) {
-	s, addr := initializeGRPCTestServer(t)
-	api_v2.RegisterSamplingManagerServer(s, &mockSamplingHandler{})
+	s, addr := initializeGRPCTestServer(t, func(s *grpc.Server) {
+		api_v2.RegisterSamplingManagerServer(s, &mockSamplingHandler{})
+	})
 	conn, err := grpc.Dial(addr.String(), grpc.WithInsecure())
 	defer conn.Close()
 	require.NoError(t, err)
@@ -64,10 +65,11 @@ func (*mockSamplingHandler) GetSamplingStrategy(context.Context, *api_v2.Samplin
 	return &api_v2.SamplingStrategyResponse{StrategyType: api_v2.SamplingStrategyType_PROBABILISTIC}, nil
 }
 
-func initializeGRPCTestServer(t *testing.T) (*grpc.Server, net.Addr) {
+func initializeGRPCTestServer(t *testing.T, beforeServe func(server *grpc.Server)) (*grpc.Server, net.Addr) {
 	server := grpc.NewServer()
 	lis, err := net.Listen("tcp", "localhost:0")
 	require.NoError(t, err)
+	beforeServe(server)
 	go func() {
 		err := server.Serve(lis)
 		require.NoError(t, err)