github.com/cloudwego/hertz@v0.9.3/pkg/protocol/response_test.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   *
    16   * The MIT License (MIT)
    17   *
    18   * Copyright (c) 2015-present Aliaksandr Valialkin, VertaMedia, Kirill Danshin, Erik Dubbelboer, FastHTTP Authors
    19   *
    20   * Permission is hereby granted, free of charge, to any person obtaining a copy
    21   * of this software and associated documentation files (the "Software"), to deal
    22   * in the Software without restriction, including without limitation the rights
    23   * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    24   * copies of the Software, and to permit persons to whom the Software is
    25   * furnished to do so, subject to the following conditions:
    26   *
    27   * The above copyright notice and this permission notice shall be included in
    28   * all copies or substantial portions of the Software.
    29   *
    30   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    31   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    32   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    33   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    34   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    35   * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    36   * THE SOFTWARE.
    37   *
    38   * This file may have been modified by CloudWeGo authors. All CloudWeGo
    39   * Modifications are Copyright 2022 CloudWeGo Authors.
    40   */
    41  
    42  package protocol
    43  
    44  import (
    45  	"bytes"
    46  	"fmt"
    47  	"math"
    48  	"reflect"
    49  	"testing"
    50  
    51  	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
    52  	"github.com/cloudwego/hertz/pkg/common/compress"
    53  	"github.com/cloudwego/hertz/pkg/common/test/assert"
    54  	"github.com/cloudwego/hertz/pkg/common/test/mock"
    55  	"github.com/cloudwego/hertz/pkg/protocol/consts"
    56  )
    57  
    58  func TestResponseCopyTo(t *testing.T) {
    59  	t.Parallel()
    60  
    61  	var resp Response
    62  
    63  	// empty copy
    64  	testResponseCopyTo(t, &resp)
    65  
    66  	// init resp
    67  	// resp.laddr = zeroTCPAddr
    68  	resp.SkipBody = true
    69  	resp.Header.SetStatusCode(consts.StatusOK)
    70  	resp.SetBodyString("test")
    71  	testResponseCopyTo(t, &resp)
    72  }
    73  
    74  func TestResponseBodyStreamMultipleBodyCalls(t *testing.T) {
    75  	t.Parallel()
    76  
    77  	var r Response
    78  
    79  	s := "foobar baz abc"
    80  	if r.IsBodyStream() {
    81  		t.Fatalf("IsBodyStream must return false")
    82  	}
    83  	r.SetBodyStream(bytes.NewBufferString(s), len(s))
    84  	if !r.IsBodyStream() {
    85  		t.Fatalf("IsBodyStream must return true")
    86  	}
    87  	for i := 0; i < 10; i++ {
    88  		body := r.Body()
    89  		if string(body) != s {
    90  			t.Fatalf("unexpected body %q. Expecting %q. iteration %d", body, s, i)
    91  		}
    92  	}
    93  }
    94  
    95  func TestResponseBodyWriteToPlain(t *testing.T) {
    96  	t.Parallel()
    97  
    98  	var r Response
    99  
   100  	expectedS := "foobarbaz"
   101  	r.AppendBodyString(expectedS)
   102  
   103  	testBodyWriteTo(t, &r, expectedS, true)
   104  }
   105  
   106  func TestResponseBodyWriteToStream(t *testing.T) {
   107  	t.Parallel()
   108  
   109  	var r Response
   110  
   111  	expectedS := "aaabbbccc"
   112  	buf := bytes.NewBufferString(expectedS)
   113  	if r.IsBodyStream() {
   114  		t.Fatalf("IsBodyStream must return false")
   115  	}
   116  	r.SetBodyStream(buf, len(expectedS))
   117  	if !r.IsBodyStream() {
   118  		t.Fatalf("IsBodyStream must return true")
   119  	}
   120  
   121  	testBodyWriteTo(t, &r, expectedS, false)
   122  }
   123  
   124  func TestResponseBodyWriter(t *testing.T) {
   125  	t.Parallel()
   126  
   127  	var r Response
   128  	w := r.BodyWriter()
   129  	for i := 0; i < 10; i++ {
   130  		fmt.Fprintf(w, "%d", i)
   131  	}
   132  	if string(r.Body()) != "0123456789" {
   133  		t.Fatalf("unexpected body %q. Expecting %q", r.Body(), "0123456789")
   134  	}
   135  }
   136  
   137  func TestResponseRawBodySet(t *testing.T) {
   138  	t.Parallel()
   139  
   140  	var resp Response
   141  
   142  	expectedS := "test"
   143  	body := []byte(expectedS)
   144  	resp.SetBodyRaw(body)
   145  
   146  	testBodyWriteTo(t, &resp, expectedS, true)
   147  }
   148  
   149  func TestResponseRawBodyReset(t *testing.T) {
   150  	t.Parallel()
   151  
   152  	var resp Response
   153  
   154  	body := []byte("test")
   155  	resp.SetBodyRaw(body)
   156  	resp.ResetBody()
   157  
   158  	testBodyWriteTo(t, &resp, "", true)
   159  }
   160  
   161  func TestResponseResetBody(t *testing.T) {
   162  	resp := Response{}
   163  	resp.BodyBuffer()
   164  	assert.NotNil(t, resp.body)
   165  	resp.maxKeepBodySize = math.MaxUint32
   166  	resp.ResetBody()
   167  	assert.NotNil(t, resp.body)
   168  	resp.maxKeepBodySize = -1
   169  	resp.ResetBody()
   170  	assert.Nil(t, resp.body)
   171  }
   172  
   173  func testResponseCopyTo(t *testing.T, src *Response) {
   174  	var dst Response
   175  	src.CopyTo(&dst)
   176  
   177  	if !reflect.DeepEqual(src, &dst) { //nolint:govet
   178  		t.Fatalf("ResponseCopyTo fail, src: \n%+v\ndst: \n%+v\n", src, &dst) //nolint:govet
   179  	}
   180  }
   181  
   182  func TestResponseMustSkipBody(t *testing.T) {
   183  	resp := Response{}
   184  	resp.SetStatusCode(consts.StatusOK)
   185  	resp.SetBodyString("test")
   186  	assert.False(t, resp.MustSkipBody())
   187  	// no content 204 means that skip body is necessary
   188  	resp.SetStatusCode(consts.StatusNoContent)
   189  	resp.ResetBody()
   190  	assert.True(t, resp.MustSkipBody())
   191  }
   192  
   193  func TestResponseBodyGunzip(t *testing.T) {
   194  	t.Parallel()
   195  	dst1 := []byte("")
   196  	src1 := []byte("hello")
   197  	res1 := compress.AppendGzipBytes(dst1, src1)
   198  	resp := Response{}
   199  	resp.SetBody(res1)
   200  	zipData, err := resp.BodyGunzip()
   201  	assert.Nil(t, err)
   202  	assert.DeepEqual(t, zipData, src1)
   203  }
   204  
   205  func TestResponseSwapResponseBody(t *testing.T) {
   206  	t.Parallel()
   207  	resp1 := Response{}
   208  	str1 := "resp1"
   209  	byteBuffer1 := &bytebufferpool.ByteBuffer{}
   210  	byteBuffer1.Set([]byte(str1))
   211  	resp1.ConstructBodyStream(byteBuffer1, bytes.NewBufferString(str1))
   212  	assert.True(t, resp1.HasBodyBytes())
   213  	resp2 := Response{}
   214  	str2 := "resp2"
   215  	byteBuffer2 := &bytebufferpool.ByteBuffer{}
   216  	byteBuffer2.Set([]byte(str2))
   217  	resp2.ConstructBodyStream(byteBuffer2, bytes.NewBufferString(str2))
   218  	SwapResponseBody(&resp1, &resp2)
   219  	assert.DeepEqual(t, resp1.body.B, []byte(str2))
   220  	assert.DeepEqual(t, resp1.BodyStream(), bytes.NewBufferString(str2))
   221  	assert.DeepEqual(t, resp2.body.B, []byte(str1))
   222  	assert.DeepEqual(t, resp2.BodyStream(), bytes.NewBufferString(str1))
   223  }
   224  
   225  func TestResponseAcquireResponse(t *testing.T) {
   226  	t.Parallel()
   227  	resp1 := AcquireResponse()
   228  	assert.NotNil(t, resp1)
   229  	resp1.SetBody([]byte("test"))
   230  	resp1.SetStatusCode(consts.StatusOK)
   231  	ReleaseResponse(resp1)
   232  	assert.Nil(t, resp1.body)
   233  }
   234  
   235  type closeBuffer struct {
   236  	*bytes.Buffer
   237  }
   238  
   239  func (b *closeBuffer) Close() error {
   240  	b.Reset()
   241  	return nil
   242  }
   243  
   244  func TestSetBodyStreamNoReset(t *testing.T) {
   245  	t.Parallel()
   246  	resp := Response{}
   247  	bsA := &closeBuffer{bytes.NewBufferString("A")}
   248  	bsB := &closeBuffer{bytes.NewBufferString("B")}
   249  	bsC := &closeBuffer{bytes.NewBufferString("C")}
   250  
   251  	resp.SetBodyStream(bsA, 1)
   252  	resp.SetBodyStreamNoReset(bsB, 1)
   253  	// resp.Body() has closed bsB
   254  	assert.DeepEqual(t, string(resp.Body()), "B")
   255  	assert.DeepEqual(t, bsA.String(), "A")
   256  
   257  	resp.bodyStream = bsA
   258  	resp.SetBodyStream(bsC, 1)
   259  	assert.DeepEqual(t, bsA.String(), "")
   260  }
   261  
   262  func TestRespSafeCopy(t *testing.T) {
   263  	resp := AcquireResponse()
   264  	resp.bodyRaw = make([]byte, 1)
   265  	resps := make([]*Response, 10)
   266  	for i := 0; i < 10; i++ {
   267  		resp.bodyRaw[0] = byte(i)
   268  		tmpResq := AcquireResponse()
   269  		resp.CopyTo(tmpResq)
   270  		resps[i] = tmpResq
   271  	}
   272  	for i := 0; i < 10; i++ {
   273  		assert.DeepEqual(t, []byte{byte(i)}, resps[i].Body())
   274  	}
   275  }
   276  
   277  func TestResponse_HijackWriter(t *testing.T) {
   278  	resp := AcquireResponse()
   279  	buf := new(bytes.Buffer)
   280  	isFinal := false
   281  	resp.HijackWriter(&mock.ExtWriter{Buf: buf, IsFinal: &isFinal})
   282  	resp.AppendBody([]byte("hello"))
   283  	assert.DeepEqual(t, 0, buf.Len())
   284  	resp.GetHijackWriter().Flush()
   285  	assert.DeepEqual(t, "hello", buf.String())
   286  	resp.AppendBodyString(", world")
   287  	assert.DeepEqual(t, "hello", buf.String())
   288  	resp.GetHijackWriter().Flush()
   289  	assert.DeepEqual(t, "hello, world", buf.String())
   290  	resp.SetBody([]byte("hello, hertz"))
   291  	resp.GetHijackWriter().Flush()
   292  	assert.DeepEqual(t, "hello, hertz", buf.String())
   293  	assert.False(t, isFinal)
   294  	resp.GetHijackWriter().Finalize()
   295  	assert.True(t, isFinal)
   296  }