Feature flags enabling for tests.

It is decided to go with all features enabled by default behaviour
for tests. This should help us identify potential problems with
feature flag combinations enabled in production. To verify implementation
without feature flag enabled it should be disabled explicitly in the
test.

Part of: https://gitlab.com/gitlab-org/gitaly/-/issues/2702

3 files changed, 134 insertions, 54 deletions

diff --git a/internal/testhelper/grpc_test.go b/internal/testhelper/grpc_test.go
index 792d7d11a..6b1272c3f 100644
--- a/internal/testhelper/grpc_test.go
+++ b/internal/testhelper/grpc_test.go
@@ -1,7 +1,6 @@
 package testhelper_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -11,7 +10,11 @@ import (
 
 func TestSetCtxGrpcMethod(t *testing.T) {
 	expectFullMethodName := "/pinkypb/TakeOverTheWorld.SNARF"
-	ctx := testhelper.SetCtxGrpcMethod(context.Background(), expectFullMethodName)
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	ctx = testhelper.SetCtxGrpcMethod(ctx, expectFullMethodName)
 
 	actualFullMethodName, ok := grpc.Method(ctx)
 	require.True(t, ok, "expected context to contain server transport stream")
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index 72d803176..ee1562b01 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -21,6 +21,7 @@ import (
 	"path"
 	"path/filepath"
 	"runtime"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -146,6 +147,26 @@ func GitalyServersMetadata(t testing.TB, serverSocketPath string) metadata.MD {
 	return metadata.Pairs("gitaly-servers", base64.StdEncoding.EncodeToString(gitalyServersJSON))
 }
 
+// MergeOutgoingMetadata merges provided metadata-s and returns context with resulting value.
+func MergeOutgoingMetadata(ctx context.Context, md ...metadata.MD) context.Context {
+	ctxmd, ok := metadata.FromOutgoingContext(ctx)
+	if !ok {
+		return metadata.NewOutgoingContext(ctx, metadata.Join(md...))
+	}
+
+	return metadata.NewOutgoingContext(ctx, metadata.Join(append(md, ctxmd)...))
+}
+
+// MergeIncomingMetadata merges provided metadata-s and returns context with resulting value.
+func MergeIncomingMetadata(ctx context.Context, md ...metadata.MD) context.Context {
+	ctxmd, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return metadata.NewIncomingContext(ctx, metadata.Join(md...))
+	}
+
+	return metadata.NewIncomingContext(ctx, metadata.Join(append(md, ctxmd)...))
+}
+
 // isValidRepoPath checks whether a valid git repository exists at the given path.
 func isValidRepoPath(absolutePath string) bool {
 	if _, err := os.Stat(filepath.Join(absolutePath, "objects")); err != nil {
@@ -438,9 +459,36 @@ func mustFindNoRunningChildProcess() {
 	panic(fmt.Errorf("%s: %v", desc, err))
 }
 
+// ContextOpt returns a new context instance with the new additions to it.
+type ContextOpt func(context.Context) (context.Context, func())
+
+// ContextWithTimeout allows to set provided timeout to the context.
+func ContextWithTimeout(duration time.Duration) ContextOpt {
+	return func(ctx context.Context) (context.Context, func()) {
+		return context.WithTimeout(ctx, duration)
+	}
+}
+
 // Context returns a cancellable context.
-func Context() (context.Context, func()) {
-	return context.WithCancel(context.Background())
+func Context(opts ...ContextOpt) (context.Context, func()) {
+	ctx, cancel := context.WithCancel(context.Background())
+	for _, ff := range featureflag.All {
+		ctx = featureflag.IncomingCtxWithFeatureFlag(ctx, ff)
+		ctx = featureflag.OutgoingCtxWithFeatureFlags(ctx, ff)
+	}
+
+	cancels := make([]func(), len(opts)+1)
+	cancels[0] = cancel
+	for i, opt := range opts {
+		ctx, cancel = opt(ctx)
+		cancels[i+1] = cancel
+	}
+
+	return ctx, func() {
+		for i := len(cancels) - 1; i >= 0; i-- {
+			cancels[i]()
+		}
+	}
 }
 
 // CreateRepo creates a temporary directory for a repo, without initializing it
@@ -733,48 +781,41 @@ func WriteBlobs(t testing.TB, testRepoPath string, n int) []string {
 	return blobIDs
 }
 
-// FeatureSet is a representation of a set of features that are enabled
-// This is useful in situations where a test needs to test any combination of features toggled on and off
+// FeatureSet is a representation of a set of features that should be disabled.
+// This is useful in situations where a test needs to test any combination of features toggled on and off.
+// It is designed to disable features as all features are enabled by default, please see: testhelper.Context()
 type FeatureSet struct {
-	features     map[featureflag.FeatureFlag]bool
-	rubyFeatures map[featureflag.FeatureFlag]bool
+	features     map[featureflag.FeatureFlag]struct{}
+	rubyFeatures map[featureflag.FeatureFlag]struct{}
 }
 
-func (f FeatureSet) IsEnabled(flag featureflag.FeatureFlag) bool {
-	on, ok := f.features[flag]
-	if !ok {
-		return flag.OnByDefault
-	}
-
-	return on
+func (f FeatureSet) IsDisabled(flag featureflag.FeatureFlag) bool {
+	_, ok := f.features[flag]
+	return ok
 }
 
 func (f FeatureSet) String() string {
-	features := make([]string, 0, len(f.enabledFeatures()))
-	for _, feature := range f.enabledFeatures() {
+	features := make([]string, 0, len(f.features))
+	for feature := range f.features {
 		features = append(features, feature.Name)
 	}
 
-	return strings.Join(features, ",")
-}
-
-func (f FeatureSet) enabledFeatures() []featureflag.FeatureFlag {
-	var enabled []featureflag.FeatureFlag
-
-	for feature := range f.features {
-		enabled = append(enabled, feature)
+	if len(features) == 0 {
+		return "none"
 	}
 
-	return enabled
+	sort.Strings(features)
+
+	return strings.Join(features, ",")
 }
 
-func (f FeatureSet) WithParent(ctx context.Context) context.Context {
-	for _, enabledFeature := range f.enabledFeatures() {
-		if _, ok := f.rubyFeatures[enabledFeature]; ok {
-			ctx = featureflag.OutgoingCtxWithRubyFeatureFlags(ctx, enabledFeature)
+func (f FeatureSet) Disable(ctx context.Context) context.Context {
+	for feature := range f.features {
+		if _, ok := f.rubyFeatures[feature]; ok {
+			ctx = featureflag.OutgoingCtxWithRubyFeatureFlagValue(ctx, feature, "false")
 			continue
 		}
-		ctx = featureflag.OutgoingCtxWithFeatureFlags(ctx, enabledFeature)
+		ctx = featureflag.OutgoingCtxWithFeatureFlagValue(ctx, feature, "false")
 	}
 
 	return ctx
@@ -786,20 +827,20 @@ type FeatureSets []FeatureSet
 // NewFeatureSets takes a slice of go feature flags, and an optional variadic set of ruby feature flags
 // and returns a FeatureSets slice
 func NewFeatureSets(goFeatures []featureflag.FeatureFlag, rubyFeatures ...featureflag.FeatureFlag) (FeatureSets, error) {
-	rubyFeatureMap := make(map[featureflag.FeatureFlag]bool)
+	rubyFeatureMap := make(map[featureflag.FeatureFlag]struct{})
 	for _, rubyFeature := range rubyFeatures {
-		rubyFeatureMap[rubyFeature] = true
+		rubyFeatureMap[rubyFeature] = struct{}{}
 	}
 
 	// start with an empty feature set
-	f := []FeatureSet{{features: make(map[featureflag.FeatureFlag]bool), rubyFeatures: rubyFeatureMap}}
+	f := []FeatureSet{{features: make(map[featureflag.FeatureFlag]struct{}), rubyFeatures: rubyFeatureMap}}
 
 	allFeatures := append(goFeatures, rubyFeatures...)
 
 	for i := range allFeatures {
-		featureMap := make(map[featureflag.FeatureFlag]bool)
+		featureMap := make(map[featureflag.FeatureFlag]struct{})
 		for j := 0; j <= i; j++ {
-			featureMap[allFeatures[j]] = true
+			featureMap[allFeatures[j]] = struct{}{}
 		}
 
 		f = append(f, FeatureSet{features: featureMap, rubyFeatures: rubyFeatureMap})
diff --git a/internal/testhelper/testserver.go b/internal/testhelper/testserver.go
index b2218c492..2a14a8ee1 100644
--- a/internal/testhelper/testserver.go
+++ b/internal/testhelper/testserver.go
@@ -41,12 +41,6 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
-// PraefectEnabled returns whether or not tests should use a praefect proxy
-func PraefectEnabled() bool {
-	_, ok := os.LookupEnv("GITALY_TEST_PRAEFECT_BIN")
-	return ok
-}
-
 // TestServerOpt is an option for TestServer
 type TestServerOpt func(t *TestServer)
 
@@ -75,15 +69,19 @@ func NewTestServer(srv *grpc.Server, opts ...TestServerOpt) *TestServer {
 		opt(ts)
 	}
 
+	// the health service needs to be registered in order to support health checks on all
+	// gitaly services that are under test.
+	// The health check is executed by the praefect in case 'test-with-praefect' verification
+	// job is running.
+	healthpb.RegisterHealthServer(srv, health.NewServer())
+
 	return ts
 }
 
 // NewServerWithAuth creates a new test server with authentication
 func NewServerWithAuth(tb testing.TB, streamInterceptors []grpc.StreamServerInterceptor, unaryInterceptors []grpc.UnaryServerInterceptor, token string, opts ...TestServerOpt) *TestServer {
 	if token != "" {
-		if PraefectEnabled() {
-			opts = append(opts, WithToken(token))
-		}
+		opts = append(opts, WithToken(token))
 		streamInterceptors = append(streamInterceptors, serverauth.StreamServerInterceptor(auth.Config{Token: token}))
 		unaryInterceptors = append(unaryInterceptors, serverauth.UnaryServerInterceptor(auth.Config{Token: token}))
 	}
@@ -160,6 +158,11 @@ func (p *TestServer) Start() error {
 			Token: p.token,
 		},
 		MemoryQueueEnabled: true,
+		Failover: praefectconfig.Failover{
+			Enabled:           true,
+			BootstrapInterval: config.Duration(time.Microsecond),
+			MonitorInterval:   config.Duration(time.Second),
+		},
 	}
 
 	for _, storage := range p.storages {
@@ -212,11 +215,11 @@ func (p *TestServer) Start() error {
 	conn, err := grpc.Dial("unix://"+praefectServerSocketPath, opts...)
 
 	if err != nil {
-		return fmt.Errorf("dial: %v", err)
+		return fmt.Errorf("dial to praefect: %v", err)
 	}
 	defer conn.Close()
 
-	if err = waitForPraefectStartup(conn); err != nil {
+	if err = WaitHealthy(conn, 3, time.Second); err != nil {
 		return err
 	}
 
@@ -234,25 +237,58 @@ func (p *TestServer) listen() (string, error) {
 	}
 
 	go p.grpcServer.Serve(listener)
+
+	opts := []grpc.DialOption{grpc.WithInsecure()}
+	if p.token != "" {
+		opts = append(opts, grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(p.token)))
+	}
+
+	conn, err := grpc.Dial("unix://"+gitalyServerSocketPath, opts...)
+
+	if err != nil {
+		return "", err
+	}
+	defer conn.Close()
+
+	if err := WaitHealthy(conn, 3, time.Second); err != nil {
+		return "", err
+	}
+
 	return gitalyServerSocketPath, nil
 }
 
-func waitForPraefectStartup(conn *grpc.ClientConn) error {
-	client := healthpb.NewHealthClient(conn)
+// WaitHealthy executes health check request `retries` times and awaits each `timeout` period to respond.
+// After `retries` unsuccessful attempts it returns an error.
+// Returns immediately without an error once get a successful health check response.
+func WaitHealthy(conn *grpc.ClientConn, retries int, timeout time.Duration) error {
+	for i := 0; i < retries; i++ {
+		if IsHealthy(conn, timeout) {
+			return nil
+		}
+	}
+
+	return errors.New("server not yet ready to serve")
+}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+// IsHealthy creates a health client to passed in connection and send `Check` request.
+// It waits for `timeout` duration to get response back.
+// It returns `true` only if remote responds with `SERVING` status.
+func IsHealthy(conn *grpc.ClientConn, timeout time.Duration) bool {
+	healthClient := healthpb.NewHealthClient(conn)
+
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
-	resp, err := client.Check(ctx, &healthpb.HealthCheckRequest{}, grpc.WaitForReady(true))
+	resp, err := healthClient.Check(ctx, &healthpb.HealthCheckRequest{}, grpc.WaitForReady(true))
 	if err != nil {
-		return err
+		return false
 	}
 
 	if resp.Status != healthpb.HealthCheckResponse_SERVING {
-		return errors.New("server not yet ready to serve")
+		return false
 	}
 
-	return nil
+	return true
 }
 
 // NewServer creates a Server for testing purposes