From af6c705d8cf86b3b8e0124bf871039104ed802ed Mon Sep 17 00:00:00 2001
From: Sasha Koshka <sashakoshka@tebibyte.media>
Date: Sat, 7 Dec 2024 00:46:24 -0500
Subject: [PATCH] Add HTTPResponseRecorder

---
 http.go | 48 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/http.go b/http.go
index d7cf68e..32d8e46 100644
--- a/http.go
+++ b/http.go
@@ -1,5 +1,7 @@
 package step
 
+import "io"
+import "bytes"
 import "net/http"
 
 // HTTPData represents information about an ongoing HTTP request that is made
@@ -7,9 +9,53 @@ import "net/http"
 type HTTPData struct {
 	// Res is like an http.ResponseWriter.
 	Res struct {
-		Header http.Header
+		Header      http.Header
 		WriteHeader func (statusCode int)
 	}
 	// Req is the HTTP request.
 	Req *http.Request
 }
+
+var _ http.ResponseWriter = new(HTTPResponseRecorder)
+
+// HTTPResponseRecorder is an http.ResponseWriter that can buffer a response to
+// be played back later.
+type HTTPResponseRecorder struct {
+	statusCode int
+	header     http.Header
+	buffer     bytes.Buffer
+}
+
+func (this *HTTPResponseRecorder) Header () http.Header {
+	if this.header == nil {
+		this.header = make(http.Header)
+	}
+	return this.header
+}
+
+func (this *HTTPResponseRecorder) SetHeader (header http.Header) {
+	this.header = header
+}
+
+func (this *HTTPResponseRecorder) Write (buffer []byte) (int, error) {
+	return this.buffer.Write(buffer)
+}
+
+func (this *HTTPResponseRecorder) WriteHeader (statusCode int) {
+	this.statusCode = statusCode
+}
+
+// Play replays the response to the given http.ResponseWriter. This resets the
+// recorder.
+func (this *HTTPResponseRecorder) Play (res http.ResponseWriter) error {
+	defer this.Reset()
+	_, err := io.Copy(res, &this.buffer)
+	return err
+}
+
+// Reset resets this response recorder so it can be used again.
+func (this *HTTPResponseRecorder) Reset () {
+	this.buffer.Reset()
+	this.header = nil
+	this.statusCode = http.StatusOK
+}