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 := range errChan {
if err != nil {
t.Fatalf(err.Error())
}
}
}
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
errChan := make(chan error)
go func() {
defer close(errChan)
conn, err := l.Accept()
l.Close()
if err != nil {
errChan <- err
return
}
server, err = yamux.Server(conn, nil)
if err != nil {
errChan <- 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
err = <-errChan
if err != nil {
t.Fatalf("err: %s", err)
}
return
}