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

Code before:
func (c *authorityCheckCreds) Info() credentials.ProtocolInfo {
	return credentials.ProtocolInfo{}
}
func (c *authorityCheckCreds) Clone() credentials.TransportCredentials {
	return c
}

// This test makes sure that the authority client handshake gets is the endpoint
// in dial target, not the resolved ip address.
func (s) TestCredsHandshakeAuthority(t *testing.T) {
	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	defer lis.Close()

	cred := &authorityCheckCreds{}
	s := grpc.NewServer()
	go s.Serve(lis)
	defer s.Stop()

	r := manual.NewBuilderWithScheme("whatever")

	r.InitialState(resolver.State{
		Addresses: []resolver.Address{{Addr: lis.Addr().String()}},
	})

	cc, err := grpc.NewClient(r.Scheme()+":///"+testAuthority, grpc.WithTransportCredentials(cred), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient(%q) = %v", lis.Addr().String(), err)
	}
	cc.Connect()
	defer cc.Close()

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	testutils.AwaitState(ctx, t, cc, connectivity.Ready)

	if cred.got != testAuthority {
		t.Fatalf("client creds got authority: %q, want: %q", cred.got, testAuthority)
	}
}

// This test makes sure that the authority client handshake gets is the endpoint
// of the ServerName of the address when it is set.
func (s) TestCredsHandshakeServerNameAuthority(t *testing.T) {
	const testServerName = "test.server.name"

	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)

Review comment (refactor):
I believe you can keep the `r.UpdateState` call below and remove this one to maintain the old test behaviour. Since a `cc.Connect()` is called, the resolver should be built before calling `r.UpdateState`. Same comment for the next test. 

Code after:
func (c *authorityCheckCreds) Info() credentials.ProtocolInfo {
	return credentials.ProtocolInfo{}
}
func (c *authorityCheckCreds) Clone() credentials.TransportCredentials {
	return c
}

// This test makes sure that the authority client handshake gets is the endpoint
// in dial target, not the resolved ip address.
func (s) TestCredsHandshakeAuthority(t *testing.T) {
	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	cred := &authorityCheckCreds{}
	s := grpc.NewServer()
	go s.Serve(lis)
	defer s.Stop()

	r := manual.NewBuilderWithScheme("whatever")

	cc, err := grpc.NewClient(r.Scheme()+":///"+testAuthority, grpc.WithTransportCredentials(cred), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient(%q) = %v", lis.Addr().String(), err)
	}
	defer cc.Close()
	cc.Connect()
	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: lis.Addr().String()}}})

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	testutils.AwaitState(ctx, t, cc, connectivity.Ready)

	if cred.got != testAuthority {
		t.Fatalf("client creds got authority: %q, want: %q", cred.got, testAuthority)
	}
}

// This test makes sure that the authority client handshake gets is the endpoint
// of the ServerName of the address when it is set.
func (s) TestCredsHandshakeServerNameAuthority(t *testing.T) {
	const testServerName = "test.server.name"

	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	cred := &authorityCheckCreds{}
	s := grpc.NewServer()
	go s.Serve(lis)
	defer s.Stop()