2024-02-10 23:13:31 +00:00
|
|
|
package main
|
|
|
|
|
|
|
|
import (
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"sync"
|
|
|
|
"testing"
|
|
|
|
)
|
|
|
|
|
|
|
|
type MockSender struct {
|
|
|
|
messages []GenericMessage
|
|
|
|
}
|
|
|
|
|
|
|
|
func (self *MockSender) Send(msg GenericMessage) error {
|
|
|
|
self.messages = append(self.messages, msg)
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2024-02-11 10:51:41 +00:00
|
|
|
type MockSleeper struct {
|
2024-02-11 13:56:03 +00:00
|
|
|
ch *chan<- interface{}
|
|
|
|
callCount int
|
2024-02-11 10:51:41 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (self *MockSleeper) Sleep(ch chan<- interface{}) {
|
|
|
|
if self.ch == nil {
|
|
|
|
self.ch = &ch
|
|
|
|
}
|
2024-02-11 13:56:03 +00:00
|
|
|
self.callCount += 1
|
2024-02-11 10:51:41 +00:00
|
|
|
}
|
|
|
|
|
2024-02-10 23:13:31 +00:00
|
|
|
func TestMessageThrottling(t *testing.T) {
|
|
|
|
testEvents := make(chan GenericMessage)
|
|
|
|
wg := sync.WaitGroup{}
|
|
|
|
mockSender := MockSender{make([]GenericMessage, 0)}
|
2024-02-11 13:56:03 +00:00
|
|
|
mockSleeper := MockSleeper{nil, 0}
|
|
|
|
go tgSenderWorker(testEvents, &mockSender, &wg, &mockSleeper)
|
2024-02-10 23:13:31 +00:00
|
|
|
testEvents <- GenericMessage{123, "test1"}
|
|
|
|
testEvents <- GenericMessage{124, "test3"}
|
|
|
|
testEvents <- GenericMessage{123, "test2"}
|
|
|
|
testEvents <- GenericMessage{124, "test4"}
|
2024-02-11 13:56:03 +00:00
|
|
|
assert.Equal(t, 1, mockSleeper.callCount)
|
2024-02-11 10:51:41 +00:00
|
|
|
*mockSleeper.ch <- nil
|
2024-02-11 13:56:03 +00:00
|
|
|
assert.Equal(t, 1, mockSleeper.callCount)
|
2024-02-10 23:13:31 +00:00
|
|
|
testEvents <- GenericMessage{123, "test5"}
|
|
|
|
close(testEvents)
|
2024-02-10 23:24:11 +00:00
|
|
|
wg.Wait()
|
2024-02-11 13:56:03 +00:00
|
|
|
assert.Equal(t, 2, mockSleeper.callCount)
|
2024-02-10 23:13:31 +00:00
|
|
|
|
|
|
|
assert.Len(t, mockSender.messages, 3)
|
|
|
|
assert.Contains(t, mockSender.messages, GenericMessage{123, "test1\ntest2\n"})
|
|
|
|
}
|