acc: Simplify writing handlers; support headers in responses (#2338)

## Changes

Handlers now receive testserver.Request and return any which could be
- string or []byte (returns it as is but sets content-type to json or
plain text depending on content)
- struct (encodes it as json and sets content-type to json)
- testserver.Response (full control over status and headers)

Note if testserver.Response is returned from the handler, it's Body
attribute can still be an object. In that case, it'll be serialized and
appropriate content-type header will be added.

The config is now using the same testserver.Response struct, the same
logic applies both configured responses and responses returned from
handlers.

As a result, one can set headers both in Golang handlers and in
test.toml.

This also fixes a bug with RecordRequest not seeing the body if it was
already consumed by the handler.

## Tests
- Existing rests. 
- acceptance/selftest/server is extended to set response header.

acceptance/acceptance_test.go

@@ -7,7 +7,6 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -267,12 +266,8 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 				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) {
+				server.Handle(items[0], items[1], func(req testserver.Request) any {
-					statusCode := http.StatusOK
+					return stub.Response
-					if stub.Response.StatusCode != 0 {
-						statusCode = stub.Response.StatusCode
-					}
-					return stub.Response.Body, statusCode
 				})
 			}
 

acceptance/bundle/debug/out.stderr.parallel.txt

@@ -9,7 +9,7 @@
 10:07:59 Debug: ApplyReadOnly pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:folder_permissions
 10:07:59 Debug: ApplyReadOnly pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:validate_sync_patterns
 10:07:59 Debug: Path /Workspace/Users/[USERNAME]/.bundle/debug/default/files has type directory (ID: 0) pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync
-10:07:59 Debug: non-retriable error:  pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true
+10:07:59 Debug: non-retriable error: Workspace path not found pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true
-< {} pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true
+< HTTP/0.0 000 OK pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true
-< {} pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true
+< } pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true
 < } pid=12345 mutator=validate mutator (read-only)=parallel mutator (read-only)=validate:files_to_sync sdk=true

acceptance/bundle/debug/out.stderr.txt

@@ -79,11 +79,12 @@
 10:07:59 Debug: Apply pid=12345 mutator=validate
 10:07:59 Debug: GET /api/2.0/workspace/get-status?path=/Workspace/Users/[USERNAME]/.bundle/debug/default/files
 < HTTP/1.1 404 Not Found
+< {
+<   "message": "Workspace path not found"
 10:07:59 Debug: POST /api/2.0/workspace/mkdirs
 > {
 >   "path": "/Workspace/Users/[USERNAME]/.bundle/debug/default/files"
 > }
-< HTTP/1.1 200 OK
 10:07:59 Debug: GET /api/2.0/workspace/get-status?path=/Workspace/Users/[USERNAME]/.bundle/debug/default/files
 < HTTP/1.1 200 OK
 < {

acceptance/cmd_server_test.go

@@ -1,8 +1,8 @@
 package acceptance_test
 
 import (
+	"context"
 	"encoding/json"
-	"net/http"
 	"os"
 	"strings"
 	"testing"
@@ -14,7 +14,7 @@ import (
 
 func StartCmdServer(t *testing.T) *testserver.Server {
 	server := testserver.New(t)
-	server.Handle("GET", "/", func(_ *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/", func(r testserver.Request) any {
 		q := r.URL.Query()
 		args := strings.Split(q.Get("args"), " ")
 
@@ -27,7 +27,7 @@ func StartCmdServer(t *testing.T) *testserver.Server {
 
 		defer Chdir(t, q.Get("cwd"))()
 
-		c := testcli.NewRunner(t, r.Context(), args...)
+		c := testcli.NewRunner(t, context.Background(), args...)
 		c.Verbose = false
 		stdout, stderr, err := c.Run()
 		result := map[string]any{
@@ -39,7 +39,7 @@ func StartCmdServer(t *testing.T) *testserver.Server {
 			exitcode = 1
 		}
 		result["exitcode"] = exitcode
-		return result, http.StatusOK
+		return result
 	})
 	return server
 }

acceptance/config_test.go

@@ -10,6 +10,7 @@ import (
 	"dario.cat/mergo"
 	"github.com/BurntSushi/toml"
 	"github.com/databricks/cli/libs/testdiff"
+	"github.com/databricks/cli/libs/testserver"
 	"github.com/stretchr/testify/require"
 )
 
@@ -56,10 +57,7 @@ type ServerStub struct {
 	Pattern string
 
 	// The response body to return.
-	Response struct {
+	Response testserver.Response
-		Body       string
-		StatusCode int
-	}
 }
 
 // FindConfigs finds all the config relevant for this test,

acceptance/selftest/server/out.requests.txt

@@ -6,3 +6,7 @@
   "method": "GET",
   "path": "/custom/endpoint"
 }
+{
+  "method": "GET",
+  "path": "/api/2.0/workspace/get-status"
+}

acceptance/selftest/server/output.txt

@@ -6,10 +6,16 @@
 }
 >>> curl -sD - [DATABRICKS_URL]/custom/endpoint?query=param
 HTTP/1.1 201 Created
-Content-Type: application/json
+X-Custom-Header: hello
 Date: (redacted)
 Content-Length: (redacted)
+Content-Type: text/plain; charset=utf-8
 
 custom
 ---
 response
+
+>>> errcode [CLI] workspace get-status /a/b/c
+Error: Workspace path not found
+
+Exit code: 1

acceptance/selftest/server/script

@@ -1,2 +1,4 @@
 trace curl -s $DATABRICKS_HOST/api/2.0/preview/scim/v2/Me
 trace curl -sD - $DATABRICKS_HOST/custom/endpoint?query=param
+
+trace errcode $CLI workspace get-status /a/b/c

acceptance/selftest/server/test.toml

@@ -12,6 +12,8 @@ Response.Body = '''custom
 response
 '''
 Response.StatusCode = 201
+[Server.Response.Headers]
+"X-Custom-Header" = ["hello"]
 
 [[Repls]]
 Old = 'Date: .*'

acceptance/server_test.go

@@ -1,14 +1,12 @@
 package acceptance_test
 
 import (
-	"bytes"
 	"encoding/json"
 	"fmt"
 	"net/http"
 
 	"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"
@@ -23,7 +21,7 @@ var testUser = iam.User{
 }
 
 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(req testserver.Request) any {
 		return compute.ListPoliciesResponse{
 			Policies: []compute.Policy{
 				{
@@ -35,10 +33,10 @@ func AddHandlers(server *testserver.Server) {
 					Name:     "some-test-cluster-policy",
 				},
 			},
-		}, 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(req testserver.Request) any {
 		return compute.ListInstancePools{
 			InstancePools: []compute.InstancePoolAndStats{
 				{
@@ -46,10 +44,10 @@ func AddHandlers(server *testserver.Server) {
 					InstancePoolId:   "1234",
 				},
 			},
-		}, 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(req testserver.Request) any {
 		return compute.ListClustersResponse{
 			Clusters: []compute.ClusterDetails{
 				{
@@ -61,74 +59,57 @@ func AddHandlers(server *testserver.Server) {
 					ClusterId:   "9876",
 				},
 			},
-		}, http.StatusOK
+		}
 	})
 
-	server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any {
-		return testUser, http.StatusOK
+		return testUser
 	})
 
-	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(req testserver.Request) any {
-		path := r.URL.Query().Get("path")
+		path := req.URL.Query().Get("path")
-
+		return req.Workspace.WorkspaceGetStatus(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(req testserver.Request) any {
-		request := workspace.Mkdirs{}
+		var request workspace.Mkdirs
-		decoder := json.NewDecoder(r.Body)
+		if err := json.Unmarshal(req.Body, &request); err != nil {
-
+			return testserver.Response{
-		err := decoder.Decode(&request)
+				Body:       fmt.Sprintf("internal error: %s", err),
-		if err != nil {
+				StatusCode: http.StatusInternalServerError,
-			return internalError(err)
+			}
 		}
 
-		return fakeWorkspace.WorkspaceMkdirs(request)
+		req.Workspace.WorkspaceMkdirs(request)
+		return ""
 	})
 
-	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(req testserver.Request) any {
-		path := r.URL.Query().Get("path")
+		path := req.URL.Query().Get("path")
-
+		return req.Workspace.WorkspaceExport(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(req testserver.Request) any {
-		path := r.URL.Query().Get("path")
+		path := req.URL.Query().Get("path")
-		recursiveStr := r.URL.Query().Get("recursive")
+		recursive := req.URL.Query().Get("recursive") == "true"
-		var recursive bool
+		req.Workspace.WorkspaceDelete(path, recursive)
-
+		return ""
-		if recursiveStr == "true" {
-			recursive = true
-		} else {
-			recursive = false
-		}
-
-		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) {
+	server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(req testserver.Request) any {
-		vars := mux.Vars(r)
+		path := req.Vars["path"]
-		path := vars["path"]
+		req.Workspace.WorkspaceFilesImportFile(path, req.Body)
-
+		return ""
-		body := new(bytes.Buffer)
-		_, err := body.ReadFrom(r.Body)
-		if err != nil {
-			return internalError(err)
-		}
-
-		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(req testserver.Request) any {
 		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) {
+	server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(req testserver.Request) any {
-		vars := mux.Vars(r)
+		objectId := req.Vars["objectId"]
-		objectId := vars["objectId"]
-
 		return workspace.WorkspaceObjectPermissions{
 			ObjectId:   objectId,
 			ObjectType: "DIRECTORY",
@@ -142,48 +123,43 @@ 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(req testserver.Request) any {
-		request := jobs.CreateJob{}
+		var request jobs.CreateJob
-		decoder := json.NewDecoder(r.Body)
+		if err := json.Unmarshal(req.Body, &request); err != nil {
-
+			return testserver.Response{
-		err := decoder.Decode(&request)
+				Body:       fmt.Sprintf("internal error: %s", err),
-		if err != nil {
+				StatusCode: 500,
-			return internalError(err)
+			}
 		}
 
-		return fakeWorkspace.JobsCreate(request)
+		return req.Workspace.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(req testserver.Request) any {
-		jobId := r.URL.Query().Get("job_id")
+		jobId := req.URL.Query().Get("job_id")
-
+		return req.Workspace.JobsGet(jobId)
-		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(req testserver.Request) any {
-		return fakeWorkspace.JobsList()
+		return req.Workspace.JobsList()
 	})
 
-	server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(_ testserver.Request) any {
 		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) {
+	server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any {
 		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) {
-	return fmt.Errorf("internal error: %w", err), http.StatusInternalServerError
-}

libs/testserver/fake_workspace.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"net/http"
 	"sort"
 	"strconv"
 	"strings"
@@ -33,40 +32,39 @@ func NewFakeWorkspace() *FakeWorkspace {
 	}
 }
 
-func (s *FakeWorkspace) WorkspaceGetStatus(path string) (workspace.ObjectInfo, int) {
+func (s *FakeWorkspace) WorkspaceGetStatus(path string) Response {
 	if s.directories[path] {
-		return workspace.ObjectInfo{
+		return Response{
-			ObjectType: "DIRECTORY",
+			Body: &workspace.ObjectInfo{
-			Path:       path,
+				ObjectType: "DIRECTORY",
-		}, http.StatusOK
+				Path:       path,
+			},
+		}
 	} else if _, ok := s.files[path]; ok {
-		return workspace.ObjectInfo{
+		return Response{
-			ObjectType: "FILE",
+			Body: &workspace.ObjectInfo{
-			Path:       path,
+				ObjectType: "FILE",
-			Language:   "SCALA",
+				Path:       path,
-		}, http.StatusOK
+				Language:   "SCALA",
+			},
+		}
 	} else {
-		return workspace.ObjectInfo{}, http.StatusNotFound
+		return Response{
+			StatusCode: 404,
+			Body:       map[string]string{"message": "Workspace path not found"},
+		}
 	}
 }
 
-func (s *FakeWorkspace) WorkspaceMkdirs(request workspace.Mkdirs) (string, int) {
+func (s *FakeWorkspace) WorkspaceMkdirs(request workspace.Mkdirs) {
 	s.directories[request.Path] = true
-
-	return "{}", http.StatusOK
 }
 
-func (s *FakeWorkspace) WorkspaceExport(path string) ([]byte, int) {
+func (s *FakeWorkspace) WorkspaceExport(path string) []byte {
-	file := s.files[path]
+	return s.files[path]
-
-	if file == nil {
-		return nil, http.StatusNotFound
-	}
-
-	return file, http.StatusOK
 }
 
-func (s *FakeWorkspace) WorkspaceDelete(path string, recursive bool) (string, int) {
+func (s *FakeWorkspace) WorkspaceDelete(path string, recursive bool) {
 	if !recursive {
 		s.files[path] = nil
 	} else {
@@ -76,28 +74,26 @@ func (s *FakeWorkspace) WorkspaceDelete(path string, recursive bool) (string, in
 			}
 		}
 	}
-
-	return "{}", http.StatusOK
 }
 
-func (s *FakeWorkspace) WorkspaceFilesImportFile(path string, body []byte) (any, int) {
+func (s *FakeWorkspace) WorkspaceFilesImportFile(path string, body []byte) {
 	if !strings.HasPrefix(path, "/") {
 		path = "/" + path
 	}
-
 	s.files[path] = body
-
-	return "{}", http.StatusOK
 }
 
-func (s *FakeWorkspace) JobsCreate(request jobs.CreateJob) (any, int) {
+func (s *FakeWorkspace) JobsCreate(request jobs.CreateJob) Response {
 	jobId := s.nextJobId
 	s.nextJobId++
 
 	jobSettings := jobs.JobSettings{}
 	err := jsonConvert(request, &jobSettings)
 	if err != nil {
-		return internalError(err)
+		return Response{
+			StatusCode: 400,
+			Body:       fmt.Sprintf("Cannot convert request to jobSettings: %s", err),
+		}
 	}
 
 	s.jobs[jobId] = jobs.Job{
@@ -105,32 +101,44 @@ func (s *FakeWorkspace) JobsCreate(request jobs.CreateJob) (any, int) {
 		Settings: &jobSettings,
 	}
 
-	return jobs.CreateResponse{JobId: jobId}, http.StatusOK
+	return Response{
+		Body: jobs.CreateResponse{JobId: jobId},
+	}
 }
 
-func (s *FakeWorkspace) JobsGet(jobId string) (any, int) {
+func (s *FakeWorkspace) JobsGet(jobId string) Response {
 	id := jobId
 
 	jobIdInt, err := strconv.ParseInt(id, 10, 64)
 	if err != nil {
-		return internalError(fmt.Errorf("failed to parse job id: %s", err))
+		return Response{
+			StatusCode: 400,
+			Body:       fmt.Sprintf("Failed to parse job id: %s: %v", err, id),
+		}
 	}
 
 	job, ok := s.jobs[jobIdInt]
 	if !ok {
-		return jobs.Job{}, http.StatusNotFound
+		return Response{
+			StatusCode: 404,
+		}
 	}
 
-	return job, http.StatusOK
+	return Response{
+		Body: job,
+	}
 }
 
-func (s *FakeWorkspace) JobsList() (any, int) {
+func (s *FakeWorkspace) JobsList() Response {
 	list := make([]jobs.BaseJob, 0, len(s.jobs))
 	for _, job := range s.jobs {
 		baseJob := jobs.BaseJob{}
 		err := jsonConvert(job, &baseJob)
 		if err != nil {
-			return internalError(fmt.Errorf("failed to convert job to base job: %w", err))
+			return Response{
+				StatusCode: 400,
+				Body:       fmt.Sprintf("failed to convert job to base job: %s", err),
+			}
 		}
 
 		list = append(list, baseJob)
@@ -141,9 +149,11 @@ func (s *FakeWorkspace) JobsList() (any, int) {
 		return list[i].JobId < list[j].JobId
 	})
 
-	return jobs.ListJobsResponse{
+	return Response{
-		Jobs: list,
+		Body: jobs.ListJobsResponse{
-	}, http.StatusOK
+			Jobs: list,
+		},
+	}
 }
 
 // jsonConvert saves input to a value pointed by output
@@ -163,7 +173,3 @@ func jsonConvert(input, output any) error {
 
 	return nil
 }
-
-func internalError(err error) (string, int) {
-	return fmt.Sprintf("internal error: %s", err), http.StatusInternalServerError
-}

libs/testserver/server.go

@@ -5,14 +5,14 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
+	"reflect"
 	"slices"
 	"strings"
 	"sync"
 
 	"github.com/gorilla/mux"
 
-	"github.com/stretchr/testify/assert"
-
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/databricks-sdk-go/apierr"
 )
@@ -29,10 +29,10 @@ type Server struct {
 	RecordRequests        bool
 	IncludeRequestHeaders []string
 
-	Requests []Request
+	Requests []LoggedRequest
 }
 
-type Request struct {
+type LoggedRequest struct {
 	Headers http.Header `json:"headers,omitempty"`
 	Method  string      `json:"method"`
 	Path    string      `json:"path"`
@@ -40,6 +40,153 @@ type Request struct {
 	RawBody string      `json:"raw_body,omitempty"`
 }
 
+type Request struct {
+	Method    string
+	URL       *url.URL
+	Headers   http.Header
+	Body      []byte
+	Vars      map[string]string
+	Workspace *FakeWorkspace
+}
+
+type Response struct {
+	StatusCode int
+	Headers    http.Header
+	Body       any
+}
+
+type encodedResponse struct {
+	StatusCode int
+	Headers    http.Header
+	Body       []byte
+}
+
+func NewRequest(t testutil.TestingT, r *http.Request, fakeWorkspace *FakeWorkspace) Request {
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		t.Fatalf("Failed to read request body: %s", err)
+	}
+
+	return Request{
+		Method:    r.Method,
+		URL:       r.URL,
+		Headers:   r.Header,
+		Body:      body,
+		Vars:      mux.Vars(r),
+		Workspace: fakeWorkspace,
+	}
+}
+
+func normalizeResponse(t testutil.TestingT, resp any) encodedResponse {
+	result := normalizeResponseBody(t, resp)
+	if result.StatusCode == 0 {
+		result.StatusCode = 200
+	}
+	return result
+}
+
+func normalizeResponseBody(t testutil.TestingT, resp any) encodedResponse {
+	if isNil(resp) {
+		t.Errorf("Handler must not return nil")
+		return encodedResponse{StatusCode: 500}
+	}
+
+	respBytes, ok := resp.([]byte)
+	if ok {
+		return encodedResponse{
+			Body:    respBytes,
+			Headers: getHeaders(respBytes),
+		}
+	}
+
+	respString, ok := resp.(string)
+	if ok {
+		return encodedResponse{
+			Body:    []byte(respString),
+			Headers: getHeaders([]byte(respString)),
+		}
+	}
+
+	respStruct, ok := resp.(Response)
+	if ok {
+		if isNil(respStruct.Body) {
+			return encodedResponse{
+				StatusCode: respStruct.StatusCode,
+				Headers:    respStruct.Headers,
+				Body:       []byte{},
+			}
+		}
+
+		bytesVal, isBytes := respStruct.Body.([]byte)
+		if isBytes {
+			return encodedResponse{
+				StatusCode: respStruct.StatusCode,
+				Headers:    respStruct.Headers,
+				Body:       bytesVal,
+			}
+		}
+
+		stringVal, isString := respStruct.Body.(string)
+		if isString {
+			return encodedResponse{
+				StatusCode: respStruct.StatusCode,
+				Headers:    respStruct.Headers,
+				Body:       []byte(stringVal),
+			}
+		}
+
+		respBytes, err := json.MarshalIndent(respStruct.Body, "", "    ")
+		if err != nil {
+			t.Errorf("JSON encoding error: %s", err)
+			return encodedResponse{
+				StatusCode: 500,
+				Body:       []byte("internal error"),
+			}
+		}
+
+		headers := respStruct.Headers
+		if headers == nil {
+			headers = getJsonHeaders()
+		}
+
+		return encodedResponse{
+			StatusCode: respStruct.StatusCode,
+			Headers:    headers,
+			Body:       respBytes,
+		}
+	}
+
+	respBytes, err := json.MarshalIndent(resp, "", "    ")
+	if err != nil {
+		t.Errorf("JSON encoding error: %s", err)
+		return encodedResponse{
+			StatusCode: 500,
+			Body:       []byte("internal error"),
+		}
+	}
+
+	return encodedResponse{
+		Body:    respBytes,
+		Headers: getJsonHeaders(),
+	}
+}
+
+func getJsonHeaders() http.Header {
+	return map[string][]string{
+		"Content-Type": {"application/json"},
+	}
+}
+
+func getHeaders(value []byte) http.Header {
+	if json.Valid(value) {
+		return getJsonHeaders()
+	} else {
+		return map[string][]string{
+			"Content-Type": {"text/plain"},
+		}
+	}
+}
+
 func New(t testutil.TestingT) *Server {
 	router := mux.NewRouter()
 	server := httptest.NewServer(router)
@@ -96,7 +243,7 @@ Response.StatusCode = <response status-code here>
 	return s
 }
 
-type HandlerFunc func(fakeWorkspace *FakeWorkspace, req *http.Request) (resp any, statusCode int)
+type HandlerFunc func(req Request) any
 
 func (s *Server) Handle(method, path string, handler HandlerFunc) {
 	s.Router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
@@ -117,56 +264,22 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) {
 			fakeWorkspace = s.fakeWorkspaces[token]
 		}
 
-		resp, statusCode := handler(fakeWorkspace, r)
+		request := NewRequest(s.t, r, fakeWorkspace)
-
 		if s.RecordRequests {
-			body, err := io.ReadAll(r.Body)
+			s.Requests = append(s.Requests, getLoggedRequest(request, s.IncludeRequestHeaders))
-			assert.NoError(s.t, err)
-
-			headers := make(http.Header)
-			for k, v := range r.Header {
-				if !slices.Contains(s.IncludeRequestHeaders, k) {
-					continue
-				}
-				for _, vv := range v {
-					headers.Add(k, vv)
-				}
-			}
-
-			req := Request{
-				Headers: headers,
-				Method:  r.Method,
-				Path:    r.URL.Path,
-			}
-
-			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")
+		respAny := handler(request)
-		w.WriteHeader(statusCode)
+		resp := normalizeResponse(s.t, respAny)
 
-		var respBytes []byte
+		for k, v := range resp.Headers {
-		var err error
+			w.Header()[k] = v
-		if respString, ok := resp.(string); ok {
-			respBytes = []byte(respString)
-		} else if respBytes0, ok := resp.([]byte); ok {
-			respBytes = respBytes0
-		} else {
-			respBytes, err = json.MarshalIndent(resp, "", "    ")
-			if err != nil {
-				http.Error(w, err.Error(), http.StatusInternalServerError)
-				return
-			}
 		}
 
-		if _, err := w.Write(respBytes); err != nil {
+		w.WriteHeader(resp.StatusCode)
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+
+		if _, err := w.Write(resp.Body); err != nil {
+			s.t.Errorf("Failed to write response: %s", err)
 			return
 		}
 	}).Methods(method)
@@ -182,3 +295,43 @@ func getToken(r *http.Request) string {
 
 	return header[len(prefix):]
 }
+
+func getLoggedRequest(req Request, includedHeaders []string) LoggedRequest {
+	result := LoggedRequest{
+		Method:  req.Method,
+		Path:    req.URL.Path,
+		Headers: filterHeaders(req.Headers, includedHeaders),
+	}
+
+	if json.Valid(req.Body) {
+		result.Body = json.RawMessage(req.Body)
+	} else {
+		result.RawBody = string(req.Body)
+	}
+
+	return result
+}
+
+func filterHeaders(h http.Header, includedHeaders []string) http.Header {
+	headers := make(http.Header)
+	for k, v := range h {
+		if !slices.Contains(includedHeaders, k) {
+			continue
+		}
+		headers[k] = v
+	}
+	return headers
+}
+
+func isNil(i any) bool {
+	if i == nil {
+		return true
+	}
+	v := reflect.ValueOf(i)
+	switch v.Kind() {
+	case reflect.Chan, reflect.Func, reflect.Map, reflect.Ptr, reflect.Interface, reflect.Slice:
+		return v.IsNil()
+	default:
+		return false
+	}
+}