From c390896497eede44a1edcb8c877ad12aeddcde8c Mon Sep 17 00:00:00 2001
From: Wit Riewrangboonya <wit@uber.com>
Date: Mon, 23 Aug 2021 12:15:39 -0700
Subject: [PATCH] Review comments - use more mocks; replace gomock.Any()

---
 encoding/thrift/inbound_nowire_test.go | 22 ++++++++--------------
 1 file changed, 8 insertions(+), 14 deletions(-)

diff --git a/encoding/thrift/inbound_nowire_test.go b/encoding/thrift/inbound_nowire_test.go
index cf391697af..08bddbad6c 100644
--- a/encoding/thrift/inbound_nowire_test.go
+++ b/encoding/thrift/inbound_nowire_test.go
@@ -23,7 +23,6 @@ package thrift
 import (
 	"context"
 	"fmt"
-	"io"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -49,14 +48,6 @@ func (br *bodyReader) Decode(sr stream.Reader) error {
 	return br.err
 }
 
-type responseWriter struct{ err error }
-
-var _ stream.ResponseWriter = (*responseWriter)(nil)
-
-func (rw responseWriter) WriteResponse(wire.EnvelopeType, io.Writer, stream.Enveloper) error {
-	return rw.err
-}
-
 type responseHandler struct {
 	t   *testing.T
 	nwc *NoWireCall
@@ -179,11 +170,16 @@ func TestNoWireHandleWriteResponseError(t *testing.T) {
 	env := streamtest.NewMockEnveloper(mockCtrl)
 	env.EXPECT().EnvelopeType().Return(wire.Reply).Times(1)
 
-	proto := streamtest.NewMockRequestReader(mockCtrl)
-	proto.EXPECT().ReadRequest(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
-		Return(responseWriter{err: fmt.Errorf("write response error")}, nil)
+	rw := new(transporttest.FakeResponseWriter)
+	streamRw := streamtest.NewMockResponseWriter(mockCtrl)
+	streamRw.EXPECT().WriteResponse(wire.Reply, rw, env).Return(fmt.Errorf("write response error")).Times(1)
 
+	req := request()
 	br := &bodyReader{}
+	proto := streamtest.NewMockRequestReader(mockCtrl)
+	proto.EXPECT().ReadRequest(gomock.Any(), wire.Call, req.Body, br).
+		Return(streamRw, nil)
+
 	rh := responseHandler{t: t, reqBody: br, body: env}
 	h := thriftNoWireHandler{
 		Handler:       &rh,
@@ -193,8 +189,6 @@ func TestNoWireHandleWriteResponseError(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testtime.Second)
 	defer cancel()
 
-	req := request()
-	rw := new(transporttest.FakeResponseWriter)
 	err := h.Handle(ctx, req, rw)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "write response error")