broker_test.go 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297
  1. package sarama
  2. import (
  3. "encoding/binary"
  4. "fmt"
  5. "io"
  6. "net"
  7. "strconv"
  8. "testing"
  9. "time"
  10. )
  11. // MockBroker is a mock Kafka broker. It consists of a TCP server on a kernel-selected localhost port that
  12. // accepts a single connection. It reads Kafka requests from that connection and returns each response
  13. // from the channel provided at creation-time (if a response has a len of 0, nothing is sent, if a response
  14. // the server sleeps for 250ms instead of reading a request).
  15. //
  16. // When running tests with one of these, it is strongly recommended to specify a timeout to `go test` so that if the broker hangs
  17. // waiting for a response, the test panics.
  18. //
  19. // It is not necessary to prefix message length or correlation ID to your response bytes, the server does that
  20. // automatically as a convenience.
  21. type MockBroker struct {
  22. port int32
  23. stopper chan bool
  24. responses chan []byte
  25. listener net.Listener
  26. t *testing.T
  27. }
  28. func (b *MockBroker) Port() int32 {
  29. return b.port
  30. }
  31. func (b *MockBroker) Addr() string {
  32. return b.listener.Addr().String()
  33. }
  34. // Close closes the response channel originally provided, then waits to make sure
  35. // that all requests/responses matched up before exiting.
  36. func (b *MockBroker) Close() {
  37. close(b.responses)
  38. <-b.stopper
  39. }
  40. func (b *MockBroker) serverLoop() {
  41. defer close(b.stopper)
  42. conn, err := b.listener.Accept()
  43. if err != nil {
  44. b.t.Error(err)
  45. conn.Close()
  46. b.listener.Close()
  47. return
  48. }
  49. reqHeader := make([]byte, 4)
  50. resHeader := make([]byte, 8)
  51. for response := range b.responses {
  52. if response == nil {
  53. time.Sleep(250 * time.Millisecond)
  54. continue
  55. }
  56. _, err := io.ReadFull(conn, reqHeader)
  57. if err != nil {
  58. b.t.Error(err)
  59. conn.Close()
  60. b.listener.Close()
  61. return
  62. }
  63. body := make([]byte, binary.BigEndian.Uint32(reqHeader))
  64. if len(body) < 10 {
  65. b.t.Error("Kafka request too short.")
  66. conn.Close()
  67. b.listener.Close()
  68. return
  69. }
  70. _, err = io.ReadFull(conn, body)
  71. if err != nil {
  72. b.t.Error(err)
  73. conn.Close()
  74. b.listener.Close()
  75. return
  76. }
  77. if len(response) == 0 {
  78. continue
  79. }
  80. binary.BigEndian.PutUint32(resHeader, uint32(len(response)+4))
  81. binary.BigEndian.PutUint32(resHeader[4:], binary.BigEndian.Uint32(body[4:]))
  82. _, err = conn.Write(resHeader)
  83. if err != nil {
  84. b.t.Error(err)
  85. conn.Close()
  86. b.listener.Close()
  87. return
  88. }
  89. _, err = conn.Write(response)
  90. if err != nil {
  91. b.t.Error(err)
  92. conn.Close()
  93. b.listener.Close()
  94. return
  95. }
  96. }
  97. err = conn.Close()
  98. if err != nil {
  99. b.t.Error(err)
  100. b.listener.Close()
  101. return
  102. }
  103. err = b.listener.Close()
  104. if err != nil {
  105. b.t.Error(err)
  106. return
  107. }
  108. }
  109. // NewMockBroker launches a fake Kafka broker. It takes a testing.T as provided by the test framework and a channel of responses to use.
  110. // If an error occurs it is simply logged to the testing.T and the broker exits.
  111. func NewMockBroker(t *testing.T, responses chan []byte) *MockBroker {
  112. var err error
  113. broker := new(MockBroker)
  114. broker.stopper = make(chan bool)
  115. broker.responses = responses
  116. broker.t = t
  117. broker.listener, err = net.Listen("tcp", "localhost:0")
  118. if err != nil {
  119. t.Fatal(err)
  120. }
  121. _, portStr, err := net.SplitHostPort(broker.listener.Addr().String())
  122. if err != nil {
  123. t.Fatal(err)
  124. }
  125. tmp, err := strconv.ParseInt(portStr, 10, 32)
  126. if err != nil {
  127. t.Fatal(err)
  128. }
  129. broker.port = int32(tmp)
  130. go broker.serverLoop()
  131. return broker
  132. }
  133. func ExampleBroker() error {
  134. broker := NewBroker("localhost:9092")
  135. err := broker.Open(4)
  136. if err != nil {
  137. return err
  138. }
  139. defer broker.Close()
  140. request := MetadataRequest{Topics: []string{"myTopic"}}
  141. response, err := broker.GetMetadata("myClient", &request)
  142. if err != nil {
  143. return err
  144. }
  145. fmt.Println("There are", len(response.Topics), "topics active in the cluster.")
  146. return nil
  147. }
  148. func TestBrokerAccessors(t *testing.T) {
  149. broker := NewBroker("abc:123")
  150. if broker.ID() != -1 {
  151. t.Error("New broker didn't have an ID of -1.")
  152. }
  153. if broker.Addr() != "abc:123" {
  154. t.Error("New broker didn't have the correct address")
  155. }
  156. broker.id = 34
  157. if broker.ID() != 34 {
  158. t.Error("Manually setting broker ID did not take effect.")
  159. }
  160. }
  161. func TestSimpleBrokerCommunication(t *testing.T) {
  162. responses := make(chan []byte)
  163. mockBroker := NewMockBroker(t, responses)
  164. defer mockBroker.Close()
  165. broker := NewBroker(mockBroker.Addr())
  166. err := broker.Open(4)
  167. if err != nil {
  168. t.Fatal(err)
  169. }
  170. go func() {
  171. for _, tt := range brokerTestTable {
  172. responses <- tt.response
  173. }
  174. }()
  175. for _, tt := range brokerTestTable {
  176. tt.runner(t, broker)
  177. }
  178. err = broker.Close()
  179. if err != nil {
  180. t.Error(err)
  181. }
  182. }
  183. // We're not testing encoding/decoding here, so most of the requests/responses will be empty for simplicity's sake
  184. var brokerTestTable = []struct {
  185. response []byte
  186. runner func(*testing.T, *Broker)
  187. }{
  188. {[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
  189. func(t *testing.T, broker *Broker) {
  190. request := MetadataRequest{}
  191. response, err := broker.GetMetadata("clientID", &request)
  192. if err != nil {
  193. t.Error(err)
  194. }
  195. if response == nil {
  196. t.Error("Metadata request got no response!")
  197. }
  198. }},
  199. {[]byte{},
  200. func(t *testing.T, broker *Broker) {
  201. request := ProduceRequest{}
  202. request.RequiredAcks = NoResponse
  203. response, err := broker.Produce("clientID", &request)
  204. if err != nil {
  205. t.Error(err)
  206. }
  207. if response != nil {
  208. t.Error("Produce request with NoResponse got a response!")
  209. }
  210. }},
  211. {[]byte{0x00, 0x00, 0x00, 0x00},
  212. func(t *testing.T, broker *Broker) {
  213. request := ProduceRequest{}
  214. request.RequiredAcks = WaitForLocal
  215. response, err := broker.Produce("clientID", &request)
  216. if err != nil {
  217. t.Error(err)
  218. }
  219. if response == nil {
  220. t.Error("Produce request without NoResponse got no response!")
  221. }
  222. }},
  223. {[]byte{0x00, 0x00, 0x00, 0x00},
  224. func(t *testing.T, broker *Broker) {
  225. request := FetchRequest{}
  226. response, err := broker.Fetch("clientID", &request)
  227. if err != nil {
  228. t.Error(err)
  229. }
  230. if response == nil {
  231. t.Error("Fetch request got no response!")
  232. }
  233. }},
  234. {[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
  235. func(t *testing.T, broker *Broker) {
  236. request := OffsetFetchRequest{}
  237. response, err := broker.FetchOffset("clientID", &request)
  238. if err != nil {
  239. t.Error(err)
  240. }
  241. if response == nil {
  242. t.Error("OffsetFetch request got no response!")
  243. }
  244. }},
  245. {[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
  246. func(t *testing.T, broker *Broker) {
  247. request := OffsetCommitRequest{}
  248. response, err := broker.CommitOffset("clientID", &request)
  249. if err != nil {
  250. t.Error(err)
  251. }
  252. if response == nil {
  253. t.Error("OffsetCommit request got no response!")
  254. }
  255. }},
  256. {[]byte{0x00, 0x00, 0x00, 0x00},
  257. func(t *testing.T, broker *Broker) {
  258. request := OffsetRequest{}
  259. response, err := broker.GetAvailableOffsets("clientID", &request)
  260. if err != nil {
  261. t.Error(err)
  262. }
  263. if response == nil {
  264. t.Error("Offset request got no response!")
  265. }
  266. }},
  267. }