66 lines
2.0 KiB
Go
66 lines
2.0 KiB
Go
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"}}})
|
|
}
|