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.
2025-02-12 14:00:57 +01:00 · 2025-02-12 14:00:57 +01:00 · 4034766c93
parent bfde3585b9
commit 4034766c93
12 changed files with 336 additions and 193 deletions
--- a/acceptance/acceptance_test.go
+++ b/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) {
-					statusCode := http.StatusOK
-					if stub.Response.StatusCode != 0 {
-						statusCode = stub.Response.StatusCode
-					}
-					return stub.Response.Body, statusCode
+				server.Handle(items[0], items[1], func(req testserver.Request) any {
+					return stub.Response
 				})
 			}

--- a/acceptance/bundle/debug/out.stderr.parallel.txt
+++ b/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
-< {} 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
+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
+< 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
--- a/acceptance/bundle/debug/out.stderr.txt
+++ b/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
 < {
--- a/acceptance/cmd_server_test.go
+++ b/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
 }
--- a/acceptance/config_test.go
+++ b/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 {
-		Body       string
-		StatusCode int
-	}
+	Response testserver.Response
 }

 // FindConfigs finds all the config relevant for this test,
--- a/acceptance/selftest/server/out.requests.txt
+++ b/acceptance/selftest/server/out.requests.txt
@ -6,3 +6,7 @@
  "method": "GET",
  "path": "/custom/endpoint"
 }
+{
+  "method": "GET",
+  "path": "/api/2.0/workspace/get-status"
+}
--- a/acceptance/selftest/server/output.txt
+++ b/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
--- a/acceptance/selftest/server/script
+++ b/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
--- a/acceptance/selftest/server/test.toml
+++ b/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: .*'
--- a/acceptance/server_test.go
+++ b/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) {
-		return testUser, http.StatusOK
+	server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any {
+		return testUser
 	})

-	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("GET", "/api/2.0/workspace/get-status", func(req testserver.Request) any {
+		path := req.URL.Query().Get("path")
+		return req.Workspace.WorkspaceGetStatus(path)
 	})

-	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)
-
-		err := decoder.Decode(&request)
-		if err != nil {
-			return internalError(err)
+	server.Handle("POST", "/api/2.0/workspace/mkdirs", func(req testserver.Request) any {
+		var request workspace.Mkdirs
+		if err := json.Unmarshal(req.Body, &request); err != nil {
+			return testserver.Response{
+				Body:       fmt.Sprintf("internal error: %s", err),
+				StatusCode: http.StatusInternalServerError,
+			}
 		}

-		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) {
-		path := r.URL.Query().Get("path")
-
-		return fakeWorkspace.WorkspaceExport(path)
+	server.Handle("GET", "/api/2.0/workspace/export", func(req testserver.Request) any {
+		path := req.URL.Query().Get("path")
+		return req.Workspace.WorkspaceExport(path)
 	})

-	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
-
-		if recursiveStr == "true" {
-			recursive = true
-		} else {
-			recursive = false
-		}
-
-		return fakeWorkspace.WorkspaceDelete(path, recursive)
+	server.Handle("POST", "/api/2.0/workspace/delete", func(req testserver.Request) any {
+		path := req.URL.Query().Get("path")
+		recursive := req.URL.Query().Get("recursive") == "true"
+		req.Workspace.WorkspaceDelete(path, recursive)
+		return ""
 	})

-	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"]
-
-		body := new(bytes.Buffer)
-		_, err := body.ReadFrom(r.Body)
-		if err != nil {
-			return internalError(err)
-		}
-
-		return fakeWorkspace.WorkspaceFilesImportFile(path, body.Bytes())
+	server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(req testserver.Request) any {
+		path := req.Vars["path"]
+		req.Workspace.WorkspaceFilesImportFile(path, req.Body)
+		return ""
 	})

-	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) {
-		vars := mux.Vars(r)
-		objectId := vars["objectId"]
-
+	server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(req testserver.Request) any {
+		objectId := req.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) {
-		request := jobs.CreateJob{}
-		decoder := json.NewDecoder(r.Body)
-
-		err := decoder.Decode(&request)
-		if err != nil {
-			return internalError(err)
+	server.Handle("POST", "/api/2.1/jobs/create", func(req testserver.Request) any {
+		var request jobs.CreateJob
+		if err := json.Unmarshal(req.Body, &request); err != nil {
+			return testserver.Response{
+				Body:       fmt.Sprintf("internal error: %s", err),
+				StatusCode: 500,
+			}
 		}

-		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) {
-		jobId := r.URL.Query().Get("job_id")
-
-		return fakeWorkspace.JobsGet(jobId)
+	server.Handle("GET", "/api/2.1/jobs/get", func(req testserver.Request) any {
+		jobId := req.URL.Query().Get("job_id")
+		return req.Workspace.JobsGet(jobId)
 	})

-	server.Handle("GET", "/api/2.1/jobs/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		return fakeWorkspace.JobsList()
+	server.Handle("GET", "/api/2.1/jobs/list", func(req testserver.Request) any {
+		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
-}
--- a/libs/testserver/fake_workspace.go
+++ b/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{
-			ObjectType: "DIRECTORY",
-			Path:       path,
-		}, http.StatusOK
+		return Response{
+			Body: &workspace.ObjectInfo{
+				ObjectType: "DIRECTORY",
+				Path:       path,
+			},
+		}
 	} else if _, ok := s.files[path]; ok {
-		return workspace.ObjectInfo{
-			ObjectType: "FILE",
-			Path:       path,
-			Language:   "SCALA",
-		}, http.StatusOK
+		return Response{
+			Body: &workspace.ObjectInfo{
+				ObjectType: "FILE",
+				Path:       path,
+				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) {
-	file := s.files[path]
-
-	if file == nil {
-		return nil, http.StatusNotFound
-	}
-
-	return file, http.StatusOK
+func (s *FakeWorkspace) WorkspaceExport(path string) []byte {
+	return s.files[path]
 }

-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{
-		Jobs: list,
-	}, http.StatusOK
+	return Response{
+		Body: jobs.ListJobsResponse{
+			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
-}
--- a/libs/testserver/server.go
+++ b/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)
-			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)
+			s.Requests = append(s.Requests, getLoggedRequest(request, s.IncludeRequestHeaders))
 		}

-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(statusCode)
+		respAny := handler(request)
+		resp := normalizeResponse(s.t, respAny)

-		var respBytes []byte
-		var err error
-		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
-			}
+		for k, v := range resp.Headers {
+			w.Header()[k] = v
 		}

-		if _, err := w.Write(respBytes); err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+		w.WriteHeader(resp.StatusCode)
+
+		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
+	}
+}