add protobuf ReadMessages function

pkg/p2p/protobuf/protobuf.go

@@ -2,12 +2,15 @@ package protobuf
 
 import (
 	ggio "github.com/gogo/protobuf/io"
+	"github.com/gogo/protobuf/proto"
 	"github.com/janos/bee/pkg/p2p"
 	"io"
 )
 
 const delimitedReaderMaxSize = 128 * 1024 // max message size
 
+type Message = proto.Message
+
 func NewWriterAndReader(s p2p.Stream) (w ggio.Writer, r ggio.Reader) {
 	r = ggio.NewDelimitedReader(s, delimitedReaderMaxSize)
 	w = ggio.NewDelimitedWriter(s)
@@ -21,3 +24,18 @@ func NewReader(r io.Reader) ggio.Reader {
 func NewWriter(w io.Writer) ggio.Writer {
 	return ggio.NewDelimitedWriter(w)
 }
+
+func ReadMessages(r io.Reader, newMessage func() Message) (m []Message, err error) {
+	pr := NewReader(r)
+	for {
+		msg := newMessage()
+		if err := pr.ReadMsg(msg); err != nil {
+			if err == io.EOF {
+				break
+			}
+			return nil, err
+		}
+		m = append(m, msg)
+	}
+	return m, nil
+}

pkg/pingpong/pingpong_test.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"io"
 	"testing"
 
 	"github.com/janos/bee/pkg/p2p/mock"
@@ -34,39 +33,38 @@ func TestPing(t *testing.T) {
 		t.Errorf("invalid RTT value %v", rtt)
 	}
 
-	// validate received ping greetings
-	r := protobuf.NewReader(bytes.NewReader(streamer.In.Bytes()))
+	// validate received ping greetings from the client
+	wantGreetings := greetings
+	messages, err := protobuf.ReadMessages(
+		bytes.NewReader(streamer.In.Bytes()),
+		func() protobuf.Message { return new(pingpong.Ping) },
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
 	var gotGreetings []string
-	for {
-		var ping pingpong.Ping
-		if err := r.ReadMsg(&ping); err != nil {
-			if err == io.EOF {
-				break
-			}
-			t.Fatal(err)
-		}
-		gotGreetings = append(gotGreetings, ping.Greeting)
+	for _, m := range messages {
+		gotGreetings = append(gotGreetings, m.(*pingpong.Ping).Greeting)
 	}
-	if fmt.Sprint(gotGreetings) != fmt.Sprint(greetings) {
-		t.Errorf("got greetings %v, want %v", gotGreetings, greetings)
+	if fmt.Sprint(gotGreetings) != fmt.Sprint(wantGreetings) {
+		t.Errorf("got greetings %v, want %v", gotGreetings, wantGreetings)
 	}
 
-	// validate send pong responses by handler
-	r = protobuf.NewReader(bytes.NewReader(streamer.Out.Bytes()))
+	// validate sent pong responses by handler
 	var wantResponses []string
 	for _, g := range greetings {
 		wantResponses = append(wantResponses, "{"+g+"}")
 	}
+	messages, err = protobuf.ReadMessages(
+		bytes.NewReader(streamer.Out.Bytes()),
+		func() protobuf.Message { return new(pingpong.Pong) },
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
 	var gotResponses []string
-	for {
-		var pong pingpong.Pong
-		if err := r.ReadMsg(&pong); err != nil {
-			if err == io.EOF {
-				break
-			}
-			t.Fatal(err)
-		}
-		gotResponses = append(gotResponses, pong.Response)
+	for _, m := range messages {
+		gotResponses = append(gotResponses, m.(*pingpong.Pong).Response)
 	}
 	if fmt.Sprint(gotResponses) != fmt.Sprint(wantResponses) {
 		t.Errorf("got responses %v, want %v", gotResponses, wantResponses)