From 55c03cc119c2792b084928f12a15b93b3a6ed14a Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Wed, 29 Jan 2025 21:24:33 +0530
Subject: [PATCH] Always close test HTTP server during cleanup (#2261)

## Changes
This PR registers the `server.Close()` function to be run during test
cleanup in the server initialization function. This ensures that all
test servers are closed as soon as the test they are scoped to finish.

Motivated by https://github.com/databricks/cli/pull/2255/files where a
regression was introduced where we did not close the test server.

## Tests
N/A
---
 acceptance/cmd_server_test.go | 3 ++-
 acceptance/server_test.go     | 9 ---------
 libs/testserver/server.go     | 5 +----
 3 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/acceptance/cmd_server_test.go b/acceptance/cmd_server_test.go
index 3f5a6356e..9af63d0db 100644
--- a/acceptance/cmd_server_test.go
+++ b/acceptance/cmd_server_test.go
@@ -13,7 +13,8 @@ import (
 )
 
 func StartCmdServer(t *testing.T) *testserver.Server {
-	server := StartServer(t)
+	server := testserver.New(t)
+
 	server.Handle("/", func(r *http.Request) (any, error) {
 		q := r.URL.Query()
 		args := strings.Split(q.Get("args"), " ")
diff --git a/acceptance/server_test.go b/acceptance/server_test.go
index 66de5dcbf..98e351739 100644
--- a/acceptance/server_test.go
+++ b/acceptance/server_test.go
@@ -2,7 +2,6 @@ package acceptance_test
 
 import (
 	"net/http"
-	"testing"
 
 	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
@@ -11,14 +10,6 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
-func StartServer(t *testing.T) *testserver.Server {
-	server := testserver.New(t)
-	t.Cleanup(func() {
-		server.Close()
-	})
-	return server
-}
-
 func AddHandlers(server *testserver.Server) {
 	server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, error) {
 		return compute.ListPoliciesResponse{
diff --git a/libs/testserver/server.go b/libs/testserver/server.go
index 10269af8f..9ebfe3ba0 100644
--- a/libs/testserver/server.go
+++ b/libs/testserver/server.go
@@ -18,6 +18,7 @@ type Server struct {
 func New(t testutil.TestingT) *Server {
 	mux := http.NewServeMux()
 	server := httptest.NewServer(mux)
+	t.Cleanup(server.Close)
 
 	return &Server{
 		Server: server,
@@ -28,10 +29,6 @@ func New(t testutil.TestingT) *Server {
 
 type HandlerFunc func(req *http.Request) (resp any, err error)
 
-func (s *Server) Close() {
-	s.Server.Close()
-}
-
 func (s *Server) Handle(pattern string, handler HandlerFunc) {
 	s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
 		resp, err := handler(r)