Repository: grpc/grpc-go
File: test/clientconn_test.go (Go)

Code before:
	rb.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: stub.Address}}})
	cc, err := grpc.NewClient(rb.Scheme()+":///test.server",
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithResolvers(rb),
		grpc.WithStatsHandler(statsHandler),
	)
	if err != nil {
		t.Fatalf("grpc.NewClient() failed: %v", err)
	}
	defer cc.Close()

	// Call an RPC to trigger waitForResolvedAddrs.
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()
	client := testgrpc.NewTestServiceClient(cc)
	// First RPC call should succeed immediately.
	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("First RPC failed unexpectedly: %v", err)
	}
	// Ensure that resolution was not delayed.
	if statsHandler.nameResolutionDelayed {
		t.Fatalf("Expected no name resolution delay, but it was detected.")
	}
}

// TestStatsHandlerDetectsResolutionDelay verifies that the StatsHandler
// detects delays in name resolution when using a manual resolver.
// Ensures that once resolution is updated, the RPC completes successfully.
// Checks that the StatsHandler correctly detects the name resolution delay.
func (s) TestStatsHandlerDetectsResolutionDelay(t *testing.T) {
	stub, cleanup := startStubServer(t)
	defer cleanup()

	statsHandler := &testStatsHandler{}
	cc, rb := createTestClient(t, "delayed", statsHandler)
	defer cc.Close()

	client := testgrpc.NewTestServiceClient(cc)
	resolutionReady := make(chan struct{})
	rpcCompleted := make(chan error, 1)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	go func() {
		_, err := client.EmptyCall(ctx, &testpb.Empty{})
		rpcCompleted <- err
	}()
	// Simulate delayed resolution and unblock it via resolutionReady
	go func() {
		<-resolutionReady
		rb.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: stub.Address}}})

Review comment (suggestion):
// TestStatsHandlerDetectsResolutionDelay verifies that if this is the 
// first time resolution is happening at the time of making RPC, nameResolutionDelayed
// flag was set indicating there was a delay in name resolution waiting for resolver to return addresses.

Code after:

// TestStatsHandlerDetectsResolutionDelay verifies that if this is the
// first time resolution is happening at the time of making RPC,
// nameResolutionDelayed flag is set indicating there was a delay in name
// resolution waiting for resolver to return addresses.
func (s) TestClientConnRPC_WithNameResolutionDelay(t *testing.T) {
	resolutionWait := grpcsync.NewEvent()
	prevHook := internal.NewStreamWaitingForResolver
	internal.NewStreamWaitingForResolver = func() { resolutionWait.Fire() }
	defer func() { internal.NewStreamWaitingForResolver = prevHook }()

	ss := &stubserver.StubServer{
		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
			return &testpb.Empty{}, nil
		},
	}
	if err := ss.Start(nil); err != nil {
		t.Fatalf("Failed to start StubServer: %v", err)
	}
	defer ss.Stop()

	statsHandler := &testStatsHandler{}
	rb := manual.NewBuilderWithScheme("delayed")
	cc, err := grpc.NewClient(rb.Scheme()+":///test.server",
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithResolvers(rb),
		grpc.WithStatsHandler(statsHandler),
	)
	if err != nil {
		t.Fatalf("grpc.NewClient() failed: %v", err)
	}
	defer cc.Close()
	go func() {
		<-resolutionWait.Done()
		rb.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: ss.Address}}})
	}()
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	client := testgrpc.NewTestServiceClient(cc)
	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("EmptyCall RPC failed: %v", err)
	}
	if !statsHandler.nameResolutionDelayed {
		t.Fatalf("statsHandler.nameResolutionDelayed = %v; want true", statsHandler.nameResolutionDelayed)
	}
}