Changed from grpc.Dial to grpc.NewClient

Author: sonamtenzin2

grpc_client.go

@@ -24,9 +24,6 @@ func dialGRPCConn(tls *tls.Config, dialer func(context.Context, string) (net.Con
 	// We use a custom dialer so that we can connect over unix domain sockets.
 	opts = append(opts, grpc.WithContextDialer(dialer))
 
-	// Fail right away
-	opts = append(opts, grpc.FailOnNonTempDialError(true))
-
 	// If we have no TLS configuration set, we need to explicitly tell grpc
 	// that we're connecting with an insecure connection.
 	if tls == nil {
@@ -45,7 +42,7 @@ func dialGRPCConn(tls *tls.Config, dialer func(context.Context, string) (net.Con
 
 	// Connect. Note the first parameter is unused because we use a custom
 	// dialer that has the state to see the address.
-	conn, err := grpc.Dial("unused", opts...)
+	conn, err := grpc.NewClient("passthrough:///unused", opts...)
 	if err != nil {
 		return nil, err
 	}

testing.go

@@ -10,10 +10,12 @@ import (
 	"net"
 	"net/rpc"
 	"testing"
+	"time"
 
 	hclog "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/go-plugin/internal/grpcmux"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
 )
 
@@ -120,15 +122,26 @@ func TestGRPCConn(t testing.TB, register func(*grpc.Server)) (*grpc.ClientConn,
 	go func() { _ = server.Serve(l) }()
 
 	// Connect to the server
-	conn, err := grpc.Dial(
+	conn, err := grpc.NewClient(
 		l.Addr().String(),
-		grpc.WithBlock(),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
 
+	// Wait for the connection to be established before closing the listener
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	conn.Connect()
+
+	// Wait until connection reaches Ready state
+	for state := conn.GetState(); state != connectivity.Ready; state = conn.GetState() {
+		if !conn.WaitForStateChange(ctx, state) {
+			t.Fatalf("failed to establish connection, final state: %v", conn.GetState())
+		}
+	}
+
 	// Connection successful, close the listener
 	_ = l.Close()