frp/test/e2e/framework/request.go

131 lines
2.7 KiB
Go
Raw Normal View History

2020-06-02 22:48:55 +08:00
package framework
import (
2021-06-18 16:48:36 +08:00
"bytes"
"net/http"
flog "github.com/fatedier/frp/pkg/util/log"
2021-03-31 16:57:39 +08:00
"github.com/fatedier/frp/test/e2e/framework/consts"
2020-06-02 22:48:55 +08:00
"github.com/fatedier/frp/test/e2e/pkg/request"
)
2021-06-18 16:48:36 +08:00
func SpecifiedHTTPBodyHandler(body []byte) http.HandlerFunc {
return func(w http.ResponseWriter, req *http.Request) {
w.Write(body)
2020-09-07 14:57:23 +08:00
}
}
2021-06-18 16:48:36 +08:00
func ExpectResponseCode(code int) EnsureFunc {
return func(resp *request.Response) bool {
2021-06-21 19:27:26 +08:00
if resp.Code == code {
return true
}
flog.Warn("Expect code %d, but got %d", code, resp.Code)
return false
2020-09-07 14:57:23 +08:00
}
}
2021-06-18 16:48:36 +08:00
// NewRequest return a default request with default timeout and content.
2021-03-31 16:57:39 +08:00
func NewRequest() *request.Request {
return request.New().
Timeout(consts.DefaultTimeout).
Body([]byte(consts.TestString))
}
2021-06-18 16:48:36 +08:00
func NewHTTPRequest() *request.Request {
return request.New().HTTP().HTTPParams("GET", "", "/", nil)
2020-09-07 14:57:23 +08:00
}
2021-03-31 16:57:39 +08:00
type RequestExpect struct {
req *request.Request
f *Framework
expectResp []byte
expectError bool
explain []interface{}
2020-09-07 14:57:23 +08:00
}
2021-03-31 16:57:39 +08:00
func NewRequestExpect(f *Framework) *RequestExpect {
return &RequestExpect{
req: NewRequest(),
f: f,
expectResp: []byte(consts.TestString),
expectError: false,
explain: make([]interface{}, 0),
}
}
2021-08-02 13:07:28 +08:00
func (e *RequestExpect) Request(req *request.Request) *RequestExpect {
e.req = req
return e
}
func (e *RequestExpect) RequestModify(f func(r *request.Request)) *RequestExpect {
2021-03-31 16:57:39 +08:00
f(e.req)
return e
}
2021-06-18 16:48:36 +08:00
func (e *RequestExpect) Protocol(protocol string) *RequestExpect {
e.req.Protocol(protocol)
return e
}
2021-03-31 16:57:39 +08:00
func (e *RequestExpect) PortName(name string) *RequestExpect {
if e.f != nil {
e.req.Port(e.f.PortByName(name))
}
return e
}
2021-06-18 16:48:36 +08:00
func (e *RequestExpect) Port(port int) *RequestExpect {
if e.f != nil {
e.req.Port(port)
}
return e
}
func (e *RequestExpect) ExpectResp(resp []byte) *RequestExpect {
e.expectResp = resp
return e
}
2021-03-31 16:57:39 +08:00
func (e *RequestExpect) ExpectError(expectErr bool) *RequestExpect {
e.expectError = expectErr
return e
}
func (e *RequestExpect) Explain(explain ...interface{}) *RequestExpect {
e.explain = explain
return e
}
2021-06-18 16:48:36 +08:00
type EnsureFunc func(*request.Response) bool
func (e *RequestExpect) Ensure(fns ...EnsureFunc) {
ret, err := e.req.Do()
2021-03-31 16:57:39 +08:00
if e.expectError {
2021-06-18 16:48:36 +08:00
ExpectErrorWithOffset(1, err, e.explain...)
return
}
ExpectNoErrorWithOffset(1, err, e.explain...)
if len(fns) == 0 {
if !bytes.Equal(e.expectResp, ret.Content) {
flog.Trace("Response info: %+v", ret)
}
ExpectEqualValuesWithOffset(1, e.expectResp, ret.Content, e.explain...)
2021-03-31 16:57:39 +08:00
} else {
2021-06-18 16:48:36 +08:00
for _, fn := range fns {
ok := fn(ret)
if !ok {
flog.Trace("Response info: %+v", ret)
}
ExpectTrueWithOffset(1, ok, e.explain...)
}
2021-03-31 16:57:39 +08:00
}
2020-06-02 22:48:55 +08:00
}
2021-06-18 16:48:36 +08:00
func (e *RequestExpect) Do() (*request.Response, error) {
return e.req.Do()
}