Make 'code' argument saner

gitaly_test.go

@@ -199,13 +199,13 @@ func TestPostUploadPackHandledLocallyDueToEmptyGitalySocketPath(t *testing.T) {
 	testhelper.AssertResponseHeader(t, resp, "Content-Type", "application/x-git-upload-pack-result")
 }
 
-func startGitalyServer(t *testing.T, code codes.Code) (*grpc.Server, string) {
+func startGitalyServer(t *testing.T, finalMessageCode codes.Code) (*grpc.Server, string) {
 	socketPath := path.Join(scratchDir, fmt.Sprintf("gitaly-%d.sock", rand.Int()))
 	server := grpc.NewServer()
 	listener, err := net.Listen("unix", socketPath)
 	require.NoError(t, err)
 
-	pb.RegisterSmartHTTPServer(server, testhelper.NewGitalyServer(code))
+	pb.RegisterSmartHTTPServer(server, testhelper.NewGitalyServer(finalMessageCode))
 
 	go server.Serve(listener)
 

internal/testhelper/gitaly.go

@@ -15,7 +15,7 @@ import (
 )
 
 type GitalyTestServer struct {
-	code codes.Code
+	finalMessageCode codes.Code
 }
 
 const GitalyInfoRefsResponseMock = "Mock Gitaly InfoRefsResponse data"
@@ -33,8 +33,8 @@ func init() {
 	}
 }
 
-func NewGitalyServer(code codes.Code) *GitalyTestServer {
-	return &GitalyTestServer{code: code}
+func NewGitalyServer(finalMessageCode codes.Code) *GitalyTestServer {
+	return &GitalyTestServer{finalMessageCode: finalMessageCode}
 }
 
 func (s *GitalyTestServer) InfoRefsUploadPack(in *pb.InfoRefsRequest, stream pb.SmartHTTP_InfoRefsUploadPackServer) error {
@@ -45,7 +45,11 @@ func (s *GitalyTestServer) InfoRefsUploadPack(in *pb.InfoRefsRequest, stream pb.
 	response := &pb.InfoRefsResponse{
 		Data: []byte(GitalyInfoRefsResponseMock),
 	}
-	return stream.Send(response)
+	if err := stream.Send(response); err != nil {
+		return err
+	}
+
+	return s.finalError()
 }
 
 func (s *GitalyTestServer) InfoRefsReceivePack(in *pb.InfoRefsRequest, stream pb.SmartHTTP_InfoRefsReceivePackServer) error {
@@ -56,7 +60,11 @@ func (s *GitalyTestServer) InfoRefsReceivePack(in *pb.InfoRefsRequest, stream pb
 	response := &pb.InfoRefsResponse{
 		Data: []byte(GitalyInfoRefsResponseMock),
 	}
-	return stream.Send(response)
+	if err := stream.Send(response); err != nil {
+		return err
+	}
+
+	return s.finalError()
 }
 
 func (s *GitalyTestServer) PostReceivePack(stream pb.SmartHTTP_PostReceivePackServer) error {
@@ -99,7 +107,7 @@ func (s *GitalyTestServer) PostReceivePack(stream pb.SmartHTTP_PostReceivePackSe
 		}
 	}
 
-	return nil
+	return s.finalError()
 }
 
 func (s *GitalyTestServer) PostUploadPack(stream pb.SmartHTTP_PostUploadPackServer) error {
@@ -141,8 +149,12 @@ func (s *GitalyTestServer) PostUploadPack(stream pb.SmartHTTP_PostUploadPackServ
 		}
 	}
 
-	if s.code != codes.OK {
-		return grpc.Errorf(s.code, "error as specified by test")
+	return s.finalError()
+}
+
+func (s *GitalyTestServer) finalError() error {
+	if code := s.finalMessageCode; code != codes.OK {
+		return grpc.Errorf(code, "error as specified by test")
 	}
 
 	return nil