/
request.go
116 lines (98 loc) · 2.78 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
package tgmock
import (
"strconv"
"strings"
"github.com/gotd/td/bin"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
)
func (i *Mock) request(fn func(b bin.Encoder)) *RequestBuilder {
return &RequestBuilder{
mock: i,
expectFn: fn,
times: 1,
}
}
// Expect creates builder of new expected call.
func (i *Mock) Expect() *RequestBuilder {
return i.request(nil)
}
// ExpectFunc creates builder of new expected call with given input checker.
func (i *Mock) ExpectFunc(fn func(b bin.Encoder)) *RequestBuilder {
return i.request(fn)
}
// ExpectCall creates builder of new expected call with given input.
func (i *Mock) ExpectCall(body bin.Encoder) *RequestBuilder {
return i.request(func(got bin.Encoder) {
i.assert.Equal(body, got)
})
}
// RequestBuilder is builder of expected RPC request.
type RequestBuilder struct {
mock *Mock
expectFn func(b bin.Encoder)
times int
}
// N sets count of expected calls.
func (b *RequestBuilder) N(times int) *RequestBuilder {
b.times = times
return b
}
// ThenResult adds call result to the end of call stack.
func (b *RequestBuilder) ThenResult(body bin.Encoder) *Mock {
return b.result(body, nil)
}
// ThenTrue adds call tg.BoolTrue result to the end of call stack.
func (b *RequestBuilder) ThenTrue() *Mock {
return b.result(&tg.BoolTrue{}, nil)
}
// ThenFalse adds call tg.BoolFalse result to the end of call stack.
func (b *RequestBuilder) ThenFalse() *Mock {
return b.result(&tg.BoolFalse{}, nil)
}
// ThenErr adds call result to the end of call stack.
func (b *RequestBuilder) ThenErr(err error) *Mock {
return b.result(nil, err)
}
// ThenRPCErr adds call result to the end of call stack.
func (b *RequestBuilder) ThenRPCErr(err *tgerr.Error) *Mock {
return b.ThenErr(err)
}
// ThenMigrate adds call result to the end of call stack.
func (b *RequestBuilder) ThenMigrate(typ string, arg int) *Mock {
t := strings.ToUpper(typ) + "_MIGRATE"
return b.ThenRPCErr(&tgerr.Error{
Code: 303,
Message: t + "_" + strconv.Itoa(arg),
Type: t,
Argument: arg,
})
}
// ThenFlood adds call result to the end of call stack.
func (b *RequestBuilder) ThenFlood(arg int) *Mock {
return b.ThenRPCErr(&tgerr.Error{
Code: 420,
Message: "FLOOD_WAIT_" + strconv.Itoa(arg),
Type: "FLOOD_WAIT",
Argument: arg,
})
}
// ThenUnregistered adds call result to the end of call stack.
func (b *RequestBuilder) ThenUnregistered() *Mock {
return b.ThenRPCErr(&tgerr.Error{
Code: 401,
Message: "AUTH_KEY_UNREGISTERED",
Type: "AUTH_KEY_UNREGISTERED",
})
}
func (b *RequestBuilder) result(r bin.Encoder, err error) *Mock {
for i := 0; i < b.times; i++ {
b.mock.add(HandlerFunc(func(id int64, body bin.Encoder) (bin.Encoder, error) {
if b.expectFn != nil {
b.expectFn(body)
}
return r, err
}))
}
return b.mock
}