praefect: Refactor tests to allow more flexible mock backend setup

We're about to migrate Praefect tests to stop using their own Protobuf
service definitions. To make the transition easier, refactor the way we
set up mock backends to be more flexible. This enables every test to set
up a different service based on their needs.

4 files changed, 43 insertions, 34 deletions

diff --git a/internal/praefect/auth_test.go b/internal/praefect/auth_test.go
index 324392dce..72edb9862 100644
--- a/internal/praefect/auth_test.go
+++ b/internal/praefect/auth_test.go
@@ -132,7 +132,9 @@ func dial(serverSocketPath string, opts []grpc.DialOption) (*grpc.ClientConn, er
 
 func runServer(t *testing.T, token string, required bool) (*grpc.Server, string, func()) {
 	backendToken := "abcxyz"
-	backend, cleanup := newMockDownstream(t, backendToken, &mockSvc{})
+	backend, cleanup := newMockDownstream(t, backendToken, func(srv *grpc.Server) {
+		mock.RegisterSimpleServiceServer(srv, &mockSvc{})
+	})
 
 	conf := config.Config{
 		Auth: auth.Config{Token: token, Transitioning: !required},
diff --git a/internal/praefect/coordinator_test.go b/internal/praefect/coordinator_test.go
index 1d1897031..f643ce0c2 100644
--- a/internal/praefect/coordinator_test.go
+++ b/internal/praefect/coordinator_test.go
@@ -1774,6 +1774,9 @@ func TestCoordinatorEnqueueFailure(t *testing.T) {
 			return &emptypb.Empty{}, nil // always succeeds
 		},
 	}
+	registrar := func(srv *grpc.Server) {
+		mock.RegisterSimpleServiceServer(srv, ms)
+	}
 
 	r, err := protoregistry.NewFromPaths("praefect/mock/mock.proto")
 	require.NoError(t, err)
@@ -1782,9 +1785,9 @@ func TestCoordinatorEnqueueFailure(t *testing.T) {
 	cc, _, cleanup := RunPraefectServer(t, ctx, conf, BuildOptions{
 		WithAnnotations: r,
 		WithQueue:       queueInterceptor,
-		WithBackends: WithMockBackends(t, map[string]mock.SimpleServiceServer{
-			conf.VirtualStorages[0].Nodes[0].Storage: ms,
-			conf.VirtualStorages[0].Nodes[1].Storage: ms,
+		WithBackends: WithMockBackends(t, map[string]func(*grpc.Server){
+			conf.VirtualStorages[0].Nodes[0].Storage: registrar,
+			conf.VirtualStorages[0].Nodes[1].Storage: registrar,
 		}),
 	})
 	defer cleanup()
diff --git a/internal/praefect/server_test.go b/internal/praefect/server_test.go
index 438b61225..5a137b220 100644
--- a/internal/praefect/server_test.go
+++ b/internal/praefect/server_test.go
@@ -894,31 +894,35 @@ func TestProxyWrites(t *testing.T) {
 func TestErrorThreshold(t *testing.T) {
 	t.Parallel()
 	backendToken := ""
-	backend, cleanup := newMockDownstream(t, backendToken, &mockSvc{
-		repoMutatorUnary: func(ctx context.Context, req *mock.RepoRequest) (*emptypb.Empty, error) {
-			md, ok := metadata.FromIncomingContext(ctx)
-			if !ok {
-				return &emptypb.Empty{}, errors.New("couldn't read metadata")
-			}
+	backend, cleanup := newMockDownstream(t, backendToken, func(srv *grpc.Server) {
+		service := &mockSvc{
+			repoMutatorUnary: func(ctx context.Context, req *mock.RepoRequest) (*emptypb.Empty, error) {
+				md, ok := metadata.FromIncomingContext(ctx)
+				if !ok {
+					return &emptypb.Empty{}, errors.New("couldn't read metadata")
+				}
 
-			if md.Get("bad-header")[0] == "true" {
-				return &emptypb.Empty{}, helper.ErrInternalf("something went wrong")
-			}
+				if md.Get("bad-header")[0] == "true" {
+					return &emptypb.Empty{}, helper.ErrInternalf("something went wrong")
+				}
 
-			return &emptypb.Empty{}, nil
-		},
-		repoAccessorUnary: func(ctx context.Context, req *mock.RepoRequest) (*emptypb.Empty, error) {
-			md, ok := metadata.FromIncomingContext(ctx)
-			if !ok {
-				return &emptypb.Empty{}, errors.New("couldn't read metadata")
-			}
+				return &emptypb.Empty{}, nil
+			},
+			repoAccessorUnary: func(ctx context.Context, req *mock.RepoRequest) (*emptypb.Empty, error) {
+				md, ok := metadata.FromIncomingContext(ctx)
+				if !ok {
+					return &emptypb.Empty{}, errors.New("couldn't read metadata")
+				}
 
-			if md.Get("bad-header")[0] == "true" {
-				return &emptypb.Empty{}, helper.ErrInternalf("something went wrong")
-			}
+				if md.Get("bad-header")[0] == "true" {
+					return &emptypb.Empty{}, helper.ErrInternalf("something went wrong")
+				}
 
-			return &emptypb.Empty{}, nil
-		},
+				return &emptypb.Empty{}, nil
+			},
+		}
+
+		mock.RegisterSimpleServiceServer(srv, service)
 	})
 	defer cleanup()
 
diff --git a/internal/praefect/testserver.go b/internal/praefect/testserver.go
index f957634cf..3a452e2b4 100644
--- a/internal/praefect/testserver.go
+++ b/internal/praefect/testserver.go
@@ -16,7 +16,6 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/datastore"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/grpc-proxy/proxy"
-	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/mock"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/nodes"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/protoregistry"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/service"
@@ -58,20 +57,21 @@ type BuildOptions struct {
 	WithChecks []service.CheckFunc
 }
 
-// WithMockBackends mocks backends with a set of passed in stubs.
-func WithMockBackends(tb testing.TB, backends map[string]mock.SimpleServiceServer) func([]*config.VirtualStorage) []testhelper.Cleanup {
+// WithMockBackends mocks backends with a set of passed in functions that know to register a gRPC
+// server.
+func WithMockBackends(tb testing.TB, backendRegistrars map[string]func(*grpc.Server)) func([]*config.VirtualStorage) []testhelper.Cleanup {
 	return func(virtualStorages []*config.VirtualStorage) []testhelper.Cleanup {
 		var cleanups []testhelper.Cleanup
 
 		for _, vs := range virtualStorages {
-			require.Equal(tb, len(backends), len(vs.Nodes),
+			require.Equal(tb, len(backendRegistrars), len(vs.Nodes),
 				"mock server count doesn't match config nodes")
 
 			for i, node := range vs.Nodes {
-				backend, ok := backends[node.Storage]
-				require.True(tb, ok, "missing backend server for node %s", node.Storage)
+				backendRegistrar, ok := backendRegistrars[node.Storage]
+				require.True(tb, ok, "missing server registrator for node %s", node.Storage)
 
-				backendAddr, cleanup := newMockDownstream(tb, node.Token, backend)
+				backendAddr, cleanup := newMockDownstream(tb, node.Token, backendRegistrar)
 				cleanups = append(cleanups, cleanup)
 
 				node.Address = backendAddr
@@ -132,9 +132,9 @@ func dialLocalPort(tb testing.TB, port int, backend bool) *grpc.ClientConn {
 	return cc
 }
 
-func newMockDownstream(tb testing.TB, token string, m mock.SimpleServiceServer) (string, func()) {
+func newMockDownstream(tb testing.TB, token string, registerService func(*grpc.Server)) (string, func()) {
 	srv := grpc.NewServer(grpc.UnaryInterceptor(auth.UnaryServerInterceptor(gitalycfgauth.Config{Token: token})))
-	mock.RegisterSimpleServiceServer(srv, m)
+	registerService(srv)
 	healthpb.RegisterHealthServer(srv, health.NewServer())
 
 	// client to backend service