From 72347ec910b9c33d603119b39e784eae1838f4e8 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 18 Nov 2015 14:17:42 +0200
Subject: [PATCH] Added BodyWrite() to Request and Response for writing request
 and response bodies

---
 http.go      | 32 ++++++++++++++++++++++++++++++++
 http_test.go | 22 ++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/http.go b/http.go
index 2eea7b9..83bac91 100644
--- a/http.go
+++ b/http.go
@@ -78,6 +78,38 @@ func (resp *Response) SetBodyStream(bodyStream io.Reader, bodySize int) {
 	resp.Header.SetContentLength(bodySize)
 }
 
+// BodyWriter returns writer for populating response body.
+func (resp *Response) BodyWriter() io.Writer {
+	return responseBodyWriter{
+		Response: resp,
+	}
+}
+
+// BodyWriter returns writer for populating request body.
+func (req *Request) BodyWriter() io.Writer {
+	return requestBodyWriter{
+		Request: req,
+	}
+}
+
+type responseBodyWriter struct {
+	*Response
+}
+
+func (w responseBodyWriter) Write(p []byte) (int, error) {
+	w.Response.body = append(w.body, p...)
+	return len(p), nil
+}
+
+type requestBodyWriter struct {
+	*Request
+}
+
+func (w requestBodyWriter) Write(p []byte) (int, error) {
+	w.Request.body = append(w.body, p...)
+	return len(p), nil
+}
+
 // Body returns response body.
 func (resp *Response) Body() []byte {
 	return resp.body
diff --git a/http_test.go b/http_test.go
index efda918..72371bb 100644
--- a/http_test.go
+++ b/http_test.go
@@ -8,6 +8,28 @@ import (
 	"testing"
 )
 
+func TestRequestBodyWriter(t *testing.T) {
+	var r Request
+	w := r.BodyWriter()
+	for i := 0; i < 10; i++ {
+		fmt.Fprintf(w, "%d", i)
+	}
+	if string(r.Body()) != "0123456789" {
+		t.Fatalf("unexpected body %q. Expecting %q", r.Body(), "0123456789")
+	}
+}
+
+func TestResponseBodyWriter(t *testing.T) {
+	var r Response
+	w := r.BodyWriter()
+	for i := 0; i < 10; i++ {
+		fmt.Fprintf(w, "%d", i)
+	}
+	if string(r.Body()) != "0123456789" {
+		t.Fatalf("unexpected body %q. Expecting %q", r.Body(), "0123456789")
+	}
+}
+
 func TestRequestWriteRequestURINoHost(t *testing.T) {
 	var req Request
 	req.Header.SetRequestURI("http://google.com/foo/bar?baz=aaa")