ssh: Use cleanup helper to simplify `newSSHClient()`

Simplify the `newSSHClient()` function so that it doesn't have to return
the gRPC connection by using `t.Cleanup()` to close it on test end.

4 files changed, 16 insertions, 25 deletions

diff --git a/internal/gitaly/service/ssh/receive_pack_test.go b/internal/gitaly/service/ssh/receive_pack_test.go
index 4c604c569..fe04dada4 100644
--- a/internal/gitaly/service/ssh/receive_pack_test.go
+++ b/internal/gitaly/service/ssh/receive_pack_test.go
@@ -46,8 +46,7 @@ func TestReceivePack_validation(t *testing.T) {
 
 	repo, _ := gittest.CreateRepository(ctx, t, cfg)
 
-	client, conn := newSSHClient(t, cfg.SocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, cfg.SocketPath)
 
 	for _, tc := range []struct {
 		desc        string
@@ -342,8 +341,7 @@ func TestReceivePack_hidesObjectPoolReferences(t *testing.T) {
 	repo := localrepo.NewTestRepo(t, cfg, repoProto)
 	txManager := transaction.NewManager(cfg, backchannel.NewRegistry())
 
-	client, conn := newSSHClient(t, cfg.SocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, cfg.SocketPath)
 
 	stream, err := client.SSHReceivePack(ctx)
 	require.NoError(t, err)
@@ -391,8 +389,7 @@ func TestReceivePack_transactional(t *testing.T) {
 
 	testcfg.BuildGitalyHooks(t, cfg)
 
-	client, conn := newSSHClient(t, cfg.SocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, cfg.SocketPath)
 
 	ctx, err := txinfo.InjectTransaction(ctx, 1, "node", true)
 	require.NoError(t, err)
diff --git a/internal/gitaly/service/ssh/testhelper_test.go b/internal/gitaly/service/ssh/testhelper_test.go
index a23087122..194fddeb8 100644
--- a/internal/gitaly/service/ssh/testhelper_test.go
+++ b/internal/gitaly/service/ssh/testhelper_test.go
@@ -5,6 +5,7 @@ package ssh
 import (
 	"testing"
 
+	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/service"
 	hookservice "gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/service/hook"
@@ -56,14 +57,12 @@ func startSSHServerWithOptions(t *testing.T, cfg config.Cfg, opts []ServerOpt, s
 	}, serverOpts...)
 }
 
-func newSSHClient(t *testing.T, serverSocketPath string) (gitalypb.SSHServiceClient, *grpc.ClientConn) {
-	connOpts := []grpc.DialOption{
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-	}
-	conn, err := grpc.Dial(serverSocketPath, connOpts...)
-	if err != nil {
-		t.Fatal(err)
-	}
+func newSSHClient(t *testing.T, serverSocketPath string) gitalypb.SSHServiceClient {
+	conn, err := grpc.Dial(serverSocketPath, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		testhelper.MustClose(t, conn)
+	})
 
-	return gitalypb.NewSSHServiceClient(conn), conn
+	return gitalypb.NewSSHServiceClient(conn)
 }
diff --git a/internal/gitaly/service/ssh/upload_archive_test.go b/internal/gitaly/service/ssh/upload_archive_test.go
index af2699a90..ea5fb77b3 100644
--- a/internal/gitaly/service/ssh/upload_archive_test.go
+++ b/internal/gitaly/service/ssh/upload_archive_test.go
@@ -30,8 +30,7 @@ func TestFailedUploadArchiveRequestDueToTimeout(t *testing.T) {
 		Seed: gittest.SeedGitLabTest,
 	})
 
-	client, conn := newSSHClient(t, cfg.SocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, cfg.SocketPath)
 
 	stream, err := client.SSHUploadArchive(ctx)
 	require.NoError(t, err)
@@ -64,8 +63,7 @@ func TestFailedUploadArchiveRequestDueToValidationError(t *testing.T) {
 
 	serverSocketPath := runSSHServer(t, cfg)
 
-	client, conn := newSSHClient(t, serverSocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, serverSocketPath)
 
 	tests := []struct {
 		Desc string
diff --git a/internal/gitaly/service/ssh/upload_pack_test.go b/internal/gitaly/service/ssh/upload_pack_test.go
index d7b311fc7..e6e2551b0 100644
--- a/internal/gitaly/service/ssh/upload_pack_test.go
+++ b/internal/gitaly/service/ssh/upload_pack_test.go
@@ -111,8 +111,7 @@ func testUploadPackTimeout(t *testing.T, ctx context.Context, opts ...testcfg.Op
 	repo, repoPath := gittest.CreateRepository(testhelper.Context(t), t, cfg)
 	gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("main"))
 
-	client, conn := newSSHClient(t, cfg.SocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, cfg.SocketPath)
 
 	stream, err := client.SSHUploadPack(ctx)
 	require.NoError(t, err)
@@ -435,8 +434,7 @@ func TestUploadPack_validation(t *testing.T) {
 
 	serverSocketPath := runSSHServer(t, cfg)
 
-	client, conn := newSSHClient(t, serverSocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, serverSocketPath)
 
 	for _, tc := range []struct {
 		desc        string
@@ -808,8 +806,7 @@ func testUploadPackGitFailure(t *testing.T, ctx context.Context) {
 		Seed: gittest.SeedGitLabTest,
 	})
 
-	client, conn := newSSHClient(t, cfg.SocketPath)
-	defer conn.Close()
+	client := newSSHClient(t, cfg.SocketPath)
 
 	// Writing an invalid config will allow repo to pass the `IsGitDirectory` check but still
 	// trigger an error when git tries to access the repo.