agent: Revert "client: remove the parameter of 'enableYamux'"

This reverts commit 717ee24.

Since most of the runtime unit tests depending on unix socket without yamux,
and that commit would break those tests, thus I have to revert this commit.

Fixes: #662

Signed-off-by: lifupan <lifupan@gmail.com>

protocols/client/client.go

@@ -83,13 +83,13 @@ type dialer func(string, time.Duration) (net.Conn, error)
 //   - hvsock://<path>:<port>. Firecracker implements the virtio-vsock device
 //     model, and mediates communication between AF_UNIX sockets (on the host end)
 //     and AF_VSOCK sockets (on the guest end).
-func NewAgentClient(ctx context.Context, sock string) (*AgentClient, error) {
+func NewAgentClient(ctx context.Context, sock string, enableYamux bool) (*AgentClient, error) {
 	grpcAddr, parsedAddr, err := parse(sock)
 	if err != nil {
 		return nil, err
 	}
 	dialOpts := []grpc.DialOption{grpc.WithInsecure(), grpc.WithBlock()}
-	dialOpts = append(dialOpts, grpc.WithDialer(agentDialer(parsedAddr)))
+	dialOpts = append(dialOpts, grpc.WithDialer(agentDialer(parsedAddr, enableYamux)))
 
 	var tracer opentracing.Tracer
 
@@ -206,10 +206,8 @@ func heartBeat(session *yamux.Session) {
 	}
 }
 
-func agentDialer(addr *url.URL) dialer {
+func agentDialer(addr *url.URL, enableYamux bool) dialer {
 	var d dialer
-	var enableYamux bool
-
 	switch addr.Scheme {
 	case VSockSocketScheme:
 		d = vsockDialer
@@ -219,7 +217,6 @@ func agentDialer(addr *url.URL) dialer {
 		fallthrough
 	default:
 		d = unixDialer
-		enableYamux = true
 	}
 
 	if !enableYamux {

protocols/client/client_test.go

@@ -91,13 +91,13 @@ func checkVersion(cli *AgentClient) error {
 	return nil
 }
 
-func agentClientTest(t *testing.T, sock string, success bool, expect string) {
+func agentClientTest(t *testing.T, sock string, success, enableYamux bool, expect string) {
 	dialTimeout := defaultDialTimeout
 	defaultDialTimeout = 1 * time.Second
 	defer func() {
 		defaultDialTimeout = dialTimeout
 	}()
-	cli, err := NewAgentClient(context.Background(), sock)
+	cli, err := NewAgentClient(context.Background(), sock, enableYamux)
 	if success {
 		assert.Nil(t, err, "Failed to create new agent client: %s", err)
 	} else if !success {
@@ -116,16 +116,17 @@ func agentClientTest(t *testing.T, sock string, success bool, expect string) {
 }
 
 func TestNewAgentClient(t *testing.T) {
-	mock, waitCh, err := startMockServer(t, true)
+	mock, waitCh, err := startMockServer(t, false)
 	assert.Nil(t, err, "failed to start mock server: %s", err)
 	defer os.Remove(mockSockAddr)
 
 	cliFunc := func(sock string, success bool, expect string) {
-		agentClientTest(t, sock, success, expect)
+		agentClientTest(t, sock, success, false, expect)
 	}
 
 	// server starts
 	<-waitCh
+	cliFunc(mockSockAddr, true, "")
 	cliFunc(unixMockAddr, true, "")
 	cliFunc(mockBadSchemeAddr, false, "Invalid scheme:")
 	cliFunc(mockBadVsockScheme, false, "Invalid vsock scheme:")
@@ -144,7 +145,7 @@ func TestNewAgentClientWithYamux(t *testing.T) {
 	defer os.Remove(mockSockAddr)
 
 	cliFunc := func(sock string, success bool, expect string) {
-		agentClientTest(t, sock, success, expect)
+		agentClientTest(t, sock, success, true, expect)
 	}
 	// server starts
 	<-waitCh