forked from fatedier/frp
/
request.go
130 lines (109 loc) · 2.73 KB
/
request.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
package framework
import (
"bytes"
"net/http"
flog "github.com/fatedier/frp/pkg/util/log"
"github.com/fatedier/frp/test/e2e/framework/consts"
"github.com/fatedier/frp/test/e2e/pkg/request"
)
func SpecifiedHTTPBodyHandler(body []byte) http.HandlerFunc {
return func(w http.ResponseWriter, req *http.Request) {
_, _ = w.Write(body)
}
}
func ExpectResponseCode(code int) EnsureFunc {
return func(resp *request.Response) bool {
if resp.Code == code {
return true
}
flog.Warn("Expect code %d, but got %d", code, resp.Code)
return false
}
}
// NewRequest return a default request with default timeout and content.
func NewRequest() *request.Request {
return request.New().
Timeout(consts.DefaultTimeout).
Body([]byte(consts.TestString))
}
func NewHTTPRequest() *request.Request {
return request.New().HTTP().HTTPParams("GET", "", "/", nil)
}
type RequestExpect struct {
req *request.Request
f *Framework
expectResp []byte
expectError bool
explain []interface{}
}
func NewRequestExpect(f *Framework) *RequestExpect {
return &RequestExpect{
req: NewRequest(),
f: f,
expectResp: []byte(consts.TestString),
expectError: false,
explain: make([]interface{}, 0),
}
}
func (e *RequestExpect) Request(req *request.Request) *RequestExpect {
e.req = req
return e
}
func (e *RequestExpect) RequestModify(f func(r *request.Request)) *RequestExpect {
f(e.req)
return e
}
func (e *RequestExpect) Protocol(protocol string) *RequestExpect {
e.req.Protocol(protocol)
return e
}
func (e *RequestExpect) PortName(name string) *RequestExpect {
if e.f != nil {
e.req.Port(e.f.PortByName(name))
}
return e
}
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
}
func (e *RequestExpect) ExpectError(expectErr bool) *RequestExpect {
e.expectError = expectErr
return e
}
func (e *RequestExpect) Explain(explain ...interface{}) *RequestExpect {
e.explain = explain
return e
}
type EnsureFunc func(*request.Response) bool
func (e *RequestExpect) Ensure(fns ...EnsureFunc) {
ret, err := e.req.Do()
if e.expectError {
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, ret.Content, e.expectResp, e.explain...)
} else {
for _, fn := range fns {
ok := fn(ret)
if !ok {
flog.Trace("Response info: %+v", ret)
}
ExpectTrueWithOffset(1, ok, e.explain...)
}
}
}
func (e *RequestExpect) Do() (*request.Response, error) {
return e.req.Do()
}