Allow test servers to return errors responses (#2291)

## Changes
The APIs at Databricks when returning a non `200` status code will
return a response body of the format:
```
{
  "error_code": "Error code",
  "message": "Human-readable error message."
}
```

This PR adds the ability to stub non-200 status codes in the test
server, allowing us to mock API errors from Databricks.
## Tests
New test

acceptance/acceptance_test.go

@@ -261,8 +261,12 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 
 		for _, stub := range config.Server {
 			require.NotEmpty(t, stub.Pattern)
-			server.Handle(stub.Pattern, func(req *http.Request) (resp any, err error) {
-				return stub.Response.Body, nil
+			server.Handle(stub.Pattern, func(req *http.Request) (any, int) {
+				statusCode := http.StatusOK
+				if stub.Response.StatusCode != 0 {
+					statusCode = stub.Response.StatusCode
+				}
+				return stub.Response.Body, statusCode
 			})
 		}
 		cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+server.URL)

acceptance/cmd_server_test.go

@@ -15,7 +15,7 @@ import (
 func StartCmdServer(t *testing.T) *testserver.Server {
 	server := testserver.New(t)
 
-	server.Handle("/", func(r *http.Request) (any, error) {
+	server.Handle("/", func(r *http.Request) (any, int) {
 		q := r.URL.Query()
 		args := strings.Split(q.Get("args"), " ")
 
@@ -40,7 +40,7 @@ func StartCmdServer(t *testing.T) *testserver.Server {
 			exitcode = 1
 		}
 		result["exitcode"] = exitcode
-		return result, nil
+		return result, http.StatusOK
 	})
 	return server
 }

acceptance/config_test.go

@@ -57,7 +57,8 @@ type ServerStub struct {
 
 	// The response body to return.
 	Response struct {
-		Body string
+		Body       string
+		StatusCode int
 	}
 }
 

acceptance/server_test.go

@@ -11,7 +11,7 @@ import (
 )
 
 func AddHandlers(server *testserver.Server) {
-	server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, int) {
 		return compute.ListPoliciesResponse{
 			Policies: []compute.Policy{
 				{
@@ -23,21 +23,21 @@ func AddHandlers(server *testserver.Server) {
 					Name:     "some-test-cluster-policy",
 				},
 			},
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/instance-pools/list", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/instance-pools/list", func(r *http.Request) (any, int) {
 		return compute.ListInstancePools{
 			InstancePools: []compute.InstancePoolAndStats{
 				{
 					InstancePoolName: "some-test-instance-pool",
 					InstancePoolId:   "1234",
 				},
 			},
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.1/clusters/list", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.1/clusters/list", func(r *http.Request) (any, int) {
 		return compute.ListClustersResponse{
 			Clusters: []compute.ClusterDetails{
 				{
@@ -49,32 +49,32 @@ func AddHandlers(server *testserver.Server) {
 					ClusterId:   "9876",
 				},
 			},
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/preview/scim/v2/Me", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/preview/scim/v2/Me", func(r *http.Request) (any, int) {
 		return iam.User{
 			Id:       "1000012345",
 			UserName: "tester@databricks.com",
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/workspace/get-status", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/workspace/get-status", func(r *http.Request) (any, int) {
 		return workspace.ObjectInfo{
 			ObjectId:   1001,
 			ObjectType: "DIRECTORY",
 			Path:       "",
 			ResourceId: "1001",
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(r *http.Request) (any, int) {
 		return catalog.MetastoreAssignment{
 			DefaultCatalogName: "main",
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/permissions/directories/1001", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/permissions/directories/1001", func(r *http.Request) (any, int) {
 		return workspace.WorkspaceObjectPermissions{
 			ObjectId:   "1001",
 			ObjectType: "DIRECTORY",
@@ -88,10 +88,10 @@ func AddHandlers(server *testserver.Server) {
 					},
 				},
 			},
-		}, nil
+		}, http.StatusOK
 	})
 
-	server.Handle("POST /api/2.0/workspace/mkdirs", func(r *http.Request) (any, error) {
-		return "{}", nil
+	server.Handle("POST /api/2.0/workspace/mkdirs", func(r *http.Request) (any, int) {
+		return "{}", http.StatusOK
 	})
 }

acceptance/workspace/jobs/create-error/out.requests.txt

@@ -0,0 +1 @@
+{"method":"POST","path":"/api/2.1/jobs/create","body":{"name":"abc"}}

acceptance/workspace/jobs/create-error/output.txt

@@ -0,0 +1,5 @@
+
+>>> [CLI] jobs create --json {"name":"abc"}
+Error: Invalid access token.
+
+Exit code: 1

acceptance/workspace/jobs/create-error/script

@@ -0,0 +1 @@
+trace $CLI jobs create --json '{"name":"abc"}'

acceptance/workspace/jobs/create-error/test.toml

@@ -0,0 +1,12 @@
+LocalOnly = true  # request recording currently does not work with cloud environment
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /api/2.1/jobs/create"
+Response.Body = '''
+{
+    "error_code": "PERMISSION_DENIED",
+    "message": "Invalid access token."
+}
+'''
+Response.StatusCode = 403

libs/testserver/server.go

@@ -40,15 +40,11 @@ func New(t testutil.TestingT) *Server {
 	}
 }
 
-type HandlerFunc func(req *http.Request) (resp any, err error)
+type HandlerFunc func(req *http.Request) (resp any, statusCode int)
 
 func (s *Server) Handle(pattern string, handler HandlerFunc) {
 	s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
-		resp, err := handler(r)
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
+		resp, statusCode := handler(r)
 
 		if s.RecordRequests {
 			body, err := io.ReadAll(r.Body)
@@ -63,9 +59,10 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) {
 		}
 
 		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(statusCode)
 
 		var respBytes []byte
-
+		var err error
 		respString, ok := resp.(string)
 		if ok {
 			respBytes = []byte(respString)