packer-cn/packer/rpc/mux_broker_test.go

97 lines
1.5 KiB
Go

package rpc
import (
"fmt"
"net"
"testing"
"github.com/hashicorp/yamux"
)
func TestMuxBroker(t *testing.T) {
c, s := testYamux(t)
defer c.Close()
defer s.Close()
bc := newMuxBroker(c)
bs := newMuxBroker(s)
go bc.Run()
go bs.Run()
errChan := make(chan error, 1)
go func() {
defer close(errChan)
c, err := bc.Dial(5)
if err != nil {
errChan <- fmt.Errorf("err dialing: %s", err.Error())
return
}
if _, err := c.Write([]byte{42}); err != nil {
errChan <- fmt.Errorf("err writing: %s", err.Error())
}
}()
client, err := bs.Accept(5)
if err != nil {
t.Fatalf("err: %s", err)
}
var data [1]byte
if _, err := client.Read(data[:]); err != nil {
t.Fatalf("err: %s", err)
}
if data[0] != 42 {
t.Fatalf("bad: %d", data[0])
}
for {
err, open := <-errChan
if !open {
if err != nil {
t.Fatalf(err.Error())
}
break
}
}
}
func testYamux(t *testing.T) (client *yamux.Session, server *yamux.Session) {
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("err: %s", err)
}
// Server side
doneCh := make(chan struct{})
go func() {
defer close(doneCh)
conn, err := l.Accept()
l.Close()
if err != nil {
t.Fatalf("err: %s", err)
}
server, err = yamux.Server(conn, nil)
if err != nil {
t.Fatalf("err: %s", err)
}
}()
// Client side
conn, err := net.Dial("tcp", l.Addr().String())
if err != nil {
t.Fatalf("err: %s", err)
}
client, err = yamux.Client(conn, nil)
if err != nil {
t.Fatalf("err: %s", err)
}
// Wait for the server
<-doneCh
return
}