Merge branch 'smarthttp-tests' into 'master'

Use stream helpers in smarthttp tests

See merge request !175

2 files changed, 38 insertions, 92 deletions

diff --git a/internal/service/smarthttp/receive_pack_test.go b/internal/service/smarthttp/receive_pack_test.go
index ca27e4317..c7e8c9b9c 100644
--- a/internal/service/smarthttp/receive_pack_test.go
+++ b/internal/service/smarthttp/receive_pack_test.go
@@ -12,7 +12,9 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 
 	pb "gitlab.com/gitlab-org/gitaly-proto/go"
+	pbhelper "gitlab.com/gitlab-org/gitaly-proto/go/helper"
 
+	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/codes"
 )
@@ -68,49 +70,29 @@ func TestSuccessfulReceivePackRequest(t *testing.T) {
 	repo := &pb.Repository{Path: remoteRepoPath}
 	rpcRequest := &pb.PostReceivePackRequest{Repository: repo, GlId: "user-123", GlRepository: "project-123"}
 	stream, err := client.PostReceivePack(context.Background())
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 
-	if err := stream.Send(rpcRequest); err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, stream.Send(rpcRequest))
+
+	sw := pbhelper.NewSendWriter(func(p []byte) error {
+		return stream.Send(&pb.PostReceivePackRequest{Data: p})
+	})
+	_, err = io.Copy(sw, requestBuffer)
+	require.NoError(t, err)
 
-	data := make([]byte, 16)
-	for {
-		n, err := requestBuffer.Read(data)
-		if err == io.EOF {
-			break
-		} else if err != nil {
-			t.Fatal(err)
-		}
-
-		rpcRequest = &pb.PostReceivePackRequest{Data: data[:n]}
-		if err := stream.Send(rpcRequest); err != nil {
-			t.Fatal(err)
-		}
-	}
 	stream.CloseSend()
 
 	// Verify everything is going as planned
 	responseBuffer := bytes.Buffer{}
-	for {
-		rpcResponse, err := stream.Recv()
-		if err != nil {
-			if err == io.EOF {
-				break
-			} else {
-				t.Fatal(err)
-			}
-		}
-
-		responseBuffer.Write(rpcResponse.GetData())
-	}
+	rr := pbhelper.NewReceiveReader(func() ([]byte, error) {
+		resp, err := stream.Recv()
+		return resp.GetData(), err
+	})
+	_, err = io.Copy(&responseBuffer, rr)
+	require.NoError(t, err)
 
 	expectedResponse := "0030\x01000eunpack ok\n0019ok refs/heads/master\n00000000"
-	if responseBuffer.String() != expectedResponse {
-		t.Errorf("Expected response to be %q, got %q", expectedResponse, responseBuffer.String())
-	}
+	require.Equal(t, expectedResponse, responseBuffer.String(), "Expected response to be %q, got %q", expectedResponse, responseBuffer.String())
 
 	// The fact that this command succeeds means that we got the commit correctly, no further checks should be needed.
 	testhelper.MustRunCommand(t, nil, "git", "-C", remoteRepoPath, "show", string(newHead))
@@ -132,13 +114,9 @@ func TestFailedReceivePackRequestDueToValidationError(t *testing.T) {
 	for _, rpcRequest := range rpcRequests {
 		t.Logf("test case: %v", rpcRequest)
 		stream, err := client.PostReceivePack(context.Background())
-		if err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, err)
 
-		if err := stream.Send(&rpcRequest); err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, stream.Send(&rpcRequest))
 		stream.CloseSend()
 
 		err = drainPostReceivePackResponse(stream)
diff --git a/internal/service/smarthttp/upload_pack_test.go b/internal/service/smarthttp/upload_pack_test.go
index 68c9b7060..2387f2f60 100644
--- a/internal/service/smarthttp/upload_pack_test.go
+++ b/internal/service/smarthttp/upload_pack_test.go
@@ -66,51 +66,29 @@ func TestSuccessfulUploadPackRequest(t *testing.T) {
 	repo := &pb.Repository{Path: path.Join(remoteRepoPath, ".git")}
 	rpcRequest := &pb.PostUploadPackRequest{Repository: repo}
 	stream, err := client.PostUploadPack(context.Background())
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if err := stream.Send(rpcRequest); err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 
-	data := make([]byte, 16)
-	for {
-		n, err := requestBuffer.Read(data)
-		if err == io.EOF {
-			break
-		} else if err != nil {
-			t.Fatal(err)
-		}
+	require.NoError(t, stream.Send(rpcRequest))
 
-		rpcRequest = &pb.PostUploadPackRequest{Data: data[:n]}
-		if err := stream.Send(rpcRequest); err != nil {
-			t.Fatal(err)
-		}
-	}
+	sw := pbhelper.NewSendWriter(func(p []byte) error {
+		return stream.Send(&pb.PostUploadPackRequest{Data: p})
+	})
+	_, err = io.Copy(sw, requestBuffer)
+	require.NoError(t, err)
 	stream.CloseSend()
 
 	responseBuffer := &bytes.Buffer{}
-	for {
-		rpcResponse, err := stream.Recv()
-		if err != nil {
-			if err == io.EOF {
-				break
-			} else {
-				t.Fatal(err)
-			}
-		}
-
-		responseBuffer.Write(rpcResponse.GetData())
-	}
+	rr := pbhelper.NewReceiveReader(func() ([]byte, error) {
+		resp, err := stream.Recv()
+		return resp.GetData(), err
+	})
+	_, err = io.Copy(responseBuffer, rr)
+	require.NoError(t, err)
 
 	// There's no git command we can pass it this response and do the work for us (extracting pack file, ...),
 	// so we have to do it ourselves.
 	pack, version, entries := extractPackDataFromResponse(t, responseBuffer)
-	if pack == nil {
-		t.Errorf("Expected to find a pack file in response, found none")
-		return
-	}
+	require.NotNil(t, pack, "Expected to find a pack file in response, found none")
 
 	testhelper.MustRunCommand(t, bytes.NewReader(pack), "git", "-C", localRepoPath, "unpack-objects", fmt.Sprintf("--pack_header=%d,%d", version, entries))
 
@@ -163,13 +141,9 @@ func TestFailedUploadPackRequestDueToValidationError(t *testing.T) {
 	for _, rpcRequest := range rpcRequests {
 		t.Logf("test case: %v", rpcRequest)
 		stream, err := client.PostUploadPack(context.Background())
-		if err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, err)
 
-		if err := stream.Send(&rpcRequest); err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, stream.Send(&rpcRequest))
 		stream.CloseSend()
 
 		err = drainPostUploadPackResponse(stream)
@@ -205,15 +179,11 @@ func extractPackDataFromResponse(t *testing.T, buf *bytes.Buffer) ([]byte, int,
 		}
 
 		pktLen, err := strconv.ParseUint(string(pktLenStr), 16, 16)
-		if err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, err)
 
 		restPktLen := int(pktLen) - 4
 		pkt := buf.Next(restPktLen)
-		if len(pkt) != restPktLen {
-			t.Fatalf("Incomplete packet read")
-		}
+		require.Equal(t, restPktLen, len(pkt), "Incomplete packet read")
 
 		// The first byte of the packet is the band designator. We only care about data in band 1.
 		if pkt[0] == 1 {
@@ -226,9 +196,7 @@ func extractPackDataFromResponse(t *testing.T, buf *bytes.Buffer) ([]byte, int,
 	// 4 bytes for header version
 	// 4 bytes for header entries
 	// The rest is the pack file
-	if string(pack[:4]) != "PACK" {
-		t.Fatalf("Invalid packet signature")
-	}
+	require.Equal(t, "PACK", string(pack[:4]), "Invalid packet signature")
 	version := int(binary.BigEndian.Uint32(pack[4:8]))
 	entries := int(binary.BigEndian.Uint32(pack[8:12]))
 	pack = pack[12:]