github.com/cloudwego/hertz@v0.9.3/pkg/common/ut/context.go (about)

     1  /*
     2   * Copyright 2023 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  	"io"
    21  	"io/ioutil"
    22  
    23  	"github.com/cloudwego/hertz/pkg/app"
    24  	"github.com/cloudwego/hertz/pkg/common/config"
    25  	"github.com/cloudwego/hertz/pkg/protocol"
    26  	"github.com/cloudwego/hertz/pkg/route"
    27  )
    28  
    29  // CreateUtRequestContext returns an app.RequestContext for testing purposes
    30  func CreateUtRequestContext(method, url string, body *Body, headers ...Header) *app.RequestContext {
    31  	engine := route.NewEngine(config.NewOptions([]config.Option{}))
    32  	return createUtRequestContext(engine, method, url, body, headers...)
    33  }
    34  
    35  func createUtRequestContext(engine *route.Engine, method, url string, body *Body, headers ...Header) *app.RequestContext {
    36  	ctx := engine.NewContext()
    37  
    38  	var r *protocol.Request
    39  	if body != nil && body.Body != nil {
    40  		r = protocol.NewRequest(method, url, body.Body)
    41  		r.CopyTo(&ctx.Request)
    42  		if engine.IsStreamRequestBody() || body.Len == -1 {
    43  			ctx.Request.SetBodyStream(body.Body, body.Len)
    44  		} else {
    45  			buf, err := ioutil.ReadAll(&io.LimitedReader{R: body.Body, N: int64(body.Len)})
    46  			ctx.Request.SetBody(buf)
    47  			if err != nil && err != io.EOF {
    48  				panic(err)
    49  			}
    50  		}
    51  	} else {
    52  		r = protocol.NewRequest(method, url, nil)
    53  		r.CopyTo(&ctx.Request)
    54  	}
    55  
    56  	for _, v := range headers {
    57  		if ctx.Request.Header.Get(v.Key) != "" {
    58  			ctx.Request.Header.Add(v.Key, v.Value)
    59  		} else {
    60  			ctx.Request.Header.Set(v.Key, v.Value)
    61  		}
    62  	}
    63  
    64  	return ctx
    65  }