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
}

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{123, "test1"}
	testEvents <- GenericMessage{124, "test3"}
	testEvents <- GenericMessage{123, "test2"}
	testEvents <- GenericMessage{124, "test4"}
	assert.Equal(t, 1, mockSleeper.callCount)
	*mockSleeper.ch <- nil
	assert.Equal(t, 1, mockSleeper.callCount)
	testEvents <- GenericMessage{123, "test5"}
	close(testEvents)
	wg.Wait()
	assert.Equal(t, 2, mockSleeper.callCount)

	assert.Len(t, mockSender.messages, 3)
	assert.Contains(t, mockSender.messages, GenericMessage{123, "test1\ntest2\n"})
}