proxy: Deduplicate proxy setup code used in tests

We have multiple functions that set up the proxy in our tests. Refactor
the code to deduplicate this logic.

2 files changed, 14 insertions, 32 deletions

diff --git a/internal/praefect/grpc-proxy/proxy/handler_ext_test.go b/internal/praefect/grpc-proxy/proxy/handler_ext_test.go
index bb36ad13c..241e3473f 100644
--- a/internal/praefect/grpc-proxy/proxy/handler_ext_test.go
+++ b/internal/praefect/grpc-proxy/proxy/handler_ext_test.go
@@ -18,14 +18,10 @@ import (
 	"testing"
 
 	"github.com/getsentry/sentry-go"
-	grpcmw "github.com/grpc-ecosystem/go-grpc-middleware"
-	grpcmwtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v15/client"
-	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/fieldextractors"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata"
-	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/sentryhandler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/grpc-proxy/proxy"
 	pb "gitlab.com/gitlab-org/gitaly/v15/internal/praefect/grpc-proxy/testdata"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
@@ -213,10 +209,7 @@ func setupProxy(t *testing.T) (context.Context, pb.TestServiceClient) {
 
 	ctx := testhelper.Context(t)
 
-	listenerProxy, err := net.Listen("tcp", "127.0.0.1:0")
-	require.NoError(t, err)
-	listenerServer, err := net.Listen("tcp", "127.0.0.1:0")
-	require.NoError(t, err)
+	listenerServer := newListener(t)
 
 	// Setup of the proxy's Director.
 	proxy2Server, err := grpc.Dial(listenerServer.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithDefaultCallOptions(grpc.ForceCodec(proxy.NewCodec())))
@@ -252,30 +245,7 @@ func setupProxy(t *testing.T) (context.Context, pb.TestServiceClient) {
 	}()
 	t.Cleanup(backendServer.Stop)
 
-	// Setup grpc-proxy server for test suite
-	proxyServer := grpc.NewServer(
-		grpc.ForceServerCodec(proxy.NewCodec()),
-		grpc.StreamInterceptor(
-			grpcmw.ChainStreamServer(
-				// context tags usage is required by sentryhandler.StreamLogHandler
-				grpcmwtags.StreamServerInterceptor(grpcmwtags.WithFieldExtractorForInitialReq(fieldextractors.FieldExtractor)),
-				// sentry middleware to capture errors
-				sentryhandler.StreamLogHandler,
-			),
-		),
-		grpc.UnknownServiceHandler(proxy.TransparentHandler(director)),
-	)
-
-	// Ping handler is handled as an explicit registration and not as a TransparentHandler.
-	proxy.RegisterService(proxyServer, director, "mwitkow.testproto.TestService", "Ping")
-	go func() {
-		proxyServer.Serve(listenerProxy)
-	}()
-	t.Cleanup(proxyServer.Stop)
-
-	client2Proxy, err := grpc.DialContext(ctx, listenerProxy.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
-	require.NoError(t, err)
-	t.Cleanup(func() { testhelper.MustClose(t, client2Proxy) })
+	client2Proxy := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "Ping")
 
 	return ctx, pb.NewTestServiceClient(client2Proxy)
 }
diff --git a/internal/praefect/grpc-proxy/proxy/testhelper_test.go b/internal/praefect/grpc-proxy/proxy/testhelper_test.go
index 5faf4ce49..bc7ffca5e 100644
--- a/internal/praefect/grpc-proxy/proxy/testhelper_test.go
+++ b/internal/praefect/grpc-proxy/proxy/testhelper_test.go
@@ -7,7 +7,11 @@ import (
 	"net"
 	"testing"
 
+	grpcmw "github.com/grpc-ecosystem/go-grpc-middleware"
+	grpcmwtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/fieldextractors"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/sentryhandler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/grpc-proxy/proxy"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 	"google.golang.org/grpc"
@@ -63,6 +67,14 @@ func newBackendPinger(tb testing.TB, ctx context.Context) (*grpc.ClientConn, *in
 func newProxy(tb testing.TB, ctx context.Context, director proxy.StreamDirector, svc, method string) *grpc.ClientConn {
 	proxySrvr := grpc.NewServer(
 		grpc.ForceServerCodec(proxy.NewCodec()),
+		grpc.StreamInterceptor(
+			grpcmw.ChainStreamServer(
+				// context tags usage is required by sentryhandler.StreamLogHandler
+				grpcmwtags.StreamServerInterceptor(grpcmwtags.WithFieldExtractorForInitialReq(fieldextractors.FieldExtractor)),
+				// sentry middleware to capture errors
+				sentryhandler.StreamLogHandler,
+			),
+		),
 		grpc.UnknownServiceHandler(proxy.TransparentHandler(director)),
 	)
 	proxy.RegisterService(proxySrvr, director, svc, method)