From cf5884bc60071ce73100af061133bd81c91bc6b6 Mon Sep 17 00:00:00 2001
From: John Guo <john@johng.cn>
Date: Fri, 17 Jun 2022 11:31:32 +0800
Subject: [PATCH] add SetBodyContent for gclient.Response

---
 net/gclient/gclient_response.go    |  8 ++++++++
 net/gclient/gclient_z_unit_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/net/gclient/gclient_response.go b/net/gclient/gclient_response.go
index 84fef34d8..27869aa85 100644
--- a/net/gclient/gclient_response.go
+++ b/net/gclient/gclient_response.go
@@ -7,6 +7,7 @@
 package gclient
 
 import (
+	"bytes"
 	"io/ioutil"
 	"net/http"
 
@@ -69,6 +70,13 @@ func (r *Response) ReadAllString() string {
 	return string(r.ReadAll())
 }
 
+// SetBodyContent overwrites response content with custom one.
+func (r *Response) SetBodyContent(content []byte) {
+	buffer := bytes.NewBuffer(content)
+	r.Body = ioutil.NopCloser(buffer)
+	r.ContentLength = int64(buffer.Len())
+}
+
 // Close closes the response when it will never be used.
 func (r *Response) Close() error {
 	if r == nil || r.Response == nil {
diff --git a/net/gclient/gclient_z_unit_test.go b/net/gclient/gclient_z_unit_test.go
index 410645449..599ce889c 100644
--- a/net/gclient/gclient_z_unit_test.go
+++ b/net/gclient/gclient_z_unit_test.go
@@ -640,3 +640,27 @@ func TestClient_RequestVar(t *testing.T) {
 		t.AssertNE(users, nil)
 	})
 }
+
+func TestClient_SetBodyContent(t *testing.T) {
+	p, _ := gtcp.GetFreePort()
+	s := g.Server(p)
+	s.BindHandler("/", func(r *ghttp.Request) {
+		r.Response.Write("hello")
+	})
+	s.SetPort(p)
+	s.SetDumpRouterMap(false)
+	s.Start()
+	defer s.Shutdown()
+
+	time.Sleep(100 * time.Millisecond)
+	gtest.C(t, func(t *gtest.T) {
+		c := g.Client()
+		c.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", p))
+		res, err := c.Get(ctx, "/")
+		t.AssertNil(err)
+		defer res.Close()
+		t.Assert(res.ReadAllString(), "hello")
+		res.SetBodyContent([]byte("world"))
+		t.Assert(res.ReadAllString(), "world")
+	})
+}