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

Code before:
		},
	}
	if err := stub.Start(nil); err != nil {
		t.Fatalf("Failed to start StubServer: %v", err)
	}
	return stub, func() { stub.Stop() }
}

// createTestClient sets up a gRPC client connection with a manual resolver.
func createTestClient(t *testing.T, scheme string, statsHandler *testStatsHandler) (*grpc.ClientConn, *manual.Resolver) {
	rb := manual.NewBuilderWithScheme(scheme)
	cc, err := grpc.NewClient(
		scheme+":///test.server",
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithResolvers(rb),
		grpc.WithStatsHandler(statsHandler),
	)
	if err != nil {
		t.Fatalf("grpc.NewClient() failed: %v", err)
	}
	return cc, rb
}

// TestRPCSucceedsWithImmediateResolution ensures that when a resolver
// provides addresses immediately, RPC calls proceed without delay.
func (s) TestRPCSucceedsWithImmediateResolution(t *testing.T) {
	stub, cleanup := startStubServer(t)
	defer cleanup()

	statsHandler := &testStatsHandler{}
	rb := manual.NewBuilderWithScheme("instant")
	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.

Review comment (suggestion):
```
// TestClientConnRPC_WithoutNameResolutionDelay verify that if the resolution
// has already happened once before at the time of making RPC, the name resolution
// flag is not set indicating there was no delay in name resolution.
TestClientConnRPC_WithoutNameResolutionDelay 
```

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

	rb := manual.NewBuilderWithScheme("instant")
	rb.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: ss.Address}}})
	cc := ss.CC
	defer cc.Close()

	cc.Connect()
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
	client := testgrpc.NewTestServiceClient(cc)
	// Verify that the RPC succeeds.
	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("First RPC failed unexpectedly: %v", err)
	}
	// verifying that RPC was not blocked on resolver indicating there was no
	// delay in name resolution.
	if statsHandler.nameResolutionDelayed {
		t.Fatalf("statsHandler.nameResolutionDelayed = %v; want false", statsHandler.nameResolutionDelayed)
	}
}

// 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()