github.com/cloudwego/hertz@v0.9.3/pkg/common/ut/response.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  
    17  package ut
    18  
    19  import (
    20  	"bytes"
    21  
    22  	"github.com/cloudwego/hertz/pkg/protocol"
    23  	"github.com/cloudwego/hertz/pkg/protocol/consts"
    24  )
    25  
    26  // ResponseRecorder records handler's response for later test
    27  type ResponseRecorder struct {
    28  	// Code is the HTTP response code set by WriteHeader.
    29  	//
    30  	// Note that if a Handler never calls WriteHeader or Write,
    31  	// this might end up being 0, rather than the implicit
    32  	// http.StatusOK. To get the implicit value, use the Result
    33  	// method.
    34  	Code int
    35  
    36  	// header contains the headers explicitly set by the Handler.
    37  	// It is an internal detail.
    38  	header *protocol.ResponseHeader
    39  
    40  	// Body is the buffer to which the Handler's Write calls are sent.
    41  	// If nil, the Writes are silently discarded.
    42  	Body *bytes.Buffer
    43  
    44  	// Flushed is whether the Handler called Flush.
    45  	Flushed bool
    46  
    47  	result      *protocol.Response // cache of Result's return value
    48  	wroteHeader bool
    49  }
    50  
    51  // NewRecorder returns an initialized ResponseRecorder.
    52  func NewRecorder() *ResponseRecorder {
    53  	return &ResponseRecorder{
    54  		header: new(protocol.ResponseHeader),
    55  		Body:   new(bytes.Buffer),
    56  		Code:   consts.StatusOK,
    57  	}
    58  }
    59  
    60  // Header returns the response headers to mutate within a handler.
    61  // To test the headers that were written after a handler completes,
    62  // use the Result method and see the returned Response value's Header.
    63  func (rw *ResponseRecorder) Header() *protocol.ResponseHeader {
    64  	m := rw.header
    65  	if m == nil {
    66  		m = new(protocol.ResponseHeader)
    67  		rw.header = m
    68  	}
    69  	return m
    70  }
    71  
    72  // Write implements io.Writer. The data in buf is written to
    73  // rw.Body, if not nil.
    74  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
    75  	if !rw.wroteHeader {
    76  		rw.WriteHeader(consts.StatusOK)
    77  	}
    78  	if rw.Body != nil {
    79  		rw.Body.Write(buf)
    80  	}
    81  	return len(buf), nil
    82  }
    83  
    84  // WriteString implements io.StringWriter. The data in str is written
    85  // to rw.Body, if not nil.
    86  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
    87  	if !rw.wroteHeader {
    88  		rw.WriteHeader(consts.StatusOK)
    89  	}
    90  	if rw.Body != nil {
    91  		rw.Body.WriteString(str)
    92  	}
    93  	return len(str), nil
    94  }
    95  
    96  // WriteHeader sends an HTTP response header with the provided
    97  // status code.
    98  func (rw *ResponseRecorder) WriteHeader(code int) {
    99  	if rw.wroteHeader {
   100  		return
   101  	}
   102  	if rw.header == nil {
   103  		rw.header = new(protocol.ResponseHeader)
   104  	}
   105  	rw.header.SetStatusCode(code)
   106  	rw.Code = code
   107  	rw.wroteHeader = true
   108  }
   109  
   110  // Flush implements http.Flusher. To test whether Flush was
   111  // called, see rw.Flushed.
   112  func (rw *ResponseRecorder) Flush() {
   113  	if !rw.wroteHeader {
   114  		rw.WriteHeader(consts.StatusOK)
   115  	}
   116  	rw.Flushed = true
   117  }
   118  
   119  // Result returns the response generated by the handler.
   120  //
   121  // The returned Response will have at least its StatusCode,
   122  // Header, Body, and optionally Trailer populated.
   123  // More fields may be populated in the future, so callers should
   124  // not DeepEqual the result in tests.
   125  //
   126  // The Response.Header is a snapshot of the headers at the time of the
   127  // first write call, or at the time of this call, if the handler never
   128  // did a write.
   129  //
   130  // The Response.Body is guaranteed to be non-nil and Body.Read call is
   131  // guaranteed to not return any error other than io.EOF.
   132  //
   133  // Result must only be called after the handler has finished running.
   134  func (rw *ResponseRecorder) Result() *protocol.Response {
   135  	if rw.result != nil {
   136  		return rw.result
   137  	}
   138  
   139  	res := new(protocol.Response)
   140  	h := rw.Header()
   141  	h.CopyTo(&res.Header)
   142  	if rw.Body != nil {
   143  		b := rw.Body.Bytes()
   144  		res.SetBody(b)
   145  		res.Header.SetContentLength(len(b))
   146  	}
   147  
   148  	rw.result = res
   149  	return res
   150  }