Repository: google/go-safeweb
Pull request #273 (2021-04-14 13:23:24)

Add support for named templates.
There is currently no way to execute associated templates.

This means loading a directory of templates is impossible, they need to be parsed one by one.

This change makes them available.


Patch:
diff --git a/safehttp/default_dispatcher.go b/safehttp/default_dispatcher.go
index bd716d7..0260259 100644
--- a/safehttp/default_dispatcher.go
+++ b/safehttp/default_dispatcher.go
@@ -70,13 +70,22 @@ func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
 			return fmt.Errorf("%T is not a safe template and it cannot be parsed and written", t)
 		}
 		if len(x.FuncMap) == 0 {
-			return t.Execute(rw, x.Data)
+			if x.Name == "" {
+				return t.Execute(rw, x.Data)
+			} else {
+				return t.ExecuteTemplate(rw, x.Name, x.Data)
+			}
 		}
 		cloned, err := t.Clone()
 		if err != nil {
 			return err
 		}
-		return cloned.Funcs(x.FuncMap).Execute(rw, x.Data)
+		cloned = cloned.Funcs(x.FuncMap)
+		if x.Name == "" {
+			return cloned.Execute(rw, x.Data)
+		} else {
+			return cloned.ExecuteTemplate(rw, x.Name, x.Data)
+		}
 	case safehtml.HTML:
 		_, err := io.WriteString(rw, x.String())
 		return err
diff --git a/safehttp/response.go b/safehttp/response.go
index 1e7c7e6..09ae990 100644
--- a/safehttp/response.go
+++ b/safehttp/response.go
@@ -51,12 +51,21 @@ type Template interface {
 	// Template fails or if an error occurs while writing the result to the
 	// io.Writer.
 	Execute(wr io.Writer, data interface{}) error
+
+	// ExecuteTemplate applies the named associated template to the specified data
+	// object and writes the output to the io.Writer.
+	//
+	// ExecuteTemplate returns an error if applying the data object to the
+	// Template fails or if an error occurs while writing the result to the
+	// io.Writer.
+	ExecuteTemplate(wr io.Writer, name string, data interface{}) error
 }
 
 // TemplateResponse bundles a Template with its data and names to function
 // mappings to be passed together to the commit phase.
 type TemplateResponse struct {
 	Template Template
+	Name     string
 	Data     interface{}
 	FuncMap  map[string]interface{}
 }
@@ -65,14 +74,31 @@ type TemplateResponse struct {
 // data and calls the Write function of the ResponseWriter, passing the
 // response.
 func ExecuteTemplate(w ResponseWriter, t Template, data interface{}) Result {
-	return w.Write(&TemplateResponse{t, data, nil})
+	return ExecuteNamedTemplate(w, t, "", data)
+}
+
+// ExecuteNamedTemplate creates a TemplateResponse from the provided Template and its
+// data and calls the Write function of the ResponseWriter, passing the
+// response.
+// Leaving name empty is valid if the template does not have associated templates.
+func ExecuteNamedTemplate(w ResponseWriter, t Template, name string, data interface{}) Result {
+	return w.Write(&TemplateResponse{Template: t, Name: name, Data: data, FuncMap: nil})
 }
 
 // ExecuteTemplateWithFuncs creates a TemplateResponse from the provided
 // Template, its data and the name to function mappings and calls the Write
 // function of the ResponseWriter, passing the response.
+// Leaving name empty is valid if the template does not have associated templates.
 func ExecuteTemplateWithFuncs(w ResponseWriter, t Template, data interface{}, fm map[string]interface{}) Result {
-	return w.Write(&TemplateResponse{t, data, fm})
+	return ExecuteNamedTemplateWithFuncs(w, t, "", data, fm)
+}
+
+// ExecuteTemplateWithFuncs creates a TemplateResponse from the provided
+// Template, its data and the name to function mappings and calls the Write
+// function of the ResponseWriter, passing the response.
+// Leaving name empty is valid if the template does not have associated templates.
+func ExecuteNamedTemplateWithFuncs(w ResponseWriter, t Template, name string, data interface{}, fm map[string]interface{}) Result {
+	return w.Write(&TemplateResponse{Template: t, Name: name, Data: data, FuncMap: fm})
 }
 
 // NoContentResponse is sent to the commit phase when it's initiated from


Test patch:
diff --git a/safehttp/default_dispatcher_test.go b/safehttp/default_dispatcher_test.go
index ae0eff3..6d54600 100644
--- a/safehttp/default_dispatcher_test.go
+++ b/safehttp/default_dispatcher_test.go
@@ -18,10 +18,11 @@ import (
 	"html/template"
 	"math"
 	"net/http"
-	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml"
 	safetemplate "github.com/google/safehtml/template"
 )
@@ -51,10 +52,24 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 						Parse("<h1>{{ . }}</h1>")))
 				var data interface{}
 				data = "This is an actual heading, though."
-				return d.Write(w, &safehttp.TemplateResponse{t, data, nil})
+				return d.Write(w, &safehttp.TemplateResponse{Template: t, Data: data})
 			},
 			wantBody: "<h1>This is an actual heading, though.</h1>",
 		},
+		{
+			name: "Named Safe HTML Template Response",
+			write: func(w http.ResponseWriter) error {
+				d := &safehttp.DefaultDispatcher{}
+				t := safehttp.Template(
+					safetemplate.Must(
+						safetemplate.Must(safetemplate.New("name").Parse("<h1>{{ . }}</h1>")).
+							New("associated").Parse("<h2>{{.}}</h2>")))
+				var data interface{}
+				data = "This is an actual heading, though."
+				return d.Write(w, &safehttp.TemplateResponse{t, "associated", data, nil})
+			},
+			wantBody: "<h2>This is an actual heading, though.</h2>",
+		},
 		{
 			name: "Safe HTML Template Response with Token",
 			write: func(w http.ResponseWriter) error {
@@ -69,7 +84,7 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 				fm := map[string]interface{}{
 					"Token": func() string { return "Token-secret" },
 				}
-				return d.Write(w, &safehttp.TemplateResponse{t, data, fm})
+				return d.Write(w, &safehttp.TemplateResponse{Template: t, Data: data, FuncMap: fm})
 			},
 			wantBody: `<form><input type="hidden" name="token" value="Token-secret">Content</form>`,
 		},
@@ -87,7 +102,7 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 				fm := map[string]interface{}{
 					"Nonce": func() string { return "Nonce-secret" },
 				}
-				return d.Write(w, &safehttp.TemplateResponse{t, data, fm})
+				return d.Write(w, &safehttp.TemplateResponse{Template: t, Data: data, FuncMap: fm})
 			},
 			wantBody: `<script nonce="Nonce-secret" type="application/javascript">alert("script")</script><h1>Content</h1>`,
 		},
@@ -105,14 +120,16 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
+
 			err := tt.write(rw)
 
 			if err != nil {
 				t.Errorf("tt.write(rw): got error %v, want nil", err)
 			}
 
-			if gotBody := rw.Body.String(); tt.wantBody != gotBody {
+			if gotBody := b.String(); tt.wantBody != gotBody {
 				t.Errorf("response body: got %q, want %q", gotBody, tt.wantBody)
 			}
 		})
@@ -142,7 +159,7 @@ func TestDefaultDispatcherInvalidResponse(t *testing.T) {
 						Parse("<h1>{{ . }}</h1>")))
 				var data interface{}
 				data = "This is an actual heading, though."
-				return d.Write(w, safehttp.TemplateResponse{t, data, nil})
+				return d.Write(w, safehttp.TemplateResponse{Template: t, Data: data})
 			},
 			want: "",
 		},
@@ -157,13 +174,14 @@ func TestDefaultDispatcherInvalidResponse(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 
 			if err := tt.write(rw); err == nil {
 				t.Error("tt.write(rw): got nil, want error")
 			}
 
-			if want, got := tt.want, rw.Body.String(); want != got {
+			if want, got := tt.want, b.String(); want != got {
 				t.Errorf("response body: got %q, want %q", got, want)
 			}
 		})
diff --git a/safehttp/fileserver_test.go b/safehttp/fileserver_test.go
index df7cd05..8fad7b4 100644
--- a/safehttp/fileserver_test.go
+++ b/safehttp/fileserver_test.go
@@ -18,10 +18,12 @@ import (
 	"io/ioutil"
 	"net/http/httptest"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 )
 
 func TestFileServer(t *testing.T) {
@@ -45,14 +47,14 @@ func TestFileServer(t *testing.T) {
 		{
 			name:     "missing file",
 			path:     "failure",
-			wantCode: safehttp.StatusNotFound,
+			wantCode: 404,
 			wantCT:   "text/plain; charset=utf-8",
 			wantBody: "Not Found\n",
 		},
 		{
 			name:     "file",
 			path:     "foo.html",
-			wantCode: safehttp.StatusOK,
+			wantCode: 200,
 			wantCT:   "text/html; charset=utf-8",
 			wantBody: "<h1>Hello world</h1>",
 		},
@@ -64,18 +66,19 @@ func TestFileServer(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			rr := httptest.NewRecorder()
+			var b strings.Builder
+			rr := safehttptest.NewTestResponseWriter(&b)
 
 			req := httptest.NewRequest(safehttp.MethodGet, "https://test.science/"+tt.path, nil)
 			m.ServeHTTP(rr, req)
 
-			if got, want := rr.Code, tt.wantCode; got != int(tt.wantCode) {
+			if got, want := rr.Status(), tt.wantCode; got != tt.wantCode {
 				t.Errorf("status code got: %v want: %v", got, want)
 			}
 			if got := rr.Header().Get("Content-Type"); tt.wantCT != got {
 				t.Errorf("Content-Type: got %q want %q", got, tt.wantCT)
 			}
-			if diff := cmp.Diff(tt.wantBody, rr.Body.String()); diff != "" {
+			if diff := cmp.Diff(tt.wantBody, b.String()); diff != "" {
 				t.Errorf("Response body diff (-want,+got): \n%s", diff)
 			}
 		})
diff --git a/safehttp/flight_test.go b/safehttp/flight_test.go
index ed0a8c4..c5fab10 100644
--- a/safehttp/flight_test.go
+++ b/safehttp/flight_test.go
@@ -17,9 +17,11 @@ package safehttp_test
 import (
 	"fmt"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml"
 )
 
@@ -70,7 +72,8 @@ func TestFlightInterceptorPanic(t *testing.T) {
 			mux := mb.Mux()
 
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/search", nil)
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 
 			defer func() {
 				r := recover()
@@ -104,7 +107,8 @@ func TestFlightHandlerPanic(t *testing.T) {
 	mux := mb.Mux()
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/search", nil)
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
 
 	defer func() {
 		r := recover()
@@ -146,7 +150,8 @@ func TestFlightDoubleWritePanics(t *testing.T) {
 				mux := mb.Mux()
 
 				req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/search", nil)
-				rw := httptest.NewRecorder()
+				b := &strings.Builder{}
+				rw := safehttptest.NewTestResponseWriter(b)
 				defer func() {
 					if r := recover(); r == nil {
 						t.Fatalf("expected panic")
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index d4d0f28..d2571ed 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -18,8 +18,11 @@ import (
 	"context"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/safehtml"
@@ -63,20 +66,21 @@ func TestMuxOneHandlerOneRequest(t *testing.T) {
 			})
 			mb.Handle("/", safehttp.MethodGet, h)
 
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 
 			mux := mb.Mux()
 			mux.ServeHTTP(rw, tt.req)
 
-			if rw.Code != int(tt.wantStatus) {
-				t.Errorf("rw.Code: got %v want %v", rw.Code, tt.wantStatus)
+			if rw.Status() != tt.wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), tt.wantStatus)
 			}
 
 			if diff := cmp.Diff(tt.wantHeader, map[string][]string(rw.Header())); diff != "" {
 				t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 			}
 
-			if got := rw.Body.String(); got != tt.wantBody {
+			if got := b.String(); got != tt.wantBody {
 				t.Errorf("response body: got %q want %q", got, tt.wantBody)
 			}
 		})
@@ -124,17 +128,18 @@ func TestMuxServeTwoHandlers(t *testing.T) {
 	mux := mb.Mux()
 
 	for _, test := range tests {
-		rw := httptest.NewRecorder()
+		b := &strings.Builder{}
+		rw := safehttptest.NewTestResponseWriter(b)
 		mux.ServeHTTP(rw, test.req)
-		if want := int(test.wantStatus); rw.Code != want {
-			t.Errorf("rw.Code: got %v want %v", rw.Code, want)
+		if want := test.wantStatus; rw.Status() != want {
+			t.Errorf("rw.Status(): got %v want %v", rw.Status(), want)
 		}
 
 		if diff := cmp.Diff(test.wantHeaders, map[string][]string(rw.Header())); diff != "" {
 			t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 		}
 
-		if got, want := rw.Body.String(), test.wantBody; got != want {
+		if got, want := b.String(), test.wantBody; got != want {
 			t.Errorf("response body: got %q want %q", got, want)
 		}
 	}
@@ -171,10 +176,11 @@ func TestMuxRegisterCorrectHandlerAllPaths(t *testing.T) {
 	mux := mb.Mux()
 
 	for _, test := range tests {
-		rw := httptest.NewRecorder()
+		b := &strings.Builder{}
+		rw := safehttptest.NewTestResponseWriter(b)
 		mux.ServeHTTP(rw, test.req)
 
-		if got, want := rw.Body.String(), test.wantBody; got != want {
+		if got, want := b.String(), test.wantBody; got != want {
 			t.Errorf("response body: got %q want %q", got, want)
 		}
 	}
@@ -356,21 +362,22 @@ func TestMuxInterceptors(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
 
 			tt.mux.ServeHTTP(rw, req)
 
-			if rw.Code != int(tt.wantStatus) {
-				t.Errorf("rw.Code: got %v want %v", rw.Code, tt.wantStatus)
+			if rw.Status() != tt.wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), tt.wantStatus)
 			}
 
 			if diff := cmp.Diff(tt.wantHeaders, map[string][]string(rw.Header())); diff != "" {
 				t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 			}
 
-			if got := rw.Body.String(); got != tt.wantBody {
+			if got := b.String(); got != tt.wantBody {
 				t.Errorf("response body: got %q want %q", got, tt.wantBody)
 			}
 		})
@@ -458,22 +465,23 @@ func TestMuxInterceptorConfigs(t *testing.T) {
 			})
 			mb.Handle("/bar", safehttp.MethodGet, registeredHandler, tt.config)
 
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 
 			req := httptest.NewRequest("GET", "http://foo.com/bar", nil)
 
 			mux := mb.Mux()
 			mux.ServeHTTP(rw, req)
 
-			if rw.Code != int(tt.wantStatus) {
-				t.Errorf("rw.Code: got %v want %v", rw.Code, tt.wantStatus)
+			if rw.Status() != tt.wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), tt.wantStatus)
 			}
 
 			if diff := cmp.Diff(tt.wantHeaders, map[string][]string(rw.Header())); diff != "" {
 				t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 			}
 
-			if got := rw.Body.String(); got != tt.wantBody {
+			if got := b.String(); got != tt.wantBody {
 				t.Errorf("response body: got %q want %q", got, tt.wantBody)
 			}
 		})
@@ -539,15 +547,16 @@ func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 	})
 	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
 
 	req := httptest.NewRequest("GET", "http://foo.com/bar", nil)
 
 	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 
-	if want := safehttp.StatusOK; rw.Code != int(want) {
-		t.Errorf("rw.Code: got %v want %v", rw.Code, want)
+	if want := safehttp.StatusOK; rw.Status() != want {
+		t.Errorf("rw.Status(): got %v want %v", rw.Status(), want)
 	}
 	wantHeaders := map[string][]string{
 		"Dessert":      {"tiramisu"},
@@ -561,7 +570,7 @@ func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 	if diff := cmp.Diff(wantHeaders, map[string][]string(rw.Header())); diff != "" {
 		t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 	}
-	if got, want := rw.Body.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
+	if got, want := b.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
 		t.Errorf(`response body: got %q want %q`, got, want)
 	}
 }
@@ -574,18 +583,19 @@ func TestMuxHandlerReturnsNotWritten(t *testing.T) {
 	mb.Handle("/bar", safehttp.MethodGet, h)
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
 
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
 
 	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 
-	if want := safehttp.StatusNoContent; rw.Code != int(want) {
-		t.Errorf("rw.Code: got %v want %v", rw.Code, want)
+	if want := safehttp.StatusNoContent; rw.Status() != want {
+		t.Errorf("rw.Status(): got %v want %v", rw.Status(), want)
 	}
 	if diff := cmp.Diff(map[string][]string{}, map[string][]string(rw.Header())); diff != "" {
 		t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 	}
-	if got := rw.Body.String(); got != "" {
+	if got := b.String(); got != "" {
 		t.Errorf(`response body got: %q want: ""`, got)
 	}
 }
diff --git a/safehttp/plugins/hostcheck/hostcheck_test.go b/safehttp/plugins/hostcheck/hostcheck_test.go
index 5aec398..9b36a3c 100644
--- a/safehttp/plugins/hostcheck/hostcheck_test.go
+++ b/safehttp/plugins/hostcheck/hostcheck_test.go
@@ -17,10 +17,12 @@ package hostcheck_test
 import (
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/hostcheck"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml"
 )
 
@@ -52,12 +54,13 @@ func TestInterceptor(t *testing.T) {
 			})
 			mb.Handle("/", safehttp.MethodGet, h)
 
-			rw := httptest.NewRecorder()
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 			mux := mb.Mux()
 			mux.ServeHTTP(rw, tt.req)
 
-			if rw.Code != int(tt.wantStatus) {
-				t.Errorf("rw.Code: got %v want %v", rw.Code, tt.wantStatus)
+			if rw.Status() != tt.wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), tt.wantStatus)
 			}
 		})
 	}
diff --git a/safehttp/safehttptest/recorder.go b/safehttp/safehttptest/recorder.go
index 89d653b..39d6051 100644
--- a/safehttp/safehttptest/recorder.go
+++ b/safehttp/safehttptest/recorder.go
@@ -15,8 +15,9 @@
 package safehttptest
 
 import (
+	"io"
 	"net/http"
-	"net/http/httptest"
+	"strings"
 
 	"github.com/google/go-safeweb/safehttp"
 )
@@ -26,18 +27,32 @@ import (
 // should be passed as part of the handler function in tests.
 type ResponseRecorder struct {
 	safehttp.ResponseWriter
-	rw *httptest.ResponseRecorder
+	rw *TestResponseWriter
+	b  *strings.Builder
 }
 
 // NewResponseRecorder creates a ResponseRecorder.
 func NewResponseRecorder() *ResponseRecorder {
-	rw := httptest.NewRecorder()
+	var b strings.Builder
+	rw := NewTestResponseWriter(&b)
 	return &ResponseRecorder{
 		rw:             rw,
+		b:              &b,
 		ResponseWriter: safehttp.DeprecatedNewResponseWriter(rw, nil),
 	}
 }
 
+// NewResponseRecorderFromDispatcher creates a ResponseRecorder.
+func NewResponseRecorderFromDispatcher(d safehttp.Dispatcher) *ResponseRecorder {
+	var b strings.Builder
+	rw := NewTestResponseWriter(&b)
+	return &ResponseRecorder{
+		rw:             rw,
+		b:              &b,
+		ResponseWriter: safehttp.DeprecatedNewResponseWriter(rw, d),
+	}
+}
+
 // Header returns the recorded response headers.
 func (r *ResponseRecorder) Header() http.Header {
 	return r.rw.Header()
@@ -45,10 +60,48 @@ func (r *ResponseRecorder) Header() http.Header {
 
 // Status returns the recorded response status code.
 func (r *ResponseRecorder) Status() safehttp.StatusCode {
-	return safehttp.StatusCode(r.rw.Code)
+	return r.rw.status
 }
 
 // Body returns the recorded response body.
 func (r *ResponseRecorder) Body() string {
-	return r.rw.Body.String()
+	return r.b.String()
+}
+
+// TestResponseWriter is an implementation of the http.ResponseWriter interface
+// used for constructing an HTTP response for testing purposes.
+type TestResponseWriter struct {
+	header http.Header
+	writer io.Writer
+	status safehttp.StatusCode
+}
+
+// NewTestResponseWriter creates a new TestResponseWriter.
+func NewTestResponseWriter(w io.Writer) *TestResponseWriter {
+	return &TestResponseWriter{
+		header: http.Header{},
+		writer: w,
+		status: safehttp.StatusOK,
+	}
+}
+
+// Status returns the response status.
+func (r *TestResponseWriter) Status() safehttp.StatusCode {
+	return r.status
+}
+
+// Header implements http.ResponseWriter. It returns the response headers that
+// could have been mutated within a handler.
+func (r *TestResponseWriter) Header() http.Header {
+	return r.header
+}
+
+// WriteHeader implements http.ResponseWriter.
+func (r *TestResponseWriter) WriteHeader(statusCode int) {
+	r.status = safehttp.StatusCode(statusCode)
+}
+
+// Write implements http.ResponseWriter.
+func (r *TestResponseWriter) Write(data []byte) (int, error) {
+	return r.writer.Write(data)
 }
diff --git a/safehttp/task_test.go b/safehttp/task_test.go
index a9aebd6..ea0010a 100644
--- a/safehttp/task_test.go
+++ b/safehttp/task_test.go
@@ -17,7 +17,7 @@ package safehttp_test
 import (
 	"math"
 	"net/http"
-	"net/http/httptest"
+	"strings"
 	"testing"
 	"text/template"
 
@@ -35,7 +35,7 @@ func stubResponseWriter(dispatcher safehttp.Dispatcher, rw http.ResponseWriter)
 }
 
 func TestTaskSetCookie(t *testing.T) {
-	testRw := httptest.NewRecorder()
+	testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
 	rw := stubResponseWriter(safehttp.DefaultDispatcher{}, testRw)
 
 	c := safehttp.NewCookie("foo", "bar")
@@ -53,7 +53,7 @@ func TestTaskSetCookie(t *testing.T) {
 }
 
 func TestTaskSetInvalidCookie(t *testing.T) {
-	testRw := httptest.NewRecorder()
+	testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
 	rw := stubResponseWriter(safehttp.DefaultDispatcher{}, testRw)
 
 	c := safehttp.NewCookie("f=oo", "bar")
@@ -101,7 +101,7 @@ func TestTaskWriteTwicePanic(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := stubResponseWriter(safehttp.DefaultDispatcher{}, httptest.NewRecorder())
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r != nil {
 					return
@@ -141,7 +141,7 @@ func TestTaskUnsafeResponse(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := stubResponseWriter(safehttp.DefaultDispatcher{}, httptest.NewRecorder())
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r != nil {
 					return
@@ -209,10 +209,10 @@ func TestTaskStatusCode(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			testRw := httptest.NewRecorder()
+			testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
 			w := stubResponseWriter(safehttp.DefaultDispatcher{}, testRw)
 			tt.write(w)
-			if gotStatus := testRw.Code; gotStatus != int(tt.wantStatus) {
+			if gotStatus := testRw.Status(); gotStatus != tt.wantStatus {
 				t.Errorf("tt.write(w): got %v, want %v", gotStatus, tt.wantStatus)
 			}
 		})
@@ -241,7 +241,7 @@ func TestTaskInvalidStatusCode(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := stubResponseWriter(safehttp.DefaultDispatcher{}, httptest.NewRecorder())
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r != nil {
 					return
diff --git a/tests/integration/csp/csp_test.go b/tests/integration/csp/csp_test.go
index a447c9c..214ce32 100644
--- a/tests/integration/csp/csp_test.go
+++ b/tests/integration/csp/csp_test.go
@@ -16,11 +16,13 @@ package csp_test
 
 import (
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/csp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml/template"
 )
 
@@ -46,7 +48,8 @@ func TestServeMuxInstallCSP(t *testing.T) {
 	})
 	mb.Handle("/bar", safehttp.MethodGet, handler)
 
-	rr := httptest.NewRecorder()
+	b := strings.Builder{}
+	rr := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/bar", nil)
 
@@ -61,8 +64,8 @@ func TestServeMuxInstallCSP(t *testing.T) {
 		t.Fatalf("csp.Nonce: got %q, want non-empty", nonce)
 	}
 
-	if got, want := rr.Code, safehttp.StatusOK; got != int(want) {
-		t.Errorf("rr.Code got: %v want: %v", got, want)
+	if want, got := rr.Status(), safehttp.StatusOK; got != want {
+		t.Errorf("rr.Status() got: %v want: %v", got, want)
 	}
 
 	wantHeaders := map[string][]string{
@@ -78,7 +81,7 @@ func TestServeMuxInstallCSP(t *testing.T) {
 
 	wantBody := `<script nonce="` + nonce +
 		`" type="application/javascript">alert("script")</script><h1>Content</h1>`
-	if gotBody := rr.Body.String(); gotBody != wantBody {
+	if gotBody := b.String(); gotBody != wantBody {
 		t.Errorf("response body: got %q, want nonce %q", gotBody, wantBody)
 	}
 
diff --git a/tests/integration/errors/errors_test.go b/tests/integration/errors/errors_test.go
index 3b23583..0bc7121 100644
--- a/tests/integration/errors/errors_test.go
+++ b/tests/integration/errors/errors_test.go
@@ -19,12 +19,14 @@ import (
 	"math"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/safehtml"
 
 	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml/template"
 )
 
@@ -72,35 +74,37 @@ func TestCustomErrors(t *testing.T) {
 	m := mb.Mux()
 
 	t.Run("correct request", func(t *testing.T) {
-		rr := httptest.NewRecorder()
+		b := strings.Builder{}
+		rr := safehttptest.NewTestResponseWriter(&b)
 
 		req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/compute?a=3", nil)
 		m.ServeHTTP(rr, req)
 
-		if got, want := rr.Code, safehttp.StatusOK; got != int(want) {
-			t.Errorf("rr.Code got: %v want: %v", got, want)
+		if got, want := rr.Status(), safehttp.StatusOK; got != want {
+			t.Errorf("rr.Status() got: %v want: %v", got, want)
 		}
 		want := "Result: 9"
-		if diff := cmp.Diff(want, rr.Body.String()); diff != "" {
-			t.Errorf("response body diff (-want,+got): \n%s\ngot %q, want %q", diff, rr.Body.String(), want)
+		if diff := cmp.Diff(want, b.String()); diff != "" {
+			t.Errorf("response body diff (-want,+got): \n%s\ngot %q, want %q", diff, b.String(), want)
 		}
 	})
 
 	t.Run("missing parameter", func(t *testing.T) {
-		rr := httptest.NewRecorder()
+		b := strings.Builder{}
+		rr := safehttptest.NewTestResponseWriter(&b)
 
 		req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/compute?foo=3", nil)
 		m.ServeHTTP(rr, req)
 
-		if got, want := rr.Code, safehttp.StatusBadRequest; got != int(want) {
-			t.Errorf("rr.Code got: %v want: %v", got, want)
+		if got, want := rr.Status(), safehttp.StatusBadRequest; got != want {
+			t.Errorf("rr.Status() got: %v want: %v", got, want)
 		}
 
 		want := `<h1>Error: Bad Request</h1>
 <p>missing parameter &#39;a&#39;</p>
 `
-		if diff := cmp.Diff(want, rr.Body.String()); diff != "" {
-			t.Errorf("response body diff (-want,+got): \n%s\ngot %q, want %q", diff, rr.Body.String(), want)
+		if diff := cmp.Diff(want, b.String()); diff != "" {
+			t.Errorf("response body diff (-want,+got): \n%s\ngot %q, want %q", diff, b.String(), want)
 		}
 	})
 
@@ -130,19 +134,20 @@ func TestCustomErrorsInBefore(t *testing.T) {
 
 	m := mb.Mux()
 	t.Run("error in Before", func(t *testing.T) {
-		rr := httptest.NewRecorder()
+		b := strings.Builder{}
+		rr := safehttptest.NewTestResponseWriter(&b)
 
 		req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/compute?a=3", nil)
 		m.ServeHTTP(rr, req)
 
-		if got, want := rr.Code, safehttp.StatusForbidden; got != int(want) {
-			t.Errorf("rr.Code got: %v want: %v", got, want)
+		if got, want := rr.Status(), safehttp.StatusForbidden; got != want {
+			t.Errorf("rr.Status() got: %v want: %v", got, want)
 		}
 		want := `<h1>Error: Forbidden</h1>
 <p>forbidden in Before</p>
 `
-		if diff := cmp.Diff(want, rr.Body.String()); diff != "" {
-			t.Errorf("response body diff (-want,+got): \n%s\ngot %q, want %q", diff, rr.Body.String(), want)
+		if diff := cmp.Diff(want, b.String()); diff != "" {
+			t.Errorf("response body diff (-want,+got): \n%s\ngot %q, want %q", diff, b.String(), want)
 		}
 	})
 }
diff --git a/tests/integration/header/header_test.go b/tests/integration/header/header_test.go
index 466dbe6..e913ea5 100644
--- a/tests/integration/header/header_test.go
+++ b/tests/integration/header/header_test.go
@@ -21,6 +21,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/safehtml"
@@ -43,7 +45,9 @@ func TestAccessIncomingHeaders(t *testing.T) {
 		t.Fatalf("http.ReadRequest() got err: %v", err)
 	}
 
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
+
 	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 }
@@ -57,7 +61,8 @@ func TestChangingResponseHeaders(t *testing.T) {
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/", nil)
 
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
 
 	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
diff --git a/tests/integration/hsts/hsts_test.go b/tests/integration/hsts/hsts_test.go
index 2697726..3590783 100644
--- a/tests/integration/hsts/hsts_test.go
+++ b/tests/integration/hsts/hsts_test.go
@@ -17,8 +17,11 @@ package hsts_test
 import (
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/hsts"
@@ -34,14 +37,15 @@ func TestHSTSServeMuxInstall(t *testing.T) {
 	})
 	mb.Handle("/asdf", safehttp.MethodGet, handler)
 
-	rw := httptest.NewRecorder()
+	b := strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
 	mb.Mux().ServeHTTP(rw, req)
 
-	if want := safehttp.StatusOK; rw.Code != int(want) {
-		t.Errorf("rw.Code got: %v want: %v", rw.Code, want)
+	if want := safehttp.StatusOK; rw.Status() != want {
+		t.Errorf("rw.Status() got: %v want: %v", rw.Status(), want)
 	}
 
 	wantHeaders := map[string][]string{
@@ -52,8 +56,8 @@ func TestHSTSServeMuxInstall(t *testing.T) {
 		t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 	}
 
-	if got, want := rw.Body.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
-		t.Errorf("Body got: %v want: %v", got, want)
+	if got, want := b.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
+		t.Errorf("b.String() got: %v want: %v", got, want)
 	}
 }
 
@@ -66,14 +70,15 @@ func TestHSTSOnErrors(t *testing.T) {
 	})
 	mb.Handle("/asdf", safehttp.MethodGet, handler)
 
-	rw := httptest.NewRecorder()
+	b := strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
 	mb.Mux().ServeHTTP(rw, req)
 
-	if want := safehttp.StatusNotFound; rw.Code != int(want) {
-		t.Errorf("rw.Code got: %v want: %v", rw.Code, want)
+	if want := safehttp.StatusNotFound; rw.Status() != want {
+		t.Errorf("rw.Status() got: %v want: %v", rw.Status(), want)
 	}
 	if got, want := rw.Header()["Strict-Transport-Security"], []string{"max-age=63072000; includeSubDomains"}; !cmp.Equal(got, want) {
 		t.Errorf("rw.Header()[\"Strict-Transport-Security\"] = %q, want %q", got, want)
diff --git a/tests/integration/mux/mux_test.go b/tests/integration/mux/mux_test.go
index 338760e..da072f3 100644
--- a/tests/integration/mux/mux_test.go
+++ b/tests/integration/mux/mux_test.go
@@ -18,10 +18,12 @@ import (
 	"html/template"
 	"math"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml"
 	safetemplate "github.com/google/safehtml/template"
 )
@@ -73,21 +75,21 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			mb := &safehttp.ServeMuxConfig{}
 			mb.Handle("/pizza", safehttp.MethodGet, tt.handler)
-
-			rw := httptest.NewRecorder()
-
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/pizza", nil)
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
+
 			mb.Mux().ServeHTTP(rw, req)
 
-			if wantStatus := safehttp.StatusOK; rw.Code != int(wantStatus) {
-				t.Errorf("rw.Code: got %v want %v", rw.Code, wantStatus)
+			if wantStatus := safehttp.StatusOK; rw.Status() != wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), wantStatus)
 			}
 
 			if diff := cmp.Diff(tt.wantHeaders, map[string][]string(rw.Header())); diff != "" {
-				t.Errorf("rw.Header mismatch (-want +got):\n%s", diff)
+				t.Errorf("rw.header mismatch (-want +got):\n%s", diff)
 			}
 
-			if gotBody := rw.Body.String(); tt.wantBody != gotBody {
+			if gotBody := b.String(); tt.wantBody != gotBody {
 				t.Errorf("response body: got %v, want %v", gotBody, tt.wantBody)
 			}
 		})
@@ -126,19 +128,17 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 			// cases from the previous test into a single table test after
 			// error-handling in the ResponseWriter has been fixed.
 			t.Skip()
+			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/pizza", nil)
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
 
 			mb := &safehttp.ServeMuxConfig{}
 			mb.Handle("/pizza", safehttp.MethodGet, tt.handler)
 			mux := mb.Mux()
-
-			rw := httptest.NewRecorder()
-
-			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/pizza", nil)
-
 			mux.ServeHTTP(rw, req)
 
-			if wantStatus := safehttp.StatusInternalServerError; rw.Code != int(wantStatus) {
-				t.Errorf("rw.Code: got %v want %v", rw.Code, wantStatus)
+			if wantStatus := safehttp.StatusInternalServerError; rw.Status() != wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), wantStatus)
 			}
 
 			wantHeaders := map[string][]string{
@@ -149,7 +149,7 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 				t.Errorf("rw.Header(): mismatch (-want +got):\n%s", diff)
 			}
 
-			if wantBody, gotBody := "Internal Server Error\n", rw.Body.String(); wantBody != gotBody {
+			if wantBody, gotBody := "Internal Server Error\n", b.String(); wantBody != gotBody {
 				t.Errorf("response body: got %v, want %v", gotBody, wantBody)
 			}
 		})
diff --git a/tests/integration/safehtml/safehtml_test.go b/tests/integration/safehtml/safehtml_test.go
index c46f71a..0889fa3 100644
--- a/tests/integration/safehtml/safehtml_test.go
+++ b/tests/integration/safehtml/safehtml_test.go
@@ -16,8 +16,11 @@ package safehtml_test
 
 import (
 	"net/http/httptest"
+	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/safehtml"
 	"github.com/google/safehtml/template"
@@ -31,11 +34,12 @@ func TestHandleRequestWrite(t *testing.T) {
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/", nil)
 
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
 
 	mb.Mux().ServeHTTP(rw, req)
 
-	body := rw.Body.String()
+	body := b.String()
 
 	if want := "&lt;h1&gt;Escaped, so not really a heading&lt;/h1&gt;"; body != want {
 		t.Errorf("body got: %q want: %q", body, want)
@@ -50,11 +54,12 @@ func TestHandleRequestWriteTemplate(t *testing.T) {
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/", nil)
 
-	rw := httptest.NewRecorder()
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
 
 	mb.Mux().ServeHTTP(rw, req)
 
-	body := rw.Body.String()
+	body := b.String()
 
 	if want := "<h1>This is an actual heading, though.</h1>"; body != want {
 		t.Errorf("body got: %q want: %q", body, want)
diff --git a/tests/integration/staticheaders/staticheaders_test.go b/tests/integration/staticheaders/staticheaders_test.go
index b730e9f..804682d 100644
--- a/tests/integration/staticheaders/staticheaders_test.go
+++ b/tests/integration/staticheaders/staticheaders_test.go
@@ -17,8 +17,11 @@ package staticheaders_test
 import (
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/staticheaders"
@@ -34,14 +37,15 @@ func TestServeMuxInstallStaticHeaders(t *testing.T) {
 	})
 	mb.Handle("/asdf", safehttp.MethodGet, handler)
 
-	rw := httptest.NewRecorder()
+	b := strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
 	mb.Mux().ServeHTTP(rw, req)
 
-	if want := safehttp.StatusOK; rw.Code != int(want) {
-		t.Errorf("rw.Code got: %v want: %v", rw.Code, want)
+	if want := safehttp.StatusOK; rw.Status() != want {
+		t.Errorf("rw.Status() got: %v want: %v", rw.Status(), want)
 	}
 
 	wantHeaders := map[string][]string{
@@ -53,8 +57,8 @@ func TestServeMuxInstallStaticHeaders(t *testing.T) {
 		t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
 	}
 
-	if got, want := rw.Body.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
-		t.Errorf("response body got: %v want: %v", got, want)
+	if got, want := b.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
+		t.Errorf("b.String() got: %v want: %v", got, want)
 	}
 }
 
@@ -67,14 +71,15 @@ func TestStaticHeadersOnError(t *testing.T) {
 	})
 	mb.Handle("/asdf", safehttp.MethodGet, handler)
 
-	rw := httptest.NewRecorder()
+	b := strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
 	mb.Mux().ServeHTTP(rw, req)
 
-	if want := safehttp.StatusNotFound; rw.Code != int(want) {
-		t.Errorf("rw.Status() got: %v want: %v", rw.Code, want)
+	if want := safehttp.StatusNotFound; rw.Status() != want {
+		t.Errorf("rw.Status() got: %v want: %v", rw.Status(), want)
 	}
 
 	wantHeaders := map[string][]string{
diff --git a/tests/integration/xsrf/xsrf_test.go b/tests/integration/xsrf/xsrf_test.go
index 7617b29..4d9fb3a 100644
--- a/tests/integration/xsrf/xsrf_test.go
+++ b/tests/integration/xsrf/xsrf_test.go
@@ -21,6 +21,7 @@ import (
 
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/xsrf/xsrfhtml"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"github.com/google/safehtml/template"
 )
 
@@ -39,17 +40,18 @@ func TestServeMuxInstallXSRF(t *testing.T) {
 	})
 	mb.Handle("/bar", safehttp.MethodGet, handler)
 
-	rr := httptest.NewRecorder()
+	b := strings.Builder{}
+	rr := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/bar", nil)
 
 	mb.Mux().ServeHTTP(rr, req)
 
-	if got, want := rr.Code, safehttp.StatusOK; got != int(want) {
-		t.Errorf("rr.Code got: %v want: %v", got, int(want))
+	if got, want := rr.Status(), safehttp.StatusOK; got != want {
+		t.Errorf("rr.Status() got: %v want: %v", got, want)
 	}
 
-	if gotBody := rr.Body.String(); strings.Contains(gotBody, "WrongToken") {
+	if gotBody := b.String(); strings.Contains(gotBody, "WrongToken") {
 		t.Errorf("response body: got %q, injection failed", gotBody)
 	}