hswro-alarm-bot/sender_worker_test.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"}}})
}