refactor: code review changes

Author: aryannr97

mcp/transport.go

@@ -422,15 +422,17 @@ func (t *ioConn) Read(ctx context.Context) (jsonrpc.Message, error) {
 	}
 
 	// Read the next byte to check if there is trailing data.
-	tr := make([]byte, 1)
-	_, err := in.Buffered().Read(tr)
-	if err != nil {
-		return nil, err
+	var tr [1]byte
+	n, err := in.Buffered().Read(tr[:])
+	if n > 0 {
+		// If read byte is not a newline, it is an error.
+		if tr[0] != '\n' {
+			return nil, fmt.Errorf("invalid trailing data at the end of stream")
+		}
 	}
-
-	// If the next byte is not a newline, it is an error.
-	if tr[0] != '\n' {
-		return nil, fmt.Errorf("invalid trailing data at the end of stream")
+	// Return error except for EOF
+	if err != nil && err != io.EOF {
+		return nil, err
 	}
 
 	msgs, batch, err := readBatch(raw)

mcp/transport_test.go

@@ -7,7 +7,6 @@ package mcp
 import (
 	"context"
 	"io"
-	"reflect"
 	"strings"
 	"testing"
 
@@ -54,43 +53,38 @@ func TestBatchFraming(t *testing.T) {
 	}
 }
 
-func Test_ioConn_Read_BadTrailingData(t *testing.T) {
-	type fields struct {
-		rwc io.ReadWriteCloser
-	}
-	type args struct {
-		ctx context.Context
-	}
+func TestIOConnRead(t *testing.T) {
 	tests := []struct {
-		name    string
-		fields  fields
-		args    args
-		want    string
-		wantErr bool
+		name  string
+		input string
+		want  string
 	}{
+
 		{
-			name: "bad data at the end of first valid json",
-			fields: fields{
-				rwc: rwc{
-					rc: io.NopCloser(strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"test","params":{}},`)),
-				},
-			},
-			args: args{
-				ctx: context.Background(),
-			},
-			want:    "invalid trailing data at the end of stream",
-			wantErr: true,
+			name:  "valid json input",
+			input: `{"jsonrpc":"2.0","id":1,"method":"test","params":{}}`,
+			want:  "",
+		},
+
+		{
+			name: "newline at the end of first valid json input",
+			input: `{"jsonrpc":"2.0","id":1,"method":"test","params":{}}
+			`,
+			want: "",
+		},
+		{
+			name:  "bad data at the end of first valid json input",
+			input: `{"jsonrpc":"2.0","id":1,"method":"test","params":{}},`,
+			want:  "invalid trailing data at the end of stream",
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			tr := newIOConn(tt.fields.rwc)
-			_, err := tr.Read(tt.args.ctx)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("ioConn.Read() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-			if !reflect.DeepEqual(err.Error(), tt.want) {
+			tr := newIOConn(rwc{
+				rc: io.NopCloser(strings.NewReader(tt.input)),
+			})
+			_, err := tr.Read(context.Background())
+			if err != nil && err.Error() != tt.want {
 				t.Errorf("ioConn.Read() = %v, want %v", err.Error(), tt.want)
 			}
 		})