// Produce is used by the source to send messages to the backend
func (f *RBForwarder) Produce(data []byte, opts map[string]interface{}, opaque interface{}) error {
if atomic.LoadUint32(&f.working) == 0 {
return errors.New("Forwarder has been closed")
}
seq := f.currentProducedID
f.currentProducedID++
m := utils.NewMessage()
r := Report{
seq: seq,
Opaque: opaque,
}
m.PushPayload(data)
m.Opts.MSet(opts)
m.Reports.Push(r)
f.p.input <- m
return nil
}
func TestBatcher(t *testing.T) {
Convey("Given a batcher", t, func() {
batcher := &Batcher{
Config: Config{
Workers: 1,
TimeoutMillis: 1000,
Limit: 10,
MaxPendingBatches: 10,
Deflate: false,
},
}
b := batcher.Spawn(0).(*Batcher)
b.clk = clock.NewMock()
Convey("When the number of workers is requested", func() {
workers := batcher.Workers()
Convey("Then the number of workers should be correct", func() {
So(workers, ShouldEqual, 1)
})
})
Convey("When a message is received with no batch group", func() {
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
d := new(Doner)
d.doneCalled = make(chan *utils.Message, 1)
d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "").Times(1)
b.OnMessage(m, d.Done)
result := <-d.doneCalled
Convey("Then the message should be sent as is", func() {
So(b.batches.Count(), ShouldEqual, 0)
payload, err := result.PopPayload()
So(err, ShouldBeNil)
So(string(payload), ShouldEqual, "Hello World")
d.AssertExpectations(t)
})
})
Convey("When a message is received by the worker, but not yet sent", func() {
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
m.Opts = cmap.New()
m.Opts.Set("batch_group", "group1")
b.OnMessage(m, nil)
Convey("Message should be present on the batch", func() {
tmp, exists := b.batches.Get("group1")
So(exists, ShouldBeTrue)
batch := tmp.(*Batch)
batch.Writer.(*bufio.Writer).Flush()
data := batch.Buf.Bytes()
So(string(data), ShouldEqual, "Hello World")
opts := batch.Message.Opts
group, ok := opts.Get("batch_group")
So(ok, ShouldBeTrue)
So(group.(string), ShouldEqual, "group1")
So(b.batches.Count(), ShouldEqual, 1)
})
})
Convey("When the max number of messages is reached", func() {
var messages []*utils.Message
for i := 0; i < int(b.Config.Limit); i++ {
m := utils.NewMessage()
m.PushPayload([]byte("ABC"))
m.Opts = cmap.New()
m.Opts.Set("batch_group", "group1")
messages = append(messages, m)
}
d := new(Doner)
d.doneCalled = make(chan *utils.Message, 1)
d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "limit").Times(1)
for i := 0; i < int(b.Config.Limit); i++ {
b.OnMessage(messages[i], d.Done)
}
Convey("Then the batch should be sent", func() {
m := <-d.doneCalled
data, err := m.PopPayload()
So(err, ShouldBeNil)
So(string(data), ShouldEqual, "ABCABCABCABCABCABCABCABCABCABC")
group1, _ := b.batches.Get("group1")
So(group1, ShouldBeNil)
So(b.batches.Count(), ShouldEqual, 0)
d.AssertExpectations(t)
})
})
Convey("When the timeout expires", func() {
var messages []*utils.Message
for i := 0; i < 5; i++ {
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
m.Opts = cmap.New()
m.Opts.Set("batch_group", "group1")
messages = append(messages, m)
}
d := new(Doner)
d.doneCalled = make(chan *utils.Message, 1)
d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "timeout").Times(1)
for i := 0; i < 5; i++ {
b.OnMessage(messages[i], d.Done)
}
clk := b.clk.(*clock.Mock)
Convey("The batch should be sent", func() {
clk.Add(500 * time.Millisecond)
group1, _ := b.batches.Get("group1")
So(group1.(*Batch), ShouldNotBeNil)
clk.Add(500 * time.Millisecond)
<-d.doneCalled
group1, _ = b.batches.Get("group1")
So(group1, ShouldBeNil)
So(b.batches.Count(), ShouldEqual, 0)
d.AssertExpectations(t)
})
})
Convey("When multiple messages are received with differents groups", func() {
m1 := utils.NewMessage()
m1.PushPayload([]byte("MESSAGE 1"))
m1.Reports.Push("Report 1")
m1.Opts = cmap.New()
m1.Opts.Set("batch_group", "group1")
m2 := utils.NewMessage()
m2.PushPayload([]byte("MESSAGE 2"))
m2.Reports.Push("Report 2")
m2.Opts = cmap.New()
m2.Opts.Set("batch_group", "group2")
m3 := utils.NewMessage()
m3.PushPayload([]byte("MESSAGE 3"))
m3.Reports.Push("Report 3")
m3.Opts = cmap.New()
m3.Opts.Set("batch_group", "group2")
d := new(Doner)
d.doneCalled = make(chan *utils.Message, 2)
d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "timeout").Times(2)
b.OnMessage(m1, d.Done)
b.OnMessage(m2, d.Done)
b.OnMessage(m3, d.Done)
Convey("Then each message should be in its group", func() {
tmp, _ := b.batches.Get("group1")
batch := tmp.(*Batch)
batch.Writer.(*bufio.Writer).Flush()
group1 := batch.Buf.Bytes()
So(string(group1), ShouldEqual, "MESSAGE 1")
tmp, _ = b.batches.Get("group2")
batch = tmp.(*Batch)
batch.Writer.(*bufio.Writer).Flush()
group2 := batch.Buf.Bytes()
So(string(group2), ShouldEqual, "MESSAGE 2MESSAGE 3")
So(b.batches.Count(), ShouldEqual, 2)
})
Convey("Then after a timeout the messages should be sent", func() {
clk := b.clk.(*clock.Mock)
So(b.batches.Count(), ShouldEqual, 2)
clk.Add(time.Duration(b.Config.TimeoutMillis) * time.Millisecond)
group1 := <-d.doneCalled
group1Data, err := group1.PopPayload()
report1 := group1.Reports.Pop().(string)
So(err, ShouldBeNil)
So(report1, ShouldEqual, "Report 1")
group2 := <-d.doneCalled
group2Data, err := group2.PopPayload()
So(err, ShouldBeNil)
report3 := group2.Reports.Pop().(string)
So(report3, ShouldEqual, "Report 3")
report2 := group2.Reports.Pop().(string)
So(report2, ShouldEqual, "Report 2")
So(string(group1Data), ShouldEqual, "MESSAGE 1")
So(string(group2Data), ShouldEqual, "MESSAGE 2MESSAGE 3")
batch, _ := b.batches.Get("group1")
So(batch, ShouldBeNil)
batch, _ = b.batches.Get("group2")
So(batch, ShouldBeNil)
So(b.batches.Count(), ShouldEqual, 0)
d.AssertExpectations(t)
})
})
})
Convey("Given a batcher with compression", t, func() {
batcher := &Batcher{
Config: Config{
TimeoutMillis: 1000,
Limit: 10,
MaxPendingBatches: 10,
Deflate: true,
},
}
b := batcher.Spawn(0).(*Batcher)
b.clk = clock.NewMock()
Convey("When the max number of messages is reached", func() {
var messages []*utils.Message
for i := 0; i < int(b.Config.Limit); i++ {
m := utils.NewMessage()
m.PushPayload([]byte("ABC"))
m.Opts = cmap.New()
m.Opts.Set("batch_group", "group1")
m.Reports.Push("Report")
messages = append(messages, m)
}
d := new(Doner)
d.doneCalled = make(chan *utils.Message, 1)
d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "limit").Times(1)
for i := 0; i < int(b.Config.Limit); i++ {
b.OnMessage(messages[i], d.Done)
}
Convey("The batch should be sent compressed", func() {
m := <-d.doneCalled
decompressed := make([]byte, 30)
data, err := m.PopPayload()
buf := bytes.NewBuffer(data)
r, err := zlib.NewReader(buf)
r.Read(decompressed)
r.Close()
So(err, ShouldBeNil)
So(string(decompressed), ShouldEqual, "ABCABCABCABCABCABCABCABCABCABC")
So(m.Reports.Size(), ShouldEqual, b.Config.Limit)
batch, _ := b.batches.Get("group1")
So(batch, ShouldBeNil)
So(b.batches.Count(), ShouldEqual, 0)
d.AssertExpectations(t)
})
})
})
}
func TestHTTPSender(t *testing.T) {
Convey("Given an Limiter with 100 messages per second without burst", t, func() {
limiter := &Limiter{
Config: Config{
MessageLimit: 100,
Burst: 1,
},
clk: clock.NewMock(),
}
limiter.Spawn(0)
Convey("When the numer of worker is requested", func() {
workers := limiter.Workers()
Convey("Should be only one", func() {
So(workers, ShouldEqual, 1)
})
})
Convey("When the limit number of messages are reached", func() {
clk := limiter.clk.(*clock.Mock)
n := Nexter{
nextCalled: make(chan *utils.Message, limiter.Config.MessageLimit*2),
}
n.On("Done", mock.AnythingOfType("*utils.Message"), 0, "")
for i := uint64(0); i < limiter.Config.MessageLimit; i++ {
limiter.OnMessage(nil, n.Done)
}
Convey("Then the limiter should be paused", func() {
So(limiter.currentMessages, ShouldEqual, limiter.Config.MessageLimit)
So(limiter.paused, ShouldBeTrue)
})
Convey("Then after 1 second the limiter should be ready again", func() {
clk.Add(1 * time.Second)
limiter.OnMessage(nil, n.Done)
So(limiter.currentMessages, ShouldEqual, 1)
So(limiter.paused, ShouldBeFalse)
})
})
})
Convey("Given an Limiter with 1000 bytes per second without burst", t, func() {
limiter := &Limiter{
Config: Config{
BytesLimit: 1000,
Burst: 1,
},
clk: clock.NewMock(),
}
limiter.Spawn(0)
Convey("When messages are sent", func() {
n := Nexter{
nextCalled: make(chan *utils.Message, 100),
}
n.On("Done", mock.AnythingOfType("*utils.Message"), 0, "")
Convey("Then the limiter should not be paused after 750 bytes", func() {
for i := uint64(0); i < 3; i++ {
m := utils.NewMessage()
payload := make([]byte, 250)
m.PushPayload(payload)
limiter.OnMessage(m, n.Done)
}
So(limiter.currentBytes, ShouldEqual, 750)
So(limiter.paused, ShouldBeFalse)
})
Convey("Then the limiter should be paused after 1000 bytes", func() {
for i := uint64(0); i < 4; i++ {
m := utils.NewMessage()
payload := make([]byte, 250)
m.PushPayload(payload)
limiter.OnMessage(m, n.Done)
}
So(limiter.currentBytes, ShouldEqual, 1000)
So(limiter.paused, ShouldBeTrue)
})
Convey("Then after 1 second the limiter should be ready again", func() {
clk := limiter.clk.(*clock.Mock)
clk.Add(1 * time.Second)
m := utils.NewMessage()
payload := make([]byte, 250)
m.PushPayload(payload)
limiter.OnMessage(m, n.Done)
So(limiter.currentBytes, ShouldEqual, 250)
So(limiter.paused, ShouldBeFalse)
})
})
})
Convey("Given a limiter with burst", t, func() {
limiter := &Limiter{
Config: Config{
MessageLimit: 100,
Burst: 2,
},
clk: clock.NewMock(),
}
limiter.Spawn(0)
clk := limiter.clk.(*clock.Mock)
clk.Add(0)
clk.Add(2 * time.Second)
Convey("When the limit number of messages are reached", func() {
n := Nexter{
nextCalled: make(chan *utils.Message, limiter.Config.MessageLimit*2),
}
n.On("Done", mock.AnythingOfType("*utils.Message"), 0, "")
for i := uint64(0); i < limiter.Config.MessageLimit; i++ {
limiter.OnMessage(nil, n.Done)
}
Convey("Then should be 2 burst available", func() {
So(len(limiter.keepSending), ShouldEqual, 2)
})
Convey("Then messages are not blocked after the limit", func() {
for i := uint64(0); i < limiter.Config.MessageLimit; i++ {
limiter.OnMessage(nil, n.Done)
}
So(limiter.currentMessages, ShouldEqual, 100)
})
Convey("Then the limiter blocks again after reaching limit a second time", func() {
So(limiter.paused, ShouldBeTrue)
})
})
})
}
func TestHTTPSender(t *testing.T) {
Convey("Given an HTTP sender with defined URL", t, func() {
sender := &HTTPSender{
Config: Config{
Workers: 1,
URL: "http://example.com",
Insecure: true,
},
}
Convey("When it is initialized", func() {
s := sender.Spawn(0).(*HTTPSender)
Convey("Then the config should be ok", func() {
So(s.Client, ShouldNotBeNil)
So(s.Workers(), ShouldEqual, 1)
})
})
Convey("When a message is sent and the response code is >= 400", func() {
var url string
s := sender.Spawn(0).(*HTTPSender)
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
s.Client = NewTestClient(401, func(req *http.Request) {
url = req.URL.String()
})
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then the reporth should contain info about the error", func() {
result := <-d.doneCalled
So(result.status, ShouldEqual, "HTTPSender error: 401 Unauthorized")
So(result.code, ShouldEqual, 2)
So(url, ShouldEqual, "http://example.com/")
d.AssertExpectations(t)
})
})
Convey("When a message is received without endpoint option", func() {
var url string
s := sender.Spawn(0).(*HTTPSender)
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
s.Client = NewTestClient(200, func(req *http.Request) {
url = req.URL.String()
})
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then the message should be sent via HTTP to the URL", func() {
result := <-d.doneCalled
So(result.status, ShouldEqual, "200 OK")
So(result.code, ShouldBeZeroValue)
So(url, ShouldEqual, "http://example.com/")
d.AssertExpectations(t)
})
})
Convey("When a message is received with endpoint option", func() {
var url string
s := sender.Spawn(0).(*HTTPSender)
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
m.Opts.Set("http_endpoint", "endpoint1")
s.Client = NewTestClient(200, func(req *http.Request) {
url = req.URL.String()
})
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then the message should be sent to the URL with endpoint as suffix", func() {
result := <-d.doneCalled
So(result.status, ShouldEqual, "200 OK")
So(result.code, ShouldBeZeroValue)
So(url, ShouldEqual, "http://example.com/endpoint1")
d.AssertExpectations(t)
})
})
Convey("When a message is received with headers", func() {
var url string
var headerValue string
s := sender.Spawn(0).(*HTTPSender)
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
m.Opts.Set("http_headers", map[string]string{
"key": "value",
})
s.Client = NewTestClient(200, func(req *http.Request) {
url = req.URL.String()
headerValue = req.Header.Get("key")
})
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then the message should be sent with headers set", func() {
result := <-d.doneCalled
So(result.status, ShouldEqual, "200 OK")
So(result.code, ShouldBeZeroValue)
So(url, ShouldEqual, "http://example.com/")
So(headerValue, ShouldEqual, "value")
d.AssertExpectations(t)
})
})
Convey("When a message without payload is received", func() {
var url string
s := sender.Spawn(0).(*HTTPSender)
m := utils.NewMessage()
s.Client = NewTestClient(200, func(req *http.Request) {
url = req.URL.String()
})
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then the message should not be sent", func() {
result := <-d.doneCalled
So(result.status, ShouldEqual, "Can't get payload of message: No payload available")
So(result.code, ShouldBeGreaterThan, 0)
So(url, ShouldBeEmpty)
d.AssertExpectations(t)
})
})
Convey("When a the HTTP client fails", func() {
s := sender.Spawn(0).(*HTTPSender)
m := utils.NewMessage()
m.PushPayload(nil)
s.Client = NewTestClient(200, func(req *http.Request) {
req.Write(nil)
})
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then the message should not be sent", func() {
result := <-d.doneCalled
So(result.status, ShouldEqual, "HTTPSender error: Post http://example.com: EOF")
So(result.code, ShouldBeGreaterThan, 0)
d.AssertExpectations(t)
})
})
})
Convey("Given an HTTP sender with invalid URL", t, func() {
sender := &HTTPSender{
Config: Config{Insecure: true},
}
s := sender.Spawn(0).(*HTTPSender)
Convey("When try to send messages", func() {
m := utils.NewMessage()
m.PushPayload([]byte("Hello World"))
m.Opts.Set("http_endpoint", "endpoint1")
d := &Doner{
doneCalled: make(chan struct {
code int
status string
}, 1),
}
d.On("Done", mock.AnythingOfType("*utils.Message"),
mock.AnythingOfType("int"), mock.AnythingOfType("string"))
s.OnMessage(m, d.Done)
Convey("Then should fail to send messages", func() {
So(s.err, ShouldNotBeNil)
result := <-d.doneCalled
So(result.status, ShouldEqual, "Invalid URL")
So(result.code, ShouldBeGreaterThan, 0)
})
})
})
}