proxy: Use `t.Cleanup()` to simplify test setups

Simplify the setup of tests by using `t.Cleanup()` to finalize resources
instead of passing a cleanup function to the caller.

2 files changed, 13 insertions, 21 deletions

diff --git a/internal/praefect/grpc-proxy/proxy/peeker_test.go b/internal/praefect/grpc-proxy/proxy/peeker_test.go
index 276f175d9..8f7d9eb9e 100644
--- a/internal/praefect/grpc-proxy/proxy/peeker_test.go
+++ b/internal/praefect/grpc-proxy/proxy/peeker_test.go
@@ -22,8 +22,7 @@ import (
 func TestStreamPeeking(t *testing.T) {
 	ctx := testhelper.Context(t)
 
-	backendCC, backendSrvr, cleanupPinger := newBackendPinger(t, ctx)
-	defer cleanupPinger()
+	backendCC, backendSrvr := newBackendPinger(t, ctx)
 
 	pingReqSent := &testservice.PingRequest{Value: "hi"}
 
@@ -53,9 +52,7 @@ func TestStreamPeeking(t *testing.T) {
 		return stream.Send(pingResp)
 	}
 
-	proxyCC, cleanupProxy := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "PingStream")
-	defer cleanupProxy()
-
+	proxyCC := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "PingStream")
 	proxyClient := testservice.NewTestServiceClient(proxyCC)
 
 	proxyClientPingStream, err := proxyClient.PingStream(ctx)
@@ -79,8 +76,7 @@ func TestStreamPeeking(t *testing.T) {
 func TestStreamInjecting(t *testing.T) {
 	ctx := testhelper.Context(t)
 
-	backendCC, backendSrvr, cleanupPinger := newBackendPinger(t, ctx)
-	defer cleanupPinger()
+	backendCC, backendSrvr := newBackendPinger(t, ctx)
 
 	pingReqSent := &testservice.PingRequest{Value: "hi"}
 	newValue := "bye"
@@ -115,9 +111,7 @@ func TestStreamInjecting(t *testing.T) {
 		return stream.Send(pingResp)
 	}
 
-	proxyCC, cleanupProxy := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "PingStream")
-	defer cleanupProxy()
-
+	proxyCC := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "PingStream")
 	proxyClient := testservice.NewTestServiceClient(proxyCC)
 
 	proxyClientPingStream, err := proxyClient.PingStream(ctx)
diff --git a/internal/praefect/grpc-proxy/proxy/testhelper_test.go b/internal/praefect/grpc-proxy/proxy/testhelper_test.go
index 850d9df10..0b0efa104 100644
--- a/internal/praefect/grpc-proxy/proxy/testhelper_test.go
+++ b/internal/praefect/grpc-proxy/proxy/testhelper_test.go
@@ -26,15 +26,15 @@ func newListener(tb testing.TB) net.Listener {
 	return listener
 }
 
-func newBackendPinger(tb testing.TB, ctx context.Context) (*grpc.ClientConn, *interceptPinger, func()) {
+func newBackendPinger(tb testing.TB, ctx context.Context) (*grpc.ClientConn, *interceptPinger) {
 	ip := &interceptPinger{}
 
-	done := make(chan struct{})
 	srvr := grpc.NewServer()
 	listener := newListener(tb)
 
 	testservice.RegisterTestServiceServer(srvr, ip)
 
+	done := make(chan struct{})
 	go func() {
 		defer close(done)
 		srvr.Serve(listener)
@@ -51,26 +51,24 @@ func newBackendPinger(tb testing.TB, ctx context.Context) (*grpc.ClientConn, *in
 	)
 	require.NoError(tb, err)
 
-	cleanup := func() {
+	tb.Cleanup(func() {
 		srvr.GracefulStop()
 		require.NoError(tb, cc.Close())
 		<-done
-	}
+	})
 
-	return cc, ip, cleanup
+	return cc, ip
 }
 
-func newProxy(tb testing.TB, ctx context.Context, director proxy.StreamDirector, svc, method string) (*grpc.ClientConn, func()) {
+func newProxy(tb testing.TB, ctx context.Context, director proxy.StreamDirector, svc, method string) *grpc.ClientConn {
 	proxySrvr := grpc.NewServer(
 		grpc.ForceServerCodec(proxy.NewCodec()),
 		grpc.UnknownServiceHandler(proxy.TransparentHandler(director)),
 	)
-
 	proxy.RegisterService(proxySrvr, director, svc, method)
 
 	done := make(chan struct{})
 	listener := newListener(tb)
-
 	go func() {
 		defer close(done)
 		proxySrvr.Serve(listener)
@@ -84,13 +82,13 @@ func newProxy(tb testing.TB, ctx context.Context, director proxy.StreamDirector,
 	)
 	require.NoError(tb, err)
 
-	cleanup := func() {
+	tb.Cleanup(func() {
 		proxySrvr.GracefulStop()
 		require.NoError(tb, proxyCC.Close())
 		<-done
-	}
+	})
 
-	return proxyCC, cleanup
+	return proxyCC
 }
 
 // interceptPinger allows an RPC to be intercepted with a custom