github.com/rothwerx/packer@v0.9.0/packer/rpc/mux_broker_test.go (about)

     1  package rpc
     2  
     3  import (
     4  	"net"
     5  	"testing"
     6  
     7  	"github.com/hashicorp/yamux"
     8  )
     9  
    10  func TestMuxBroker(t *testing.T) {
    11  	c, s := testYamux(t)
    12  	defer c.Close()
    13  	defer s.Close()
    14  
    15  	bc := newMuxBroker(c)
    16  	bs := newMuxBroker(s)
    17  	go bc.Run()
    18  	go bs.Run()
    19  
    20  	go func() {
    21  		c, err := bc.Dial(5)
    22  		if err != nil {
    23  			t.Fatalf("err: %s", err)
    24  		}
    25  
    26  		if _, err := c.Write([]byte{42}); err != nil {
    27  			t.Fatalf("err: %s", err)
    28  		}
    29  	}()
    30  
    31  	client, err := bs.Accept(5)
    32  	if err != nil {
    33  		t.Fatalf("err: %s", err)
    34  	}
    35  
    36  	var data [1]byte
    37  	if _, err := client.Read(data[:]); err != nil {
    38  		t.Fatalf("err: %s", err)
    39  	}
    40  
    41  	if data[0] != 42 {
    42  		t.Fatalf("bad: %d", data[0])
    43  	}
    44  }
    45  
    46  func testYamux(t *testing.T) (client *yamux.Session, server *yamux.Session) {
    47  	l, err := net.Listen("tcp", "127.0.0.1:0")
    48  	if err != nil {
    49  		t.Fatalf("err: %s", err)
    50  	}
    51  
    52  	// Server side
    53  	doneCh := make(chan struct{})
    54  	go func() {
    55  		defer close(doneCh)
    56  		conn, err := l.Accept()
    57  		l.Close()
    58  		if err != nil {
    59  			t.Fatalf("err: %s", err)
    60  		}
    61  
    62  		server, err = yamux.Server(conn, nil)
    63  		if err != nil {
    64  			t.Fatalf("err: %s", err)
    65  		}
    66  	}()
    67  
    68  	// Client side
    69  	conn, err := net.Dial("tcp", l.Addr().String())
    70  	if err != nil {
    71  		t.Fatalf("err: %s", err)
    72  	}
    73  	client, err = yamux.Client(conn, nil)
    74  	if err != nil {
    75  		t.Fatalf("err: %s", err)
    76  	}
    77  
    78  	// Wait for the server
    79  	<-doneCh
    80  
    81  	return
    82  }