From 7d9b72dbbe0fdc49f668b263797f34dfcd680f73 Mon Sep 17 00:00:00 2001
From: Daulet Zhanguzin <daulet@cohere.com>
Date: Mon, 29 Aug 2022 20:43:35 -0700
Subject: [PATCH] Better formatting of FprintReader output

---
 grpcreplay/grpcreplay.go       |  5 +---
 grpcreplay/grpcreplay_test.go  | 17 ++++++++++++++
 grpcreplay/testdata/fprint.txt | 42 ++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 4 deletions(-)
 create mode 100644 grpcreplay/testdata/fprint.txt

diff --git a/grpcreplay/grpcreplay.go b/grpcreplay/grpcreplay.go
index 04ca0be..41eba07 100644
--- a/grpcreplay/grpcreplay.go
+++ b/grpcreplay/grpcreplay.go
@@ -573,10 +573,7 @@ func FprintReader(w io.Writer, r io.Reader) error {
 		switch {
 		case e.msg.msg != nil:
 			fmt.Fprintf(w, ", message:\n")
-			buf, err := prototext.Marshal(e.msg.msg)
-			if err != nil {
-				return err
-			}
+			buf := []byte(prototext.Format(e.msg.msg))
 			if _, err := w.Write(buf); err != nil {
 				return err
 			}
diff --git a/grpcreplay/grpcreplay_test.go b/grpcreplay/grpcreplay_test.go
index 6b426ae..aa9e8d6 100644
--- a/grpcreplay/grpcreplay_test.go
+++ b/grpcreplay/grpcreplay_test.go
@@ -19,6 +19,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"os"
 	"strings"
 	"testing"
 
@@ -619,3 +620,19 @@ func TestOutOfOrderStreamReplay(t *testing.T) {
 	buf = record(t, func(t *testing.T, conn *grpc.ClientConn) { run(t, conn, 1, 2) })
 	replay(t, buf, func(t *testing.T, conn *grpc.ClientConn) { run(t, conn, 2, 1) })
 }
+
+func TestFprintReader(t *testing.T) {
+	buf := record(t, testService)
+	got := &bytes.Buffer{}
+	err := FprintReader(got, buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want, err := os.ReadFile("testdata/fprint.txt")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(got.Bytes(), want) {
+		t.Errorf("got:\n%v\nwant:\n%v\n", got.String(), string(want))
+	}
+}
diff --git a/grpcreplay/testdata/fprint.txt b/grpcreplay/testdata/fprint.txt
new file mode 100644
index 0000000..fef1878
--- /dev/null
+++ b/grpcreplay/testdata/fprint.txt
@@ -0,0 +1,42 @@
+initial state: "\x01\x02\x03"
+#1: kind: REQUEST, method: /intstore.IntStore/Set, ref index: 0, message:
+name:  "a"
+value:  1
+#2: kind: RESPONSE, method: , ref index: 1, message:
+#3: kind: REQUEST, method: /intstore.IntStore/Get, ref index: 0, message:
+name:  "a"
+#4: kind: RESPONSE, method: , ref index: 3, message:
+name:  "a"
+value:  1
+#5: kind: REQUEST, method: /intstore.IntStore/Get, ref index: 0, message:
+name:  "x"
+#6: kind: RESPONSE, method: , ref index: 5, error: rpc error: code = NotFound desc = "x"
+#7: kind: CREATE_STREAM, method: /intstore.IntStore/ListItems, ref index: 0
+#8: kind: SEND, method: , ref index: 7, message:
+#9: kind: RECV, method: , ref index: 7, message:
+name:  "a"
+value:  1
+#10: kind: RECV, method: , ref index: 7, error: EOF
+#11: kind: CREATE_STREAM, method: /intstore.IntStore/SetStream, ref index: 0
+#12: kind: SEND, method: , ref index: 11, message:
+name:  "b"
+value:  2
+#13: kind: SEND, method: , ref index: 11, message:
+name:  "c"
+value:  3
+#14: kind: RECV, method: , ref index: 11, message:
+count:  2
+#15: kind: CREATE_STREAM, method: /intstore.IntStore/StreamChat, ref index: 0
+#16: kind: SEND, method: , ref index: 15, message:
+name:  "d"
+value:  4
+#17: kind: RECV, method: , ref index: 15, message:
+name:  "d"
+value:  4
+#18: kind: SEND, method: , ref index: 15, message:
+name:  "e"
+value:  5
+#19: kind: RECV, method: , ref index: 15, message:
+name:  "e"
+value:  5
+#20: kind: RECV, method: , ref index: 15, error: EOF