From 90baa054bca117a5d6c7e1fd338a3d6a5fe39bee Mon Sep 17 00:00:00 2001 From: Aliaksandr Valialkin Date: Fri, 3 Jun 2016 17:07:10 +0300 Subject: [PATCH] Export IsBodyStream on RequestCtx, Request and Response These methods may help writing proper request handler wrappers for request handlers, which may set body either via SetBodyStream* or via usual body methods. --- http.go | 10 ++++++++++ http_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++ server.go | 5 +++++ server_test.go | 12 ++++++++++++ 4 files changed, 69 insertions(+) diff --git a/http.go b/http.go index 1aee3db..9b29593 100644 --- a/http.go +++ b/http.go @@ -183,6 +183,16 @@ func (resp *Response) SetBodyStream(bodyStream io.Reader, bodySize int) { resp.Header.SetContentLength(bodySize) } +// IsBodyStream returns true if body is set via SetBodyStream* +func (req *Request) IsBodyStream() bool { + return req.bodyStream != nil +} + +// IsBodyStream returns true if body is set via SetBodyStream* +func (resp *Response) IsBodyStream() bool { + return resp.bodyStream != nil +} + // SetBodyStreamWriter registers the given sw for populating request body. // // This function may be used in the following cases: diff --git a/http_test.go b/http_test.go index 4d43133..df371da 100644 --- a/http_test.go +++ b/http_test.go @@ -178,7 +178,13 @@ func TestRequestBodyStreamMultipleBodyCalls(t *testing.T) { var r Request s := "foobar baz abc" + if r.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } r.SetBodyStream(bytes.NewBufferString(s), len(s)) + if !r.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } for i := 0; i < 10; i++ { body := r.Body() if string(body) != s { @@ -191,7 +197,13 @@ func TestResponseBodyStreamMultipleBodyCalls(t *testing.T) { var r Response s := "foobar baz abc" + if r.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } r.SetBodyStream(bytes.NewBufferString(s), len(s)) + if !r.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } for i := 0; i < 10; i++ { body := r.Body() if string(body) != s { @@ -223,7 +235,13 @@ func TestResponseBodyWriteToStream(t *testing.T) { expectedS := "aaabbbccc" buf := bytes.NewBufferString(expectedS) + if r.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } r.SetBodyStream(buf, len(expectedS)) + if !r.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } testBodyWriteTo(t, &r, expectedS, false) } @@ -530,6 +548,9 @@ func TestRequestMayContinue(t *testing.T) { func TestResponseGzipStream(t *testing.T) { var r Response + if r.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } r.SetBodyStreamWriter(func(w *bufio.Writer) { fmt.Fprintf(w, "foo") w.Flush() @@ -540,11 +561,17 @@ func TestResponseGzipStream(t *testing.T) { t.Fatalf("unexpected error: %s", err) } }) + if !r.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } testResponseGzipExt(t, &r, "foobarbaz1234") } func TestResponseDeflateStream(t *testing.T) { var r Response + if r.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } r.SetBodyStreamWriter(func(w *bufio.Writer) { w.Write([]byte("foo")) w.Flush() @@ -555,6 +582,9 @@ func TestResponseDeflateStream(t *testing.T) { t.Fatalf("unexpected error: %s", err) } }) + if !r.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } testResponseDeflateExt(t, &r, "foobarbaz1234") } @@ -901,7 +931,13 @@ func testSetRequestBodyStream(t *testing.T, body string, chunked bool) { if chunked { bodySize = -1 } + if req.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } req.SetBodyStream(bytes.NewBufferString(body), bodySize) + if !req.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } var w bytes.Buffer bw := bufio.NewWriter(&w) @@ -928,7 +964,13 @@ func testSetResponseBodyStream(t *testing.T, body string, chunked bool) { if chunked { bodySize = -1 } + if resp.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } resp.SetBodyStream(bytes.NewBufferString(body), bodySize) + if !resp.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } var w bytes.Buffer bw := bufio.NewWriter(&w) diff --git a/server.go b/server.go index 9ef4cf9..924dabd 100644 --- a/server.go +++ b/server.go @@ -1005,6 +1005,11 @@ func (ctx *RequestCtx) SetBodyStreamWriter(sw StreamWriter) { ctx.Response.SetBodyStreamWriter(sw) } +// IsBodyStream returns true if response body is set via SetBodyStream*. +func (ctx *RequestCtx) IsBodyStream() bool { + return ctx.Response.IsBodyStream() +} + // Logger returns logger, which may be used for logging arbitrary // request-specific messages inside RequestHandler. // diff --git a/server_test.go b/server_test.go index ae87d2e..74ca344 100644 --- a/server_test.go +++ b/server_test.go @@ -92,6 +92,9 @@ func TestServerResponseBodyStream(t *testing.T) { readyCh := make(chan struct{}) h := func(ctx *RequestCtx) { ctx.SetConnectionClose() + if ctx.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } ctx.SetBodyStreamWriter(func(w *bufio.Writer) { fmt.Fprintf(w, "first") if err := w.Flush(); err != nil { @@ -102,6 +105,9 @@ func TestServerResponseBodyStream(t *testing.T) { // there is no need to flush w here, since it will // be flushed automatically after returning from StreamWriter. }) + if !ctx.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } } serverCh := make(chan struct{}) @@ -1256,6 +1262,9 @@ func TestRequestCtxSetBodyStreamWriter(t *testing.T) { var req Request ctx.Init(&req, nil, defaultLogger) + if ctx.IsBodyStream() { + t.Fatalf("IsBodyStream must return false") + } ctx.SetBodyStreamWriter(func(w *bufio.Writer) { fmt.Fprintf(w, "body writer line 1\n") if err := w.Flush(); err != nil { @@ -1263,6 +1272,9 @@ func TestRequestCtxSetBodyStreamWriter(t *testing.T) { } fmt.Fprintf(w, "body writer line 2\n") }) + if !ctx.IsBodyStream() { + t.Fatalf("IsBodyStream must return true") + } s := ctx.Response.String()