Merge ca08796f77 into f2096eddcc

NOTICE

@@ -114,7 +114,3 @@ dario.cat/mergo
 Copyright (c) 2013 Dario Castañé. All rights reserved.
 Copyright (c) 2012 The Go Authors. All rights reserved.
 https://github.com/darccio/mergo/blob/master/LICENSE
-
-https://github.com/gorilla/mux
-Copyright (c) 2023 The Gorilla Authors. All rights reserved.
-https://github.com/gorilla/mux/blob/main/LICENSE

acceptance/acceptance_test.go

@@ -11,7 +11,6 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"regexp"
 	"runtime"
 	"slices"
 	"sort"
@@ -27,7 +26,6 @@ import (
 	"github.com/databricks/cli/libs/testdiff"
 	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go"
-	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/stretchr/testify/require"
 )
 
@@ -74,8 +72,7 @@ func TestInprocessMode(t *testing.T) {
 	if InprocessMode {
 		t.Skip("Already tested by TestAccept")
 	}
-	require.Equal(t, 1, testAccept(t, true, "selftest/basic"))
-	require.Equal(t, 1, testAccept(t, true, "selftest/server"))
+	require.Equal(t, 1, testAccept(t, true, "selftest"))
 }
 
 func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
@@ -121,12 +118,14 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	uvCache := getUVDefaultCacheDir(t)
 	t.Setenv("UV_CACHE_DIR", uvCache)
 
+	ctx := context.Background()
 	cloudEnv := os.Getenv("CLOUD_ENV")
 
 	if cloudEnv == "" {
 		defaultServer := testserver.New(t)
 		AddHandlers(defaultServer)
-		t.Setenv("DATABRICKS_DEFAULT_HOST", defaultServer.URL)
+		// Redirect API access to local server:
+		t.Setenv("DATABRICKS_HOST", defaultServer.URL)
 
 		homeDir := t.TempDir()
 		// Do not read user's ~/.databrickscfg
@@ -149,12 +148,27 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	// do it last so that full paths match first:
 	repls.SetPath(buildDir, "[BUILD_DIR]")
 
+	var config databricks.Config
+	if cloudEnv == "" {
+		// use fake token for local tests
+		config = databricks.Config{Token: "dbapi1234"}
+	} else {
+		// non-local tests rely on environment variables
+		config = databricks.Config{}
+	}
+	workspaceClient, err := databricks.NewWorkspaceClient(&config)
+	require.NoError(t, err)
+
+	user, err := workspaceClient.CurrentUser.Me(ctx)
+	require.NoError(t, err)
+	require.NotNil(t, user)
+	testdiff.PrepareReplacementsUser(t, &repls, *user)
+	testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
+	testdiff.PrepareReplacementsUUID(t, &repls)
 	testdiff.PrepareReplacementsDevVersion(t, &repls)
 	testdiff.PrepareReplacementSdkVersion(t, &repls)
 	testdiff.PrepareReplacementsGoVersion(t, &repls)
 
-	repls.Repls = append(repls.Repls, testdiff.Replacement{Old: regexp.MustCompile("dbapi[0-9a-f]+"), New: "[DATABRICKS_TOKEN]"})
-
 	testDirs := getTests(t)
 	require.NotEmpty(t, testDirs)
 
@@ -166,7 +180,8 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	}
 
 	for _, dir := range testDirs {
-		t.Run(dir, func(t *testing.T) {
+		testName := strings.ReplaceAll(dir, "\\", "/")
+		t.Run(testName, func(t *testing.T) {
 			if !InprocessMode {
 				t.Parallel()
 			}
@@ -188,8 +203,7 @@ func getTests(t *testing.T) []string {
 		name := filepath.Base(path)
 		if name == EntryPointScript {
 			// Presence of 'script' marks a test case in this directory
-			testName := filepath.ToSlash(filepath.Dir(path))
-			testDirs = append(testDirs, testName)
+			testDirs = append(testDirs, filepath.Dir(path))
 		}
 		return nil
 	})
@@ -225,6 +239,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	}
 
 	repls.SetPathWithParents(tmpDir, "[TMPDIR]")
+	repls.Repls = append(repls.Repls, config.Repls...)
 
 	scriptContents := readMergedScriptContents(t, dir)
 	testutil.WriteFile(t, filepath.Join(tmpDir, EntryPointScript), scriptContents)
@@ -238,79 +253,38 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	cmd := exec.Command(args[0], args[1:]...)
 	cmd.Env = os.Environ()
 
-	var workspaceClient *databricks.WorkspaceClient
-	var user iam.User
-
 	// Start a new server with a custom configuration if the acceptance test
 	// specifies a custom server stubs.
 	var server *testserver.Server
 
-	if cloudEnv == "" {
-		// Start a new server for this test if either:
-		// 1. A custom server spec is defined in the test configuration.
-		// 2. The test is configured to record requests and assert on them. We need
-		//    a duplicate of the default server to record requests because the default
-		//    server otherwise is a shared resource.
+	// Start a new server for this test if either:
+	// 1. A custom server spec is defined in the test configuration.
+	// 2. The test is configured to record requests and assert on them. We need
+	//    a duplicate of the default server to record requests because the default
+	//    server otherwise is a shared resource.
+	if len(config.Server) > 0 || config.RecordRequests {
+		server = testserver.New(t)
+		server.RecordRequests = config.RecordRequests
+		server.IncludeRequestHeaders = config.IncludeRequestHeaders
 
-		databricksLocalHost := os.Getenv("DATABRICKS_DEFAULT_HOST")
-
-		if len(config.Server) > 0 || config.RecordRequests {
-			server = testserver.New(t)
-			server.RecordRequests = config.RecordRequests
-			server.IncludeRequestHeaders = config.IncludeRequestHeaders
-
-			for _, stub := range config.Server {
-				require.NotEmpty(t, stub.Pattern)
-				items := strings.Split(stub.Pattern, " ")
-				require.Len(t, items, 2)
-				server.Handle(items[0], items[1], func(fakeWorkspace *testserver.FakeWorkspace, req *http.Request) (any, int) {
-					statusCode := http.StatusOK
-					if stub.Response.StatusCode != 0 {
-						statusCode = stub.Response.StatusCode
-					}
-					return stub.Response.Body, statusCode
-				})
-			}
-
-			// The earliest handlers take precedence, add default handlers last
+		// If no custom server stubs are defined, add the default handlers.
+		if len(config.Server) == 0 {
 			AddHandlers(server)
-			databricksLocalHost = server.URL
 		}
 
-		// Each local test should use a new token that will result into a new fake workspace,
-		// so that test don't interfere with each other.
-		tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "")
-		config := databricks.Config{
-			Host:  databricksLocalHost,
-			Token: "dbapi" + tokenSuffix,
+		for _, stub := range config.Server {
+			require.NotEmpty(t, stub.Pattern)
+			server.Handle(stub.Pattern, func(fakeWorkspace *testserver.FakeWorkspace, req *http.Request) (any, int) {
+				statusCode := http.StatusOK
+				if stub.Response.StatusCode != 0 {
+					statusCode = stub.Response.StatusCode
+				}
+				return stub.Response.Body, statusCode
+			})
 		}
-		workspaceClient, err = databricks.NewWorkspaceClient(&config)
-		require.NoError(t, err)
-
-		cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+config.Host)
-		cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+config.Token)
-
-		// For the purposes of replacements, use testUser.
-		// Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement:
-		user = testUser
-	} else {
-		// Use whatever authentication mechanism is configured by the test runner.
-		workspaceClient, err = databricks.NewWorkspaceClient(&databricks.Config{})
-		require.NoError(t, err)
-		pUser, err := workspaceClient.CurrentUser.Me(context.Background())
-		require.NoError(t, err, "Failed to get current user")
-		user = *pUser
+		cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+server.URL)
 	}
 
-	testdiff.PrepareReplacementsUser(t, &repls, user)
-	testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
-
-	// Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME]
-	testdiff.PrepareReplacementsUUID(t, &repls)
-
-	// User replacements come last:
-	repls.Repls = append(repls.Repls, config.Repls...)
-
 	if coverDir != "" {
 		// Creating individual coverage directory for each test, because writing to the same one
 		// results in sporadic failures like this one (only if tests are running in parallel):
@@ -321,6 +295,15 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 		cmd.Env = append(cmd.Env, "GOCOVERDIR="+coverDir)
 	}
 
+	// Each local test should use a new token that will result into a new fake workspace,
+	// so that test don't interfere with each other.
+	if cloudEnv == "" {
+		tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "")
+		token := "dbapi" + tokenSuffix
+		cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+token)
+		repls.Set(token, "[DATABRICKS_TOKEN]")
+	}
+
 	// Write combined output to a file
 	out, err := os.Create(filepath.Join(tmpDir, "output.txt"))
 	require.NoError(t, err)
@@ -337,7 +320,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 
 		for _, req := range server.Requests {
 			reqJson, err := json.MarshalIndent(req, "", "  ")
-			require.NoErrorf(t, err, "Failed to indent: %#v", req)
+			require.NoError(t, err)
 
 			reqJsonWithRepls := repls.Replace(string(reqJson))
 			_, err = f.WriteString(reqJsonWithRepls + "\n")

acceptance/auth/bundle_and_profile/output.txt

@@ -13,13 +13,13 @@
 
 === Inside the bundle, profile flag not matching bundle host. Badness: should use profile from flag instead and not fail
 >>> errcode [CLI] current-user me -p profile_name
-Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_TARGET]
+Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_URL]
 
 Exit code: 1
 
 === Inside the bundle, target and not matching profile
 >>> errcode [CLI] current-user me -t dev -p profile_name
-Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_TARGET]
+Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_URL]
 
 Exit code: 1
 

acceptance/auth/bundle_and_profile/test.toml

@@ -5,8 +5,4 @@ Badness = "When -p flag is used inside the bundle folder for any CLI commands, C
 # This is a workaround to replace DATABRICKS_URL with DATABRICKS_HOST
 [[Repls]]
 Old='DATABRICKS_HOST'
-New='DATABRICKS_TARGET'
-
-[[Repls]]
-Old='DATABRICKS_URL'
-New='DATABRICKS_TARGET'
+New='DATABRICKS_URL'

acceptance/auth/credentials/basic/out.requests.txt

@@ -1,12 +0,0 @@
-{
-  "headers": {
-    "Authorization": [
-      "Basic [ENCODED_AUTH]"
-    ],
-    "User-Agent": [
-      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/basic"
-    ]
-  },
-  "method": "GET",
-  "path": "/api/2.0/preview/scim/v2/Me"
-}

acceptance/auth/credentials/basic/output.txt

@@ -1,4 +0,0 @@
-{
-  "id":"[USERID]",
-  "userName":"[USERNAME]"
-}

acceptance/auth/credentials/basic/script

@@ -1,8 +0,0 @@
-# Unset the token which is configured by default
-# in acceptance tests
-export DATABRICKS_TOKEN=""
-
-export DATABRICKS_USERNAME=username
-export DATABRICKS_PASSWORD=password
-
-$CLI current-user me

acceptance/auth/credentials/basic/test.toml

@@ -1,4 +0,0 @@
-# "username:password" in base64 is dXNlcm5hbWU6cGFzc3dvcmQ=, expect to see this in Authorization header
-[[Repls]]
-Old = "dXNlcm5hbWU6cGFzc3dvcmQ="
-New = "[ENCODED_AUTH]"

acceptance/auth/credentials/oauth/out.requests.txt

@@ -1,34 +0,0 @@
-{
-  "headers": {
-    "User-Agent": [
-      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS]"
-    ]
-  },
-  "method": "GET",
-  "path": "/oidc/.well-known/oauth-authorization-server"
-}
-{
-  "headers": {
-    "Authorization": [
-      "Basic [ENCODED_AUTH]"
-    ],
-    "User-Agent": [
-      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS]"
-    ]
-  },
-  "method": "POST",
-  "path": "/oidc/v1/token",
-  "raw_body": "grant_type=client_credentials\u0026scope=all-apis"
-}
-{
-  "headers": {
-    "Authorization": [
-      "Bearer oauth-token"
-    ],
-    "User-Agent": [
-      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/oauth-m2m"
-    ]
-  },
-  "method": "GET",
-  "path": "/api/2.0/preview/scim/v2/Me"
-}

acceptance/auth/credentials/oauth/output.txt

@@ -1,4 +0,0 @@
-{
-  "id":"[USERID]",
-  "userName":"[USERNAME]"
-}

acceptance/auth/credentials/oauth/script

@@ -1,8 +0,0 @@
-# Unset the token which is configured by default
-# in acceptance tests
-export DATABRICKS_TOKEN=""
-
-export DATABRICKS_CLIENT_ID=client_id
-export DATABRICKS_CLIENT_SECRET=client_secret
-
-$CLI current-user me

acceptance/auth/credentials/oauth/test.toml

@@ -1,5 +0,0 @@
-# "client_id:client_secret" in base64 is Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ=, expect to
-# see this in Authorization header
-[[Repls]]
-Old = "Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ="
-New = "[ENCODED_AUTH]"

acceptance/auth/credentials/pat/out.requests.txt

@@ -1,12 +0,0 @@
-{
-  "headers": {
-    "Authorization": [
-      "Bearer dapi1234"
-    ],
-    "User-Agent": [
-      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/pat"
-    ]
-  },
-  "method": "GET",
-  "path": "/api/2.0/preview/scim/v2/Me"
-}

acceptance/auth/credentials/pat/output.txt

@@ -1,4 +0,0 @@
-{
-  "id":"[USERID]",
-  "userName":"[USERNAME]"
-}

acceptance/auth/credentials/pat/script

@@ -1,3 +0,0 @@
-export DATABRICKS_TOKEN=dapi1234
-
-$CLI current-user me

acceptance/auth/credentials/test.toml

@@ -1,20 +0,0 @@
-LocalOnly = true
-
-RecordRequests = true
-IncludeRequestHeaders = ["Authorization", "User-Agent"]
-
-[[Repls]]
-Old = '(linux|darwin|windows)'
-New = '[OS]'
-
-[[Repls]]
-Old = " upstream/[A-Za-z0-9.-]+"
-New = ""
-
-[[Repls]]
-Old = " upstream-version/[A-Za-z0-9.-]+"
-New = ""
-
-[[Repls]]
-Old = " cicd/[A-Za-z0-9.-]+"
-New = ""

acceptance/cmd_server_test.go

@@ -14,7 +14,11 @@ import (
 
 func StartCmdServer(t *testing.T) *testserver.Server {
 	server := testserver.New(t)
-	server.Handle("GET", "/", func(_ *testserver.FakeWorkspace, r *http.Request) (any, int) {
+
+	// {$} is a wildcard that only matches the end of the URL. We explicitly use
+	// /{$} to disambiguate it from the generic handler for '/' which is used to
+	// identify unhandled API endpoints in the test server.
+	server.Handle("/{$}", func(w *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		q := r.URL.Query()
 		args := strings.Split(q.Get("args"), " ")
 

acceptance/selftest/server/out.requests.txt

@@ -1,8 +0,0 @@
-{
-  "method": "GET",
-  "path": "/api/2.0/preview/scim/v2/Me"
-}
-{
-  "method": "GET",
-  "path": "/custom/endpoint"
-}

acceptance/selftest/server/output.txt

@@ -1,15 +0,0 @@
-
->>> curl -s [DATABRICKS_URL]/api/2.0/preview/scim/v2/Me
-{
-    "id": "[USERID]",
-    "userName": "[USERNAME]"
-}
->>> curl -sD - [DATABRICKS_URL]/custom/endpoint?query=param
-HTTP/1.1 201 Created
-Content-Type: application/json
-Date: (redacted)
-Content-Length: (redacted)
-
-custom
----
-response

acceptance/selftest/server/script

@@ -1,2 +0,0 @@
-trace curl -s $DATABRICKS_HOST/api/2.0/preview/scim/v2/Me
-trace curl -sD - $DATABRICKS_HOST/custom/endpoint?query=param

acceptance/selftest/server/test.toml

@@ -1,18 +0,0 @@
-LocalOnly = true
-RecordRequests = true
-
-[[Server]]
-Pattern = "GET /custom/endpoint"
-Response.Body = '''custom
----
-response
-'''
-Response.StatusCode = 201
-
-[[Repls]]
-Old = 'Date: .*'
-New = 'Date: (redacted)'
-
-[[Repls]]
-Old = 'Content-Length: [0-9]*'
-New = 'Content-Length: (redacted)'

acceptance/server_test.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/iam"
-	"github.com/gorilla/mux"
 
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -17,13 +16,8 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
-var testUser = iam.User{
-	Id:       "1000012345",
-	UserName: "tester@databricks.com",
-}
-
 func AddHandlers(server *testserver.Server) {
-	server.Handle("GET", "/api/2.0/policies/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.0/policies/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return compute.ListPoliciesResponse{
 			Policies: []compute.Policy{
 				{
@@ -38,7 +32,7 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("GET", "/api/2.0/instance-pools/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.0/instance-pools/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return compute.ListInstancePools{
 			InstancePools: []compute.InstancePoolAndStats{
 				{
@@ -49,7 +43,7 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("GET", "/api/2.1/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.1/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return compute.ListClustersResponse{
 			Clusters: []compute.ClusterDetails{
 				{
@@ -64,17 +58,20 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		return testUser, http.StatusOK
+	server.Handle("GET /api/2.0/preview/scim/v2/Me", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+		return iam.User{
+			Id:       "1000012345",
+			UserName: "tester@databricks.com",
+		}, http.StatusOK
 	})
 
-	server.Handle("GET", "/api/2.0/workspace/get-status", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.0/workspace/get-status", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		path := r.URL.Query().Get("path")
 
 		return fakeWorkspace.WorkspaceGetStatus(path)
 	})
 
-	server.Handle("POST", "/api/2.0/workspace/mkdirs", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST /api/2.0/workspace/mkdirs", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		request := workspace.Mkdirs{}
 		decoder := json.NewDecoder(r.Body)
 
@@ -86,13 +83,13 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.WorkspaceMkdirs(request)
 	})
 
-	server.Handle("GET", "/api/2.0/workspace/export", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.0/workspace/export", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		path := r.URL.Query().Get("path")
 
 		return fakeWorkspace.WorkspaceExport(path)
 	})
 
-	server.Handle("POST", "/api/2.0/workspace/delete", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST /api/2.0/workspace/delete", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		path := r.URL.Query().Get("path")
 		recursiveStr := r.URL.Query().Get("recursive")
 		var recursive bool
@@ -106,9 +103,8 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.WorkspaceDelete(path, recursive)
 	})
 
-	server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		vars := mux.Vars(r)
-		path := vars["path"]
+	server.Handle("POST /api/2.0/workspace-files/import-file/{path}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+		path := r.PathValue("path")
 
 		body := new(bytes.Buffer)
 		_, err := body.ReadFrom(r.Body)
@@ -119,15 +115,14 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.WorkspaceFilesImportFile(path, body.Bytes())
 	})
 
-	server.Handle("GET", "/api/2.1/unity-catalog/current-metastore-assignment", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return catalog.MetastoreAssignment{
 			DefaultCatalogName: "main",
 		}, http.StatusOK
 	})
 
-	server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		vars := mux.Vars(r)
-		objectId := vars["objectId"]
+	server.Handle("GET /api/2.0/permissions/directories/{objectId}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+		objectId := r.PathValue("objectId")
 
 		return workspace.WorkspaceObjectPermissions{
 			ObjectId:   objectId,
@@ -145,7 +140,7 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("POST", "/api/2.1/jobs/create", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST /api/2.1/jobs/create", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		request := jobs.CreateJob{}
 		decoder := json.NewDecoder(r.Body)
 
@@ -157,31 +152,15 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.JobsCreate(request)
 	})
 
-	server.Handle("GET", "/api/2.1/jobs/get", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.1/jobs/get", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		jobId := r.URL.Query().Get("job_id")
 
 		return fakeWorkspace.JobsGet(jobId)
 	})
 
-	server.Handle("GET", "/api/2.1/jobs/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET /api/2.1/jobs/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return fakeWorkspace.JobsList()
 	})
-
-	server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		return map[string]string{
-			"authorization_endpoint": server.URL + "oidc/v1/authorize",
-			"token_endpoint":         server.URL + "/oidc/v1/token",
-		}, http.StatusOK
-	})
-
-	server.Handle("POST", "/oidc/v1/token", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		return map[string]string{
-			"access_token": "oauth-token",
-			"expires_in":   "3600",
-			"scope":        "all-apis",
-			"token_type":   "Bearer",
-		}, http.StatusOK
-	})
 }
 
 func internalError(err error) (any, int) {

go.mod

@@ -12,7 +12,6 @@ require (
 	github.com/databricks/databricks-sdk-go v0.57.0 // Apache 2.0
 	github.com/fatih/color v1.18.0 // MIT
 	github.com/google/uuid v1.6.0 // BSD-3-Clause
-	github.com/gorilla/mux v1.8.1 // BSD 3-Clause
 	github.com/hashicorp/go-version v1.7.0 // MPL 2.0
 	github.com/hashicorp/hc-install v0.9.1 // MPL 2.0
 	github.com/hashicorp/terraform-exec v0.22.0 // MPL 2.0

go.sum

@@ -97,8 +97,6 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfF
 github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
 github.com/googleapis/gax-go/v2 v2.12.4 h1:9gWcmF85Wvq4ryPFvGFaOgPIs1AQX0d0bcbGw4Z96qg=
 github.com/googleapis/gax-go/v2 v2.12.4/go.mod h1:KYEYLorsnIGDi/rPC8b5TdlB9kbKoFubselGIoBMCwI=
-github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
-github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
 github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
 github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
 github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB11/k=

libs/testserver/server.go

@@ -9,8 +9,6 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/gorilla/mux"
-
 	"github.com/stretchr/testify/assert"
 
 	"github.com/databricks/cli/internal/testutil"
@@ -19,7 +17,7 @@ import (
 
 type Server struct {
 	*httptest.Server
-	Router *mux.Router
+	Mux *http.ServeMux
 
 	t testutil.TestingT
 
@@ -36,25 +34,26 @@ type Request struct {
 	Headers http.Header `json:"headers,omitempty"`
 	Method  string      `json:"method"`
 	Path    string      `json:"path"`
-	Body    any         `json:"body,omitempty"`
-	RawBody string      `json:"raw_body,omitempty"`
+	Body    any         `json:"body"`
 }
 
 func New(t testutil.TestingT) *Server {
-	router := mux.NewRouter()
-	server := httptest.NewServer(router)
+	mux := http.NewServeMux()
+	server := httptest.NewServer(mux)
 	t.Cleanup(server.Close)
 
 	s := &Server{
 		Server:         server,
-		Router:         router,
+		Mux:            mux,
 		t:              t,
 		mu:             &sync.Mutex{},
 		fakeWorkspaces: map[string]*FakeWorkspace{},
 	}
 
-	// Set up the not found handler as fallback
-	router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	// The server resolves conflicting handlers by using the one with higher
+	// specificity. This handler is the least specific, so it will be used as a
+	// fallback when no other handlers match.
+	s.Handle("/", func(fakeWorkspace *FakeWorkspace, r *http.Request) (any, int) {
 		pattern := r.Method + " " + r.URL.Path
 
 		t.Errorf(`
@@ -75,22 +74,9 @@ Response.StatusCode = <response status-code here>
 
 `, pattern, pattern)
 
-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(http.StatusNotImplemented)
-
-		resp := apierr.APIError{
+		return apierr.APIError{
 			Message: "No stub found for pattern: " + pattern,
-		}
-
-		respBytes, err := json.Marshal(resp)
-		if err != nil {
-			t.Errorf("JSON encoding error: %s", err)
-			respBytes = []byte("{\"message\": \"JSON encoding error\"}")
-		}
-
-		if _, err := w.Write(respBytes); err != nil {
-			t.Errorf("Response write error: %s", err)
-		}
+		}, http.StatusNotImplemented
 	})
 
 	return s
@@ -98,8 +84,8 @@ Response.StatusCode = <response status-code here>
 
 type HandlerFunc func(fakeWorkspace *FakeWorkspace, req *http.Request) (resp any, statusCode int)
 
-func (s *Server) Handle(method, path string, handler HandlerFunc) {
-	s.Router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
+func (s *Server) Handle(pattern string, handler HandlerFunc) {
+	s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
 		// For simplicity we process requests sequentially. It's fast enough because
 		// we don't do any IO except reading and writing request/response bodies.
 		s.mu.Lock()
@@ -133,19 +119,13 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) {
 				}
 			}
 
-			req := Request{
+			s.Requests = append(s.Requests, Request{
 				Headers: headers,
 				Method:  r.Method,
 				Path:    r.URL.Path,
-			}
+				Body:    json.RawMessage(body),
+			})
 
-			if json.Valid(body) {
-				req.Body = json.RawMessage(body)
-			} else {
-				req.RawBody = string(body)
-			}
-
-			s.Requests = append(s.Requests, req)
 		}
 
 		w.Header().Set("Content-Type", "application/json")
@@ -169,7 +149,7 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
-	}).Methods(method)
+	})
 }
 
 func getToken(r *http.Request) string {