golang · dunglas · May 15, 2023 · May 16, 2023 · May 17, 2023
diff --git a/src/net/http/httptest/example_test.go b/src/net/http/httptest/example_test.go
@@ -10,6 +10,8 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"strings"
+	"time"
 )
 
 func ExampleResponseRecorder() {
@@ -34,6 +36,32 @@ func ExampleResponseRecorder() {
 	// <html><body>Hello World!</body></html>
 }
 
+func ExampleResponseRecorder_requestController() {
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		rc := http.NewResponseController(w)
+		rc.SetReadDeadline(time.Now().Add(1 * time.Second))
+		rc.SetWriteDeadline(time.Now().Add(3 * time.Second))
+
+		io.WriteString(w, "<html><body>Hello, with deadlines!</body></html>")
+	}
+
+	req := httptest.NewRequest("GET", "http://example.com/bar", strings.NewReader("bar"))
+	w := httptest.NewRecorder()
+	handler(w, req)
+
+	resp := w.Result()
+	body, _ := io.ReadAll(resp.Body)
+
+	fmt.Println(resp.StatusCode)
+	fmt.Println(resp.Header.Get("Content-Type"))
+	fmt.Println(string(body))
+
+	// Output:
+	// 200
+	// text/html; charset=utf-8
+	// <html><body>Hello, with deadlines!</body></html>
+}
+
 func ExampleServer() {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		fmt.Fprintln(w, "Hello, client")

diff --git a/src/net/http/httptest/recorder.go b/src/net/http/httptest/recorder.go
@@ -12,6 +12,7 @@ import (
 	"net/textproto"
 	"strconv"
 	"strings"
+	"time"
 
 	"golang.org/x/net/http/httpguts"
 )
@@ -42,6 +43,14 @@ type ResponseRecorder struct {
 	// Flushed is whether the Handler called Flush.
 	Flushed bool
 
+	// ReadDeadline is the last read deadline that has been set using
+	// "net/http".ResponseController
+	ReadDeadline time.Time
+
+	// WriteDeadline is the last write deadline that has been set using
+	// "net/http".ResponseController
+	WriteDeadline time.Time
+
 	result      *http.Response // cache of Result's return value
 	snapHeader  http.Header    // snapshot of HeaderMap at first Write
 	wroteHeader bool
@@ -154,13 +163,50 @@ func (rw *ResponseRecorder) WriteHeader(code int) {
 	rw.snapHeader = rw.HeaderMap.Clone()
 }
 
-// Flush implements http.Flusher. To test whether Flush was
+// FlushError allows using "net/http".ResponseController.Flush()
+// with the recorder. To test whether Flush was
 // called, see rw.Flushed.
-func (rw *ResponseRecorder) Flush() {
+func (rw *ResponseRecorder) FlushError() error {
 	if !rw.wroteHeader {
 		rw.WriteHeader(200)
 	}
 	rw.Flushed = true
+
+	return nil
+}
+
+// Flush implements http.Flusher. To test whether Flush was
+// called, see rw.Flushed.
+func (rw *ResponseRecorder) Flush() {
+	rw.FlushError()
+}
+
+// SetReadDeadline allows using "net/http".ResponseController.SetReadDeadline()
+// with the recorder.
+//
+// The deadline is recorded but is not enforced.
+// To prevent flaky tests reads made after the deadline will work
+// as if no deadline was set.
+//
+// To retrieve the deadline, use rw.ReadDeadline.
+func (rw *ResponseRecorder) SetReadDeadline(deadline time.Time) error {
+	rw.ReadDeadline = deadline
+
+	return nil
+}
+
+// SetWriteDeadline allows using "net/http".ResponseController.SetWriteDeadline()
+// with the recorder.
+//
+// The deadline is recorded but is not enforced.
+// To prevent flaky tests writes made after the deadline will work
+// as if no deadline was set.
+//
+// To retrieve the deadline, use rw.WriteDeadline.
+func (rw *ResponseRecorder) SetWriteDeadline(deadline time.Time) error {
+	rw.WriteDeadline = deadline
+
+	return nil
 }
 
 // Result returns the response generated by the handler.

diff --git a/src/net/http/httptest/recorder_test.go b/src/net/http/httptest/recorder_test.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"net/http"
 	"testing"
+	"time"
 )
 
 func TestRecorder(t *testing.T) {
@@ -369,3 +370,30 @@ func TestRecorderPanicsOnNonXXXStatusCode(t *testing.T) {
 		})
 	}
 }
+
+func TestRecorderSetWriteDeadline(t *testing.T) {
+	rw := NewRecorder()
+	rc := http.NewResponseController(rw)
+
+	expected := time.Now().Add(1 * time.Second)
+	if err := rc.SetWriteDeadline(expected); err != nil {
+		t.Errorf(`"ResponseController.WriteDeadline(): got unexpected error %q`, err)
+	}
+
+	if rw.WriteDeadline != expected {
+		t.Errorf(`"ResponseRecorder.WriteDeadline: got %q want %q`, rw.WriteDeadline, expected)
+	}
+}
+
+func TestRecorderSetReadDeadline(t *testing.T) {
+	rw := NewRecorder()
+
+	expected := time.Now().Add(1 * time.Second)
+	if err := rw.SetReadDeadline(expected); err != nil {
+		t.Errorf(`"ResponseRecorder.SetReadDeadline(): got unexpected error %q`, err)
+	}
+
+	if rw.ReadDeadline != expected {
+		t.Errorf(`"ResponseRecorder.ReadDeadline: got %q want %q`, rw.ReadDeadline, expected)
+	}
+}