Merge branch 'pr-5353'

pkg/beam/unix_test.go

@@ -25,6 +25,30 @@ func TestSocketPair(t *testing.T) {
 	fmt.Printf("still open: %v\n", a.Fd())
 }
 
+func TestUSocketPair(t *testing.T) {
+	a, b, err := USocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	data := "hello world!"
+	go func() {
+		a.Write([]byte(data))
+		a.Close()
+	}()
+	res := make([]byte, 1024)
+	size, err := b.Read(res)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if size != len(data) {
+		t.Fatal("Unexpected size")
+	}
+	if string(res[0:size]) != data {
+		t.Fatal("Unexpected data")
+	}
+}
+
 func TestSendUnixSocket(t *testing.T) {
 	a1, a2, err := USocketPair()
 	if err != nil {
@@ -83,4 +107,131 @@ func TestSendUnixSocket(t *testing.T) {
 		t.Fatal(err)
 	}
 	fmt.Printf("---> %s\n", data)
+
+}
+
+// Ensure we get proper segmenting of messages
+func TestSendSegmenting(t *testing.T) {
+	a, b, err := USocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer a.Close()
+	defer b.Close()
+
+	extrafd1, extrafd2, err := SocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	extrafd2.Close()
+
+	go func() {
+		a.Send([]byte("message 1"), nil)
+		a.Send([]byte("message 2"), extrafd1)
+		a.Send([]byte("message 3"), nil)
+	}()
+
+	msg1, file1, err := b.Receive()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(msg1) != "message 1" {
+		t.Fatal("unexpected msg1:", string(msg1))
+	}
+	if file1 != nil {
+		t.Fatal("unexpectedly got file1")
+	}
+
+	msg2, file2, err := b.Receive()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(msg2) != "message 2" {
+		t.Fatal("unexpected msg2:", string(msg2))
+	}
+	if file2 == nil {
+		t.Fatal("didn't get file2")
+	}
+	file2.Close()
+
+	msg3, file3, err := b.Receive()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(msg3) != "message 3" {
+		t.Fatal("unexpected msg3:", string(msg3))
+	}
+	if file3 != nil {
+		t.Fatal("unexpectedly got file3")
+	}
+
+}
+
+// Test sending a zero byte message
+func TestSendEmpty(t *testing.T) {
+	a, b, err := USocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer a.Close()
+	defer b.Close()
+	go func() {
+		a.Send([]byte{}, nil)
+	}()
+
+	msg, file, err := b.Receive()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(msg) != 0 {
+		t.Fatalf("unexpected non-empty message: %v", msg)
+	}
+	if file != nil {
+		t.Fatal("unexpectedly got file")
+	}
+
+}
+
+func makeLarge(size int) []byte {
+	res := make([]byte, size)
+	for i := range res {
+		res[i] = byte(i % 255)
+	}
+	return res
+}
+
+func verifyLarge(data []byte, size int) bool {
+	if len(data) != size {
+		return false
+	}
+	for i := range data {
+		if data[i] != byte(i%255) {
+			return false
+		}
+	}
+	return true
+}
+
+// Test sending a large message
+func TestSendLarge(t *testing.T) {
+	a, b, err := USocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer a.Close()
+	defer b.Close()
+	go func() {
+		a.Send(makeLarge(100000), nil)
+	}()
+
+	msg, file, err := b.Receive()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !verifyLarge(msg, 100000) {
+		t.Fatalf("unexpected message (size %d)", len(msg))
+	}
+	if file != nil {
+		t.Fatal("unexpectedly got file")
+	}
 }