package main import ( "io" "log" "sync" "testing" "github.com/stretchr/testify/assert" ) type MockSender struct { messages []GenericMessage } func (self *MockSender) Send(msg GenericMessage) error { self.messages = append(self.messages, msg) return nil } type MockSleeper struct { ch *chan<- interface{} callCount int } func (self *MockSleeper) Sleep(ch chan<- interface{}) { if self.ch == nil { self.ch = &ch } self.callCount += 1 } type FakeChatId struct { FakeId int64 } func (self FakeChatId) GetTgIds() *[]int64 { return nil } func TestMessageThrottling(t *testing.T) { testEvents := make(chan GenericMessage) wg := sync.WaitGroup{} mockSender := MockSender{make([]GenericMessage, 0)} mockSleeper := MockSleeper{nil, 0} Consume(SendToTg(tgSenderWorker(testEvents, &wg, &mockSleeper, log.New(io.Discard, "", log.Ltime)), &mockSender, &wg, log.New(io.Discard, "", log.Ltime))) testEvents <- GenericMessage{TgChatId{123}, []MsgContent{StringMsgContent{"test1"}}} testEvents <- GenericMessage{TgChatId{124}, []MsgContent{StringMsgContent{"test3"}}} testEvents <- GenericMessage{TgChatId{123}, []MsgContent{StringMsgContent{"test2"}}} testEvents <- GenericMessage{TgChatId{124}, []MsgContent{StringMsgContent{"test4"}}} testEvents <- GenericMessage{FakeChatId{123}, []MsgContent{StringMsgContent{"testFake"}}} assert.Equal(t, 1, mockSleeper.callCount) *mockSleeper.ch <- nil assert.Equal(t, 1, mockSleeper.callCount) testEvents <- GenericMessage{TgChatId{123}, []MsgContent{StringMsgContent{"test5"}}} close(testEvents) wg.Wait() assert.Equal(t, 2, mockSleeper.callCount) assert.Len(t, mockSender.messages, 4) assert.Contains(t, mockSender.messages, GenericMessage{TgChatId{123}, []MsgContent{StringMsgContent{"test1"}, StringMsgContent{"test2"}}}) assert.Contains(t, mockSender.messages, GenericMessage{TgChatId{124}, []MsgContent{StringMsgContent{"test3"}, StringMsgContent{"test4"}}}) assert.Contains(t, mockSender.messages, GenericMessage{FakeChatId{123}, []MsgContent{StringMsgContent{"testFake"}}}) }