2 files changed, 22 insertions, 16 deletions
diff --git a/client/testserver/gitalyserver.go b/client/testserver/gitalyserver.go
index 4bf14f3..301a737 100644
--- a/client/testserver/gitalyserver.go
+++ b/client/testserver/gitalyserver.go
@@ -58,9 +58,12 @@ func (s *TestGitalyServer) SSHUploadArchive(stream pb.SSHService_SSHUploadArchiv
 	return nil
 }
 
-func StartGitalyServer(t *testing.T) (string, *TestGitalyServer, func()) {
+func StartGitalyServer(t *testing.T) (string, *TestGitalyServer) {
+	t.Helper()
+
 	tempDir, _ := ioutil.TempDir("", "gitlab-shell-test-api")
 	gitalySocketPath := path.Join(tempDir, "gitaly.sock")
+	t.Cleanup(func() { os.RemoveAll(tempDir) })
 
 	err := os.MkdirAll(filepath.Dir(gitalySocketPath), 0700)
 	require.NoError(t, err)
@@ -74,12 +77,9 @@ func StartGitalyServer(t *testing.T) (string, *TestGitalyServer, func()) {
 	pb.RegisterSSHServiceServer(server, &testServer)
 
 	go server.Serve(listener)
+	t.Cleanup(func() { server.Stop() })
 
 	gitalySocketUrl := "unix:" + gitalySocketPath
-	cleanup := func() {
-		server.Stop()
-		os.RemoveAll(tempDir)
-	}
 
-	return gitalySocketUrl, &testServer, cleanup
+	return gitalySocketUrl, &testServer
 }
diff --git a/client/testserver/testserver.go b/client/testserver/testserver.go
index dee051b..0667cd4 100644
--- a/client/testserver/testserver.go
+++ b/client/testserver/testserver.go
@@ -27,9 +27,12 @@ type TestRequestHandler struct {
 	Handler func(w http.ResponseWriter, r *http.Request)
 }
 
-func StartSocketHttpServer(t *testing.T, handlers []TestRequestHandler) (string, func()) {
+func StartSocketHttpServer(t *testing.T, handlers []TestRequestHandler) string {
+	t.Helper()
+
 	err := os.MkdirAll(filepath.Dir(testSocket), 0700)
 	require.NoError(t, err)
+	t.Cleanup(func() { os.RemoveAll(tempDir) })
 
 	socketListener, err := net.Listen("unix", testSocket)
 	require.NoError(t, err)
@@ -44,16 +47,21 @@ func StartSocketHttpServer(t *testing.T, handlers []TestRequestHandler) (string,
 
 	url := "http+unix://" + testSocket
 
-	return url, cleanupSocket
+	return url
 }
 
-func StartHttpServer(t *testing.T, handlers []TestRequestHandler) (string, func()) {
+func StartHttpServer(t *testing.T, handlers []TestRequestHandler) string {
+	t.Helper()
+
 	server := httptest.NewServer(buildHandler(handlers))
+	t.Cleanup(func() { server.Close() })
 
-	return server.URL, server.Close
+	return server.URL
 }
 
-func StartHttpsServer(t *testing.T, handlers []TestRequestHandler, clientCAPath string) (string, func()) {
+func StartHttpsServer(t *testing.T, handlers []TestRequestHandler, clientCAPath string) string {
+	t.Helper()
+
 	crt := path.Join(testhelper.TestRoot, "certs/valid/server.crt")
 	key := path.Join(testhelper.TestRoot, "certs/valid/server.key")
 
@@ -63,7 +71,7 @@ func StartHttpsServer(t *testing.T, handlers []TestRequestHandler, clientCAPath
 
 	server.TLS = &tls.Config{
 		Certificates: []tls.Certificate{cer},
-		MinVersion: tls.VersionTLS12,
+		MinVersion:   tls.VersionTLS12,
 	}
 	server.TLS.BuildNameToCertificate()
 
@@ -80,11 +88,9 @@ func StartHttpsServer(t *testing.T, handlers []TestRequestHandler, clientCAPath
 
 	server.StartTLS()
 
-	return server.URL, server.Close
-}
+	t.Cleanup(func() { server.Close() })
 
-func cleanupSocket() {
-	os.RemoveAll(tempDir)
+	return server.URL
 }
 
 func buildHandler(handlers []TestRequestHandler) http.Handler {