refs: Add 'LocalBranches' to 'FindLocalBranchesResponse'1294-simplify-response-type-of-findlocalbranches

Using a feature flag 'SimplifyFindLocalBranchesResponse', let's add
'LocalBranches' to the 'FindLocalBranchesResponse' structure. After the
rollout we'll deprecate the usage of 'Branches' field and only use
'LocalBranches'.

Modify the tests to accommodate this new flag.

2 files changed, 240 insertions, 98 deletions

diff --git a/internal/gitaly/service/ref/refs_test.go b/internal/gitaly/service/ref/refs_test.go
index 02d1163bd..1a09d9fa2 100644
--- a/internal/gitaly/service/ref/refs_test.go
+++ b/internal/gitaly/service/ref/refs_test.go
@@ -5,6 +5,7 @@ package ref
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"strings"
@@ -16,6 +17,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/updateref"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
@@ -252,47 +254,77 @@ func TestInvalidRepoFindDefaultBranchNameRequest(t *testing.T) {
 }
 
 func TestSuccessfulFindLocalBranches(t *testing.T) {
-	ctx := testhelper.Context(t)
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testSuccessfulFindLocalBranches)
+}
+
+func testSuccessfulFindLocalBranches(t *testing.T, ctx context.Context) {
 	_, repo, _, client := setupRefService(ctx, t)
 
 	rpcRequest := &gitalypb.FindLocalBranchesRequest{Repository: repo}
 	c, err := client.FindLocalBranches(ctx, rpcRequest)
 	require.NoError(t, err)
 
-	var branches []*gitalypb.FindLocalBranchResponse
-	for {
-		r, err := c.Recv()
-		if err == io.EOF {
-			break
+	if featureflag.SimplifyFindLocalBranchesResponse.IsEnabled(ctx) {
+		var branches []*gitalypb.Branch
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			branches = append(branches, r.GetLocalBranches()...)
 		}
-		require.NoError(t, err)
-		branches = append(branches, r.GetBranches()...)
-	}
 
-	for name, target := range localBranches {
-		localBranch := &gitalypb.FindLocalBranchResponse{
-			Name:          []byte(name),
-			CommitId:      target.Id,
-			CommitSubject: target.Subject,
-			CommitAuthor: &gitalypb.FindLocalBranchCommitAuthor{
-				Name:  target.Author.Name,
-				Email: target.Author.Email,
-				Date:  target.Author.Date,
-			},
-			CommitCommitter: &gitalypb.FindLocalBranchCommitAuthor{
-				Name:  target.Committer.Name,
-				Email: target.Committer.Email,
-				Date:  target.Committer.Date,
-			},
-			Commit: target,
+		for name, target := range localBranches {
+			localBranch := &gitalypb.Branch{
+				Name:         []byte(name),
+				TargetCommit: target,
+			}
+
+			assertContainsBranch(t, branches, localBranch)
+		}
+
+	} else {
+		var branches []*gitalypb.FindLocalBranchResponse
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			branches = append(branches, r.GetBranches()...)
 		}
 
-		assertContainsLocalBranch(t, branches, localBranch)
+		for name, target := range localBranches {
+			localBranch := &gitalypb.FindLocalBranchResponse{
+				Name:          []byte(name),
+				CommitId:      target.Id,
+				CommitSubject: target.Subject,
+				CommitAuthor: &gitalypb.FindLocalBranchCommitAuthor{
+					Name:  target.Author.Name,
+					Email: target.Author.Email,
+					Date:  target.Author.Date,
+				},
+				CommitCommitter: &gitalypb.FindLocalBranchCommitAuthor{
+					Name:  target.Committer.Name,
+					Email: target.Committer.Email,
+					Date:  target.Committer.Date,
+				},
+				Commit: target,
+			}
+
+			assertContainsLocalBranch(t, branches, localBranch)
+		}
 	}
 }
 
-func TestFindLocalBranches_huge_committer(t *testing.T) {
-	ctx := testhelper.Context(t)
+func TestFindLocalBranchesHugeCommitter(t *testing.T) {
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testFindLocalBranchesHugeCommitter)
+}
+
+func testFindLocalBranchesHugeCommitter(t *testing.T, ctx context.Context) {
 	cfg, repo, repoPath, client := setupRefService(ctx, t)
 
 	gittest.WriteCommit(t, cfg, repoPath,
@@ -315,7 +347,11 @@ func TestFindLocalBranches_huge_committer(t *testing.T) {
 }
 
 func TestFindLocalBranchesPagination(t *testing.T) {
-	ctx := testhelper.Context(t)
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testFindLocalBranchesPagination)
+}
+
+func testFindLocalBranchesPagination(t *testing.T, ctx context.Context) {
 	_, repo, _, client := setupRefService(ctx, t)
 
 	limit := 1
@@ -329,42 +365,66 @@ func TestFindLocalBranchesPagination(t *testing.T) {
 	c, err := client.FindLocalBranches(ctx, rpcRequest)
 	require.NoError(t, err)
 
-	var branches []*gitalypb.FindLocalBranchResponse
-	for {
-		r, err := c.Recv()
-		if err == io.EOF {
-			break
+	expectedBranch := "refs/heads/improve/awesome"
+	target := localBranches[expectedBranch]
+
+	if featureflag.SimplifyFindLocalBranchesResponse.IsEnabled(ctx) {
+		var branches []*gitalypb.Branch
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			branches = append(branches, r.GetLocalBranches()...)
 		}
-		require.NoError(t, err)
-		branches = append(branches, r.GetBranches()...)
-	}
 
-	require.Len(t, branches, limit)
+		require.Len(t, branches, limit)
 
-	expectedBranch := "refs/heads/improve/awesome"
-	target := localBranches[expectedBranch]
+		branch := &gitalypb.Branch{
+			Name:         []byte(expectedBranch),
+			TargetCommit: target,
+		}
+		assertContainsBranch(t, branches, branch)
+	} else {
+		var branches []*gitalypb.FindLocalBranchResponse
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			branches = append(branches, r.GetBranches()...)
+		}
 
-	branch := &gitalypb.FindLocalBranchResponse{
-		Name:          []byte(expectedBranch),
-		CommitId:      target.Id,
-		CommitSubject: target.Subject,
-		CommitAuthor: &gitalypb.FindLocalBranchCommitAuthor{
-			Name:  target.Author.Name,
-			Email: target.Author.Email,
-			Date:  target.Author.Date,
-		},
-		CommitCommitter: &gitalypb.FindLocalBranchCommitAuthor{
-			Name:  target.Committer.Name,
-			Email: target.Committer.Email,
-			Date:  target.Committer.Date,
-		},
-		Commit: target,
+		require.Len(t, branches, limit)
+
+		branch := &gitalypb.FindLocalBranchResponse{
+			Name:          []byte(expectedBranch),
+			CommitId:      target.Id,
+			CommitSubject: target.Subject,
+			CommitAuthor: &gitalypb.FindLocalBranchCommitAuthor{
+				Name:  target.Author.Name,
+				Email: target.Author.Email,
+				Date:  target.Author.Date,
+			},
+			CommitCommitter: &gitalypb.FindLocalBranchCommitAuthor{
+				Name:  target.Committer.Name,
+				Email: target.Committer.Email,
+				Date:  target.Committer.Date,
+			},
+			Commit: target,
+		}
+		assertContainsLocalBranch(t, branches, branch)
 	}
-	assertContainsLocalBranch(t, branches, branch)
 }
 
 func TestFindLocalBranchesPaginationSequence(t *testing.T) {
-	ctx := testhelper.Context(t)
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testFindLocalBranchesPaginationSequence)
+}
+
+func testFindLocalBranchesPaginationSequence(t *testing.T, ctx context.Context) {
 	_, repo, _, client := setupRefService(ctx, t)
 
 	limit := 2
@@ -377,44 +437,85 @@ func TestFindLocalBranchesPaginationSequence(t *testing.T) {
 	c, err := client.FindLocalBranches(ctx, firstRPCRequest)
 	require.NoError(t, err)
 
-	var firstResponseBranches []*gitalypb.FindLocalBranchResponse
-	for {
-		r, err := c.Recv()
-		if err == io.EOF {
-			break
+	if featureflag.SimplifyFindLocalBranchesResponse.IsEnabled(ctx) {
+		var firstResponseBranches []*gitalypb.Branch
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			firstResponseBranches = append(firstResponseBranches, r.GetLocalBranches()...)
 		}
+
+		require.Len(t, firstResponseBranches, limit)
+
+		secondRPCRequest := &gitalypb.FindLocalBranchesRequest{
+			Repository: repo,
+			PaginationParams: &gitalypb.PaginationParameter{
+				Limit:     1,
+				PageToken: string(firstResponseBranches[0].Name),
+			},
+		}
+		c, err = client.FindLocalBranches(ctx, secondRPCRequest)
 		require.NoError(t, err)
-		firstResponseBranches = append(firstResponseBranches, r.GetBranches()...)
-	}
 
-	require.Len(t, firstResponseBranches, limit)
+		var secondResponseBranches []*gitalypb.Branch
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			secondResponseBranches = append(secondResponseBranches, r.GetLocalBranches()...)
+		}
+
+		require.Len(t, secondResponseBranches, 1)
+		require.Equal(t, firstResponseBranches[1], secondResponseBranches[0])
+	} else {
+		var firstResponseBranches []*gitalypb.FindLocalBranchResponse
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			firstResponseBranches = append(firstResponseBranches, r.GetBranches()...)
+		}
 
-	secondRPCRequest := &gitalypb.FindLocalBranchesRequest{
-		Repository: repo,
-		PaginationParams: &gitalypb.PaginationParameter{
-			Limit:     1,
-			PageToken: string(firstResponseBranches[0].Name),
-		},
-	}
-	c, err = client.FindLocalBranches(ctx, secondRPCRequest)
-	require.NoError(t, err)
+		require.Len(t, firstResponseBranches, limit)
 
-	var secondResponseBranches []*gitalypb.FindLocalBranchResponse
-	for {
-		r, err := c.Recv()
-		if err == io.EOF {
-			break
+		secondRPCRequest := &gitalypb.FindLocalBranchesRequest{
+			Repository: repo,
+			PaginationParams: &gitalypb.PaginationParameter{
+				Limit:     1,
+				PageToken: string(firstResponseBranches[0].Name),
+			},
 		}
+		c, err = client.FindLocalBranches(ctx, secondRPCRequest)
 		require.NoError(t, err)
-		secondResponseBranches = append(secondResponseBranches, r.GetBranches()...)
-	}
 
-	require.Len(t, secondResponseBranches, 1)
-	require.Equal(t, firstResponseBranches[1], secondResponseBranches[0])
+		var secondResponseBranches []*gitalypb.FindLocalBranchResponse
+		for {
+			r, err := c.Recv()
+			if err == io.EOF {
+				break
+			}
+			require.NoError(t, err)
+			secondResponseBranches = append(secondResponseBranches, r.GetBranches()...)
+		}
+
+		require.Len(t, secondResponseBranches, 1)
+		require.Equal(t, firstResponseBranches[1], secondResponseBranches[0])
+	}
 }
 
 func TestFindLocalBranchesPaginationWithIncorrectToken(t *testing.T) {
-	ctx := testhelper.Context(t)
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testFindLocalBranchesPaginationWithIncorrectToken)
+}
+
+func testFindLocalBranchesPaginationWithIncorrectToken(t *testing.T, ctx context.Context) {
 	_, repo, _, client := setupRefService(ctx, t)
 
 	limit := 1
@@ -452,7 +553,11 @@ func isOrderedSubset(subset, set []string) bool {
 }
 
 func TestFindLocalBranchesSort(t *testing.T) {
-	ctx := testhelper.Context(t)
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testFindLocalBranchesSort)
+}
+
+func testFindLocalBranchesSort(t *testing.T, ctx context.Context) {
 	testCases := []struct {
 		desc          string
 		relativeOrder []string
@@ -492,9 +597,16 @@ func TestFindLocalBranchesSort(t *testing.T) {
 				}
 				require.NoError(t, err)
 
-				for _, branch := range r.GetBranches() {
-					branches = append(branches, string(branch.Name))
+				if featureflag.SimplifyFindLocalBranchesResponse.IsEnabled(ctx) {
+					for _, branch := range r.GetLocalBranches() {
+						branches = append(branches, string(branch.Name))
+					}
+				} else {
+					for _, branch := range r.GetBranches() {
+						branches = append(branches, string(branch.Name))
+					}
 				}
+
 			}
 
 			if !isOrderedSubset(testCase.relativeOrder, branches) {
@@ -505,10 +617,14 @@ func TestFindLocalBranchesSort(t *testing.T) {
 }
 
 func TestEmptyFindLocalBranchesRequest(t *testing.T) {
+	t.Parallel()
+	testhelper.NewFeatureSets(featureflag.SimplifyFindLocalBranchesResponse).Run(t, testEmptyFindLocalBranchesRequest)
+}
+
+func testEmptyFindLocalBranchesRequest(t *testing.T, ctx context.Context) {
 	_, client := setupRefServiceWithoutRepo(t)
 
 	rpcRequest := &gitalypb.FindLocalBranchesRequest{}
-	ctx := testhelper.Context(t)
 	c, err := client.FindLocalBranches(ctx, rpcRequest)
 	require.NoError(t, err)
 
diff --git a/internal/gitaly/service/ref/util.go b/internal/gitaly/service/ref/util.go
index 9a9304fda..309a5ad77 100644
--- a/internal/gitaly/service/ref/util.go
+++ b/internal/gitaly/service/ref/util.go
@@ -8,6 +8,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/catfile"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/lines"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 )
 
@@ -84,23 +85,48 @@ func buildBranch(ctx context.Context, objectReader catfile.ObjectReader, element
 
 func newFindLocalBranchesWriter(stream gitalypb.RefService_FindLocalBranchesServer, objectReader catfile.ObjectReader) lines.Sender {
 	return func(refs [][]byte) error {
-		var branches []*gitalypb.FindLocalBranchResponse
 		ctx := stream.Context()
+		var response *gitalypb.FindLocalBranchesResponse
 
-		for _, ref := range refs {
-			elements, err := parseRef(ref)
-			if err != nil {
-				return err
+		if featureflag.SimplifyFindLocalBranchesResponse.IsEnabled(ctx) {
+			var branches []*gitalypb.Branch
+
+			for _, ref := range refs {
+				elements, err := parseRef(ref)
+				if err != nil {
+					return err
+				}
+
+				branch, err := buildBranch(ctx, objectReader, elements)
+				if err != nil {
+					return err
+				}
+
+				branches = append(branches, branch)
 			}
 
-			target, err := catfile.GetCommit(ctx, objectReader, git.Revision(elements[1]))
-			if err != nil {
-				return err
+			response = &gitalypb.FindLocalBranchesResponse{LocalBranches: branches}
+		} else {
+			var branches []*gitalypb.FindLocalBranchResponse
+
+			for _, ref := range refs {
+				elements, err := parseRef(ref)
+				if err != nil {
+					return err
+				}
+
+				target, err := catfile.GetCommit(ctx, objectReader, git.Revision(elements[1]))
+				if err != nil {
+					return err
+				}
+
+				branches = append(branches, buildLocalBranch(elements[0], target))
 			}
 
-			branches = append(branches, buildLocalBranch(elements[0], target))
+			response = &gitalypb.FindLocalBranchesResponse{Branches: branches}
 		}
-		return stream.Send(&gitalypb.FindLocalBranchesResponse{Branches: branches})
+
+		return stream.Send(response)
 	}
 }