blob: 55ef1e2920f6c4b89c4ff916a5d67650a4f8ddc5 [file] [log] [blame]
Scott Bakere7144bc2019-10-01 14:16:47 -07001package sarama
2
3import (
4 "bytes"
5 "encoding/binary"
6 "fmt"
7 "io"
8 "net"
9 "reflect"
10 "strconv"
11 "sync"
12 "time"
13
14 "github.com/davecgh/go-spew/spew"
15)
16
17const (
18 expectationTimeout = 500 * time.Millisecond
19)
20
21type requestHandlerFunc func(req *request) (res encoder)
22
23// RequestNotifierFunc is invoked when a mock broker processes a request successfully
24// and will provides the number of bytes read and written.
25type RequestNotifierFunc func(bytesRead, bytesWritten int)
26
27// MockBroker is a mock Kafka broker that is used in unit tests. It is exposed
28// to facilitate testing of higher level or specialized consumers and producers
29// built on top of Sarama. Note that it does not 'mimic' the Kafka API protocol,
30// but rather provides a facility to do that. It takes care of the TCP
31// transport, request unmarshaling, response marshaling, and makes it the test
32// writer responsibility to program correct according to the Kafka API protocol
33// MockBroker behaviour.
34//
35// MockBroker is implemented as a TCP server listening on a kernel-selected
36// localhost port that can accept many connections. It reads Kafka requests
37// from that connection and returns responses programmed by the SetHandlerByMap
38// function. If a MockBroker receives a request that it has no programmed
39// response for, then it returns nothing and the request times out.
40//
41// A set of MockRequest builders to define mappings used by MockBroker is
42// provided by Sarama. But users can develop MockRequests of their own and use
43// them along with or instead of the standard ones.
44//
45// When running tests with MockBroker it is strongly recommended to specify
46// a timeout to `go test` so that if the broker hangs waiting for a response,
47// the test panics.
48//
49// It is not necessary to prefix message length or correlation ID to your
50// response bytes, the server does that automatically as a convenience.
51type MockBroker struct {
52 brokerID int32
53 port int32
54 closing chan none
55 stopper chan none
56 expectations chan encoder
57 listener net.Listener
58 t TestReporter
59 latency time.Duration
60 handler requestHandlerFunc
61 notifier RequestNotifierFunc
62 history []RequestResponse
63 lock sync.Mutex
64}
65
66// RequestResponse represents a Request/Response pair processed by MockBroker.
67type RequestResponse struct {
68 Request protocolBody
69 Response encoder
70}
71
72// SetLatency makes broker pause for the specified period every time before
73// replying.
74func (b *MockBroker) SetLatency(latency time.Duration) {
75 b.latency = latency
76}
77
78// SetHandlerByMap defines mapping of Request types to MockResponses. When a
79// request is received by the broker, it looks up the request type in the map
80// and uses the found MockResponse instance to generate an appropriate reply.
81// If the request type is not found in the map then nothing is sent.
82func (b *MockBroker) SetHandlerByMap(handlerMap map[string]MockResponse) {
83 b.setHandler(func(req *request) (res encoder) {
84 reqTypeName := reflect.TypeOf(req.body).Elem().Name()
85 mockResponse := handlerMap[reqTypeName]
86 if mockResponse == nil {
87 return nil
88 }
89 return mockResponse.For(req.body)
90 })
91}
92
93// SetNotifier set a function that will get invoked whenever a request has been
94// processed successfully and will provide the number of bytes read and written
95func (b *MockBroker) SetNotifier(notifier RequestNotifierFunc) {
96 b.lock.Lock()
97 b.notifier = notifier
98 b.lock.Unlock()
99}
100
101// BrokerID returns broker ID assigned to the broker.
102func (b *MockBroker) BrokerID() int32 {
103 return b.brokerID
104}
105
106// History returns a slice of RequestResponse pairs in the order they were
107// processed by the broker. Note that in case of multiple connections to the
108// broker the order expected by a test can be different from the order recorded
109// in the history, unless some synchronization is implemented in the test.
110func (b *MockBroker) History() []RequestResponse {
111 b.lock.Lock()
112 history := make([]RequestResponse, len(b.history))
113 copy(history, b.history)
114 b.lock.Unlock()
115 return history
116}
117
118// Port returns the TCP port number the broker is listening for requests on.
119func (b *MockBroker) Port() int32 {
120 return b.port
121}
122
123// Addr returns the broker connection string in the form "<address>:<port>".
124func (b *MockBroker) Addr() string {
125 return b.listener.Addr().String()
126}
127
128// Close terminates the broker blocking until it stops internal goroutines and
129// releases all resources.
130func (b *MockBroker) Close() {
131 close(b.expectations)
132 if len(b.expectations) > 0 {
133 buf := bytes.NewBufferString(fmt.Sprintf("mockbroker/%d: not all expectations were satisfied! Still waiting on:\n", b.BrokerID()))
134 for e := range b.expectations {
135 _, _ = buf.WriteString(spew.Sdump(e))
136 }
137 b.t.Error(buf.String())
138 }
139 close(b.closing)
140 <-b.stopper
141}
142
143// setHandler sets the specified function as the request handler. Whenever
144// a mock broker reads a request from the wire it passes the request to the
145// function and sends back whatever the handler function returns.
146func (b *MockBroker) setHandler(handler requestHandlerFunc) {
147 b.lock.Lock()
148 b.handler = handler
149 b.lock.Unlock()
150}
151
152func (b *MockBroker) serverLoop() {
153 defer close(b.stopper)
154 var err error
155 var conn net.Conn
156
157 go func() {
158 <-b.closing
159 err := b.listener.Close()
160 if err != nil {
161 b.t.Error(err)
162 }
163 }()
164
165 wg := &sync.WaitGroup{}
166 i := 0
167 for conn, err = b.listener.Accept(); err == nil; conn, err = b.listener.Accept() {
168 wg.Add(1)
169 go b.handleRequests(conn, i, wg)
170 i++
171 }
172 wg.Wait()
173 Logger.Printf("*** mockbroker/%d: listener closed, err=%v", b.BrokerID(), err)
174}
175
176func (b *MockBroker) handleRequests(conn net.Conn, idx int, wg *sync.WaitGroup) {
177 defer wg.Done()
178 defer func() {
179 _ = conn.Close()
180 }()
181 Logger.Printf("*** mockbroker/%d/%d: connection opened", b.BrokerID(), idx)
182 var err error
183
184 abort := make(chan none)
185 defer close(abort)
186 go func() {
187 select {
188 case <-b.closing:
189 _ = conn.Close()
190 case <-abort:
191 }
192 }()
193
194 resHeader := make([]byte, 8)
195 for {
196 req, bytesRead, err := decodeRequest(conn)
197 if err != nil {
198 Logger.Printf("*** mockbroker/%d/%d: invalid request: err=%+v, %+v", b.brokerID, idx, err, spew.Sdump(req))
199 b.serverError(err)
200 break
201 }
202
203 if b.latency > 0 {
204 time.Sleep(b.latency)
205 }
206
207 b.lock.Lock()
208 res := b.handler(req)
209 b.history = append(b.history, RequestResponse{req.body, res})
210 b.lock.Unlock()
211
212 if res == nil {
213 Logger.Printf("*** mockbroker/%d/%d: ignored %v", b.brokerID, idx, spew.Sdump(req))
214 continue
215 }
216 Logger.Printf("*** mockbroker/%d/%d: served %v -> %v", b.brokerID, idx, req, res)
217
218 encodedRes, err := encode(res, nil)
219 if err != nil {
220 b.serverError(err)
221 break
222 }
223 if len(encodedRes) == 0 {
224 b.lock.Lock()
225 if b.notifier != nil {
226 b.notifier(bytesRead, 0)
227 }
228 b.lock.Unlock()
229 continue
230 }
231
232 binary.BigEndian.PutUint32(resHeader, uint32(len(encodedRes)+4))
233 binary.BigEndian.PutUint32(resHeader[4:], uint32(req.correlationID))
234 if _, err = conn.Write(resHeader); err != nil {
235 b.serverError(err)
236 break
237 }
238 if _, err = conn.Write(encodedRes); err != nil {
239 b.serverError(err)
240 break
241 }
242
243 b.lock.Lock()
244 if b.notifier != nil {
245 b.notifier(bytesRead, len(resHeader)+len(encodedRes))
246 }
247 b.lock.Unlock()
248 }
249 Logger.Printf("*** mockbroker/%d/%d: connection closed, err=%v", b.BrokerID(), idx, err)
250}
251
252func (b *MockBroker) defaultRequestHandler(req *request) (res encoder) {
253 select {
254 case res, ok := <-b.expectations:
255 if !ok {
256 return nil
257 }
258 return res
259 case <-time.After(expectationTimeout):
260 return nil
261 }
262}
263
264func (b *MockBroker) serverError(err error) {
265 isConnectionClosedError := false
266 if _, ok := err.(*net.OpError); ok {
267 isConnectionClosedError = true
268 } else if err == io.EOF {
269 isConnectionClosedError = true
270 } else if err.Error() == "use of closed network connection" {
271 isConnectionClosedError = true
272 }
273
274 if isConnectionClosedError {
275 return
276 }
277
278 b.t.Errorf(err.Error())
279}
280
281// NewMockBroker launches a fake Kafka broker. It takes a TestReporter as provided by the
282// test framework and a channel of responses to use. If an error occurs it is
283// simply logged to the TestReporter and the broker exits.
284func NewMockBroker(t TestReporter, brokerID int32) *MockBroker {
285 return NewMockBrokerAddr(t, brokerID, "localhost:0")
286}
287
288// NewMockBrokerAddr behaves like newMockBroker but listens on the address you give
289// it rather than just some ephemeral port.
290func NewMockBrokerAddr(t TestReporter, brokerID int32, addr string) *MockBroker {
291 listener, err := net.Listen("tcp", addr)
292 if err != nil {
293 t.Fatal(err)
294 }
295 return NewMockBrokerListener(t, brokerID, listener)
296}
297
298// NewMockBrokerListener behaves like newMockBrokerAddr but accepts connections on the listener specified.
299func NewMockBrokerListener(t TestReporter, brokerID int32, listener net.Listener) *MockBroker {
300 var err error
301
302 broker := &MockBroker{
303 closing: make(chan none),
304 stopper: make(chan none),
305 t: t,
306 brokerID: brokerID,
307 expectations: make(chan encoder, 512),
308 listener: listener,
309 }
310 broker.handler = broker.defaultRequestHandler
311
312 Logger.Printf("*** mockbroker/%d listening on %s\n", brokerID, broker.listener.Addr().String())
313 _, portStr, err := net.SplitHostPort(broker.listener.Addr().String())
314 if err != nil {
315 t.Fatal(err)
316 }
317 tmp, err := strconv.ParseInt(portStr, 10, 32)
318 if err != nil {
319 t.Fatal(err)
320 }
321 broker.port = int32(tmp)
322
323 go broker.serverLoop()
324
325 return broker
326}
327
328func (b *MockBroker) Returns(e encoder) {
329 b.expectations <- e
330}