Repository: grpc/grpc-go
File: balancer/pickfirst/pickfirstleaf/metrics_test.go (Go)

Code before:
// TestPickFirstMetrics tests pick first metrics. It configures a pick first
// balancer, causes it to connect and then disconnect, and expects the
// subsequent metrics to emit from that.
func (s) TestPickFirstMetrics(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	ss := &stubserver.StubServer{
		EmptyCallF: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
			return &testpb.Empty{}, nil
		},
	}
	ss.StartServer()
	defer ss.Stop()

	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(pfConfig)

	r := manual.NewBuilderWithScheme("whatever")
	r.InitialState(resolver.State{
		ServiceConfig: sc,
		Addresses:     []resolver.Address{{Addr: ss.Address}}},
	)

	grpcTarget := r.Scheme() + ":///"
	tmr := stats.NewTestMetricsRecorder()
	cc, err := grpc.NewClient(grpcTarget, grpc.WithStatsHandler(tmr), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("NewClient() failed with error: %v", err)
	}
	defer cc.Close()

	tsc := testgrpc.NewTestServiceClient(cc)
	if _, err := tsc.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("EmptyCall() failed: %v", err)
	}

	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 1 {
		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 1)
	}
	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 0 {
		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 0)
	}
	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
	}

	ss.Stop()
	if err = pollForDisconnectedMetrics(ctx, tmr); err != nil {
		t.Fatal(err)
	}
}

Review comment (security):
nit: The opts can be refactored into a temp variable to shorten the line:
```go
opts := []grpc.DialOption{
   grpc.WithStatsHandler(tmr),
   grpc.WithTransportCredentials(insecure.NewCredentials()),
   grpc.WithResolvers(r)
}
```

Code after:
	defer cancel()

	ss := &stubserver.StubServer{
		EmptyCallF: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
			return &testpb.Empty{}, nil
		},
	}
	ss.StartServer()
	defer ss.Stop()

	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(pfConfig)

	r := manual.NewBuilderWithScheme("whatever")
	r.InitialState(resolver.State{
		ServiceConfig: sc,
		Addresses:     []resolver.Address{{Addr: ss.Address}}},
	)

	tmr := stats.NewTestMetricsRecorder()
	cc, err := grpc.NewClient(r.Scheme()+":///", grpc.WithStatsHandler(tmr), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("NewClient() failed with error: %v", err)
	}
	defer cc.Close()

	tsc := testgrpc.NewTestServiceClient(cc)
	if _, err := tsc.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("EmptyCall() failed: %v", err)
	}

	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 1 {
		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 1)
	}
	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 0 {
		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 0)
	}
	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
	}

	ss.Stop()
	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 1 {
		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 1)
	}
}

// TestPickFirstMetricsFailure tests the connection attempts failed metric. It
// configures a channel and scenario that causes a pick first connection attempt
// to fail, and then expects that metric to emit.
func (s) TestPickFirstMetricsFailure(t *testing.T) {