From d42be83709e59c50cf6189939978926d95310e2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Rudowicz?= <michal.rudowicz@fl9.eu>
Date: Mon, 13 Jan 2025 20:37:50 +0100
Subject: [PATCH] Preparation for command tests

---
 command_handler.go |  4 +++-
 main.go            |  5 +++--
 telegram_utils.go  | 18 ++++++++++++++----
 test_utils.go      | 15 +++++++++++++--
 4 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/command_handler.go b/command_handler.go
index 08a9264..fe9d898 100644
--- a/command_handler.go
+++ b/command_handler.go
@@ -13,7 +13,7 @@ func getCurrentStatus(dataStore *DataStore, logger *log.Logger, config AppConfig
 	return msg.Format(tpl, s, logger)
 }
 
-func HandleTelegramCommands(bot tgbotapi.BotAPI, logger *log.Logger, dataStore *DataStore, config AppConfig, tpl *template.Template, s SatelNameGetter) {
+func HandleTelegramCommands(bot TelegramBot, logger *log.Logger, dataStore *DataStore, config AppConfig, tpl *template.Template, s SatelNameGetter) {
 	u := tgbotapi.NewUpdate(0)
 	u.Timeout = 60
 
@@ -42,4 +42,6 @@ func HandleTelegramCommands(bot tgbotapi.BotAPI, logger *log.Logger, dataStore *
 			logger.Print(err)
 		}
 	}
+
+	logger.Println("Stopping command handler")
 }
diff --git a/main.go b/main.go
index a7f255d..0ac8d7f 100644
--- a/main.go
+++ b/main.go
@@ -61,7 +61,7 @@ func main() {
 	s := makeSatel(config.SatelAddr, config.PoolInterval.GetDuration())
 	logger.Printf("Connected to Satel: %s", config.SatelAddr)
 
-	var bot TelegramBotSender = nil
+	var bot TelegramBot = nil
 	if len(config.TelegramApiKey) != 0 {
 		b, err := tgbotapi.NewBotAPI(config.TelegramApiKey)
 		if err != nil {
@@ -70,7 +70,7 @@ func main() {
 		logger.Print("Created Telegram Bot API client")
 		bot = b
 	} else {
-		bot = EmptySender{}
+		bot = EmptyBot{}
 	}
 
 	tgSender := TgSender{bot, s, log.New(os.Stderr, "TgFormatter", log.Lmicroseconds), config.ChatIds}
@@ -114,6 +114,7 @@ func main() {
 	}
 
 	logger.Print("Closing...")
+	bot.StopReceivingUpdates()
 	close(closeDebugTools)
 	close(tgEvents)
 	wg.Wait()
diff --git a/telegram_utils.go b/telegram_utils.go
index 32832e7..b8e0735 100644
--- a/telegram_utils.go
+++ b/telegram_utils.go
@@ -7,18 +7,28 @@ import (
 	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
 )
 
-type TelegramBotSender interface {
+type TelegramBot interface {
 	Send(c tgbotapi.Chattable) (tgbotapi.Message, error)
+	GetUpdatesChan(u tgbotapi.UpdateConfig) tgbotapi.UpdatesChannel
+	StopReceivingUpdates()
 }
 
-type EmptySender struct{}
+type EmptyBot struct{}
 
-func (self EmptySender) Send(_ tgbotapi.Chattable) (tgbotapi.Message, error) {
+func (_ EmptyBot) Send(_ tgbotapi.Chattable) (tgbotapi.Message, error) {
 	return tgbotapi.Message{}, nil
 }
 
+func (_ EmptyBot) GetUpdatesChan(u tgbotapi.UpdateConfig) tgbotapi.UpdatesChannel {
+	ch := make(chan tgbotapi.Update)
+	close(ch)
+	return ch
+}
+
+func (_ EmptyBot) StopReceivingUpdates() {}
+
 type TgSender struct {
-	bot     TelegramBotSender
+	bot     TelegramBot
 	s       SatelNameGetter
 	logger  *log.Logger
 	chatIds []int64
diff --git a/test_utils.go b/test_utils.go
index 0dddc7c..9cfb89b 100644
--- a/test_utils.go
+++ b/test_utils.go
@@ -17,13 +17,24 @@ func makeGenericMessage(changeType satel.ChangeType, index int, val bool) Generi
 
 type MockTgBotAPI struct {
 	messages []tgbotapi.Chattable
+	updates  chan tgbotapi.Update
 }
 
-func (self *MockTgBotAPI) Send(c tgbotapi.Chattable) (tgbotapi.Message, error) {
-	self.messages = append(self.messages, c)
+func (botMock *MockTgBotAPI) Send(c tgbotapi.Chattable) (tgbotapi.Message, error) {
+	botMock.messages = append(botMock.messages, c)
 	return tgbotapi.Message{}, nil
 }
 
+func (botMock *MockTgBotAPI) GetUpdatesChan(u tgbotapi.UpdateConfig) tgbotapi.UpdatesChannel {
+	if botMock.updates != nil {
+		panic("Called mock GetUpdatesChan more than once")
+	}
+	botMock.updates = make(chan tgbotapi.Update)
+	return botMock.updates
+}
+
+func (botMock *MockTgBotAPI) StopReceivingUpdates() { close(botMock.updates) }
+
 type MockSatelNameGetter struct {
 	name string
 }