[agent] Use grpc.NewClient (jaegertracing#5392)

Pushkarm029 · May 4, 2024 · 67f2203 · 67f2203
1 parent cb0d857
commit 67f2203
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 79 deletions.
diff --git a/cmd/agent/app/reporter/grpc/builder.go b/cmd/agent/app/reporter/grpc/builder.go
@@ -102,8 +102,7 @@ func (b *ConnBuilder) CreateConnection(ctx context.Context, logger *zap.Logger,
 	dialOptions = append(dialOptions, grpc.WithUnaryInterceptor(grpc_retry.UnaryClientInterceptor(grpc_retry.WithMax(b.MaxRetry))))
 	dialOptions = append(dialOptions, b.AdditionalDialOptions...)
 
-	// TODO: Need to replace grpc.Dial with grpc.NewClient and pass test
-	conn, err := grpc.Dial(dialTarget, dialOptions...)
+	conn, err := grpc.NewClient(dialTarget, dialOptions...)
 	if err != nil {
 		return nil, err
 	}

diff --git a/cmd/agent/app/reporter/grpc/builder_test.go b/cmd/agent/app/reporter/grpc/builder_test.go
@@ -24,7 +24,6 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials"
 	yaml "gopkg.in/yaml.v2"
 
@@ -69,78 +68,52 @@ func TestBuilderFromConfig(t *testing.T) {
 
 func TestBuilderWithCollectors(t *testing.T) {
 	spanHandler1 := &mockSpanHandler{}
-	s1, addr1 := initializeGRPCTestServer(t, func(s *grpc.Server) {
+	s1, _ := initializeGRPCTestServer(t, func(s *grpc.Server) {
 		api_v2.RegisterCollectorServiceServer(s, spanHandler1)
 	})
 	defer s1.Stop()
 
 	tests := []struct {
-		target               string
-		name                 string
-		hostPorts            []string
-		checkSuffixOnly      bool
-		notifier             discovery.Notifier
-		discoverer           discovery.Discoverer
-		expectedError        string
-		checkConnectionState bool
-		expectedState        string
+		target          string
+		name            string
+		hostPorts       []string
+		checkSuffixOnly bool
+		notifier        discovery.Notifier
+		discoverer      discovery.Discoverer
+		expectedError   string
 	}{
 		{
-			target:               "///round_robin",
-			name:                 "with roundrobin schema",
-			hostPorts:            []string{"127.0.0.1:9876", "127.0.0.1:9877", "127.0.0.1:9878"},
-			checkSuffixOnly:      true,
-			notifier:             nil,
-			discoverer:           nil,
-			checkConnectionState: false,
+			target:          "///round_robin",
+			name:            "with roundrobin schema",
+			hostPorts:       []string{"127.0.0.1:9876", "127.0.0.1:9877", "127.0.0.1:9878"},
+			checkSuffixOnly: true,
+			notifier:        nil,
+			discoverer:      nil,
 		},
 		{
-			target:               "127.0.0.1:9876",
-			name:                 "with single host",
-			hostPorts:            []string{"127.0.0.1:9876"},
-			checkSuffixOnly:      false,
-			notifier:             nil,
-			discoverer:           nil,
-			checkConnectionState: false,
+			target:          "127.0.0.1:9876",
+			name:            "with single host",
+			hostPorts:       []string{"127.0.0.1:9876"},
+			checkSuffixOnly: false,
+			notifier:        nil,
+			discoverer:      nil,
 		},
 		{
-			target:               "///round_robin",
-			name:                 "with custom resolver and fixed discoverer",
-			hostPorts:            []string{"dns://random_stuff"},
-			checkSuffixOnly:      true,
-			notifier:             noopNotifier{},
-			discoverer:           discovery.FixedDiscoverer{},
-			checkConnectionState: false,
+			target:          "///round_robin",
+			name:            "with custom resolver and fixed discoverer",
+			hostPorts:       []string{"dns://random_stuff"},
+			checkSuffixOnly: true,
+			notifier:        noopNotifier{},
+			discoverer:      discovery.FixedDiscoverer{},
 		},
 		{
-			target:               "",
-			name:                 "without collectorPorts and resolver",
-			hostPorts:            nil,
-			checkSuffixOnly:      false,
-			notifier:             nil,
-			discoverer:           nil,
-			expectedError:        "at least one collector hostPort address is required when resolver is not available",
-			checkConnectionState: false,
-		},
-		{
-			target:               addr1.String(),
-			name:                 "with collector connection status ready",
-			hostPorts:            []string{addr1.String()},
-			checkSuffixOnly:      false,
-			notifier:             nil,
-			discoverer:           nil,
-			checkConnectionState: true,
-			expectedState:        "READY",
-		},
-		{
-			target:               "random_stuff",
-			name:                 "with collector connection status failure",
-			hostPorts:            []string{"random_stuff"},
-			checkSuffixOnly:      false,
-			notifier:             nil,
-			discoverer:           nil,
-			checkConnectionState: true,
-			expectedState:        "TRANSIENT_FAILURE",
+			target:          "",
+			name:            "without collectorPorts and resolver",
+			hostPorts:       nil,
+			checkSuffixOnly: false,
+			notifier:        nil,
+			discoverer:      nil,
+			expectedError:   "at least one collector hostPort address is required when resolver is not available",
 		},
 	}
 
@@ -159,9 +132,6 @@ func TestBuilderWithCollectors(t *testing.T) {
 				require.NoError(t, err)
 				defer conn.Close()
 				require.NotNil(t, conn)
-				if test.checkConnectionState {
-					assertConnectionState(t, conn, test.expectedState)
-				}
 				if test.checkSuffixOnly {
 					assert.True(t, strings.HasSuffix(conn.Target(), test.target))
 				} else {
@@ -395,20 +365,6 @@ func TestProxyClientTLS(t *testing.T) {
 	}
 }
 
-func assertConnectionState(t *testing.T, conn *grpc.ClientConn, expectedState string) {
-	for {
-		s := conn.GetState()
-		if s == connectivity.Ready {
-			assert.Equal(t, expectedState, s.String())
-			break
-		}
-		if s == connectivity.TransientFailure {
-			assert.Equal(t, expectedState, s.String())
-			break
-		}
-	}
-}
-
 type fakeInterceptor struct {
 	isCalled bool
 }