From c103a7760618995528dbaf7c55c61a8d05a7cb3a Mon Sep 17 00:00:00 2001
From: zstone12 <522089185@qq.com>
Date: Wed, 29 Nov 2023 16:03:12 +0800
Subject: [PATCH] Feature/add req param (#12)

* feat: add req param

* feat: add get & set

* feat: add set formdata

* style: improve comment

* style: remove space

* chore: remove redundant func

* test:optimize POST body test

---------

Co-authored-by: violapioggia <2604296771@qq.com>
---
 client.go      | 19 +++++++++++++++++--
 client_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/client.go b/client.go
index bfac987..1e5b2c8 100644
--- a/client.go
+++ b/client.go
@@ -58,6 +58,7 @@ type Client struct {
 	connected          bool
 	encodingBase64     bool
 	lastEventID        atomic.Value // []byte
+	body               []byte
 }
 
 var defaultClient, _ = client.NewClient(client.WithDialer(standard.NewDialer()), client.WithResponseBodyStream(true))
@@ -170,16 +171,21 @@ func (c *Client) SetOnConnectCallback(fn ConnCallback) {
 	c.connectedCallback = fn
 }
 
-// SetMaxBufferSize  set sse client MaxBufferSize
+// SetMaxBufferSize set sse client MaxBufferSize
 func (c *Client) SetMaxBufferSize(size int) {
 	c.maxBufferSize = size
 }
 
-// SetURL  set sse client url
+// SetURL set sse client url
 func (c *Client) SetURL(url string) {
 	c.url = url
 }
 
+// SetBody set sse client request body
+func (c *Client) SetBody(body []byte) {
+	c.body = body
+}
+
 // SetMethod set sse client request method
 func (c *Client) SetMethod(method string) {
 	c.method = method
@@ -225,6 +231,11 @@ func (c *Client) GetHertzClient() *client.Client {
 	return c.hertzClient
 }
 
+// GetBody get sse client request body
+func (c *Client) GetBody() []byte {
+	return c.body
+}
+
 // GetLastEventID get sse client lastEventID
 func (c *Client) GetLastEventID() []byte {
 	return c.lastEventID.Load().([]byte)
@@ -247,6 +258,10 @@ func (c *Client) request(ctx context.Context, req *protocol.Request, resp *proto
 		req.Header.Set(k, v)
 	}
 
+	if len(c.body) != 0 {
+		req.SetBody(c.body)
+	}
+
 	err := c.hertzClient.Do(ctx, req, resp)
 	return err
 }
diff --git a/client_test.go b/client_test.go
index ed5601b..c2df161 100644
--- a/client_test.go
+++ b/client_test.go
@@ -25,6 +25,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cloudwego/hertz/pkg/protocol/consts"
+
 	"github.com/cloudwego/hertz/pkg/app"
 	"github.com/cloudwego/hertz/pkg/app/server"
 	"github.com/cloudwego/hertz/pkg/common/hlog"
@@ -119,6 +121,28 @@ func newServer401(port string) {
 	h.Spin()
 }
 
+func newServerWithPOSTBody(empty bool, port string) {
+	h := server.Default(server.WithHostPorts("0.0.0.0:" + port))
+
+	h.POST("/sse", func(ctx context.Context, c *app.RequestContext) {
+		// client can tell server last event it received with Last-Event-ID header
+		lastEventID := GetLastEventID(c)
+		hlog.CtxInfof(ctx, "last event ID: %s", lastEventID)
+
+		// you must set status code and response headers before first render call
+		c.SetStatusCode(http.StatusOK)
+		s := NewStream(c)
+		body, err := c.Body()
+		if err != nil {
+			return
+		}
+		for a := 0; a < 10; a++ {
+			s.Publish(&Event{Data: body})
+		}
+	})
+	h.Run()
+}
+
 func publishMsgs(s *Stream, empty bool, count int) {
 	for a := 0; a < count; a++ {
 		if empty {
@@ -275,3 +299,29 @@ func TestTrimHeader(t *testing.T) {
 		assert.DeepEqual(t, tc.want, got)
 	}
 }
+
+func TestRequestWithBody(t *testing.T) {
+	go newServerWithPOSTBody(false, "9006")
+	time.Sleep(time.Second)
+	c := NewClient("http://127.0.0.1:9006/sse")
+	c.SetMethod(consts.MethodPost)
+	c.body = []byte(`{"msg":"echo"}`)
+	events := make(chan *Event)
+	var cErr error
+	go func() {
+		cErr = c.Subscribe(func(msg *Event) {
+			if msg.Data != nil {
+				events <- msg
+				return
+			}
+		})
+	}()
+
+	for i := 0; i < 5; i++ {
+		msg, err := wait(events, time.Second*1)
+		assert.Nil(t, err)
+		assert.DeepEqual(t, []byte(`{"msg":"echo"}`), msg)
+	}
+
+	assert.Nil(t, cErr)
+}