Merge branch 'main' of github.com:yihong1120/Construction-Hazard-Dete…

…ction

src/telegram_notifier.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import os
 from io import BytesIO
 
@@ -15,7 +16,7 @@ def __init__(self, bot_token: str | None = None):
         self.bot_token = bot_token or os.getenv('TELEGRAM_BOT_TOKEN')
         self.bot = Bot(token=self.bot_token)
 
-    def send_notification(
+    async def send_notification(
         self, chat_id: str,
         message: str,
         image: np.ndarray | None = None,
@@ -25,21 +26,28 @@ def send_notification(
             buffer = BytesIO()
             image_pil.save(buffer, format='PNG')
             buffer.seek(0)
-            self.bot.send_photo(chat_id=chat_id, photo=buffer, caption=message)
+            response = await self.bot.send_photo(
+                chat_id=chat_id,
+                photo=buffer,
+                caption=message,
+            )
         else:
-            self.bot.send_message(chat_id=chat_id, text=message)
-        return 'Message sent'
+            response = await self.bot.send_message(
+                chat_id=chat_id,
+                text=message,
+            )
+        return response
 
 
 # Example usage
-def main():
+async def main():
     notifier = TelegramNotifier()
     chat_id = 'your_chat_id_here'
     message = 'Hello, Telegram!'
     image = np.zeros((100, 100, 3), dtype=np.uint8)
-    response = notifier.send_notification(chat_id, message, image=image)
+    response = await notifier.send_notification(chat_id, message, image=image)
     print(response)
 
 
 if __name__ == '__main__':
-    main()
+    asyncio.run(main())

tests/telegram_notifier_test.py

@@ -2,14 +2,15 @@
 
 import unittest
 from io import BytesIO
+from unittest.mock import AsyncMock
 from unittest.mock import patch
 
 import numpy as np
 
 from src.telegram_notifier import TelegramNotifier
 
 
-class TestTelegramNotifier(unittest.TestCase):
+class TestTelegramNotifier(unittest.IsolatedAsyncioTestCase):
 
     @classmethod
     def setUpClass(cls):
@@ -19,26 +20,29 @@ def test_init(self):
         """Test if the TelegramNotifier instance is initialised correctly."""
         self.assertEqual(self.telegram_notifier.bot_token, 'test_bot_token')
 
-    @patch('telegram.Bot.send_message')
-    def test_send_notification_no_image(self, mock_send_message):
+    @patch('telegram.Bot.send_message', new_callable=AsyncMock)
+    async def test_send_notification_no_image(self, mock_send_message):
         """Test sending a notification without an image."""
         mock_send_message.return_value = 'Message sent'
         chat_id = 'test_chat_id'
         message = 'Hello, Telegram!'
-        response = self.telegram_notifier.send_notification(chat_id, message)
+        response = await self.telegram_notifier.send_notification(
+            chat_id,
+            message,
+        )
         self.assertEqual(response, 'Message sent')
         mock_send_message.assert_called_once_with(
             chat_id=chat_id, text=message,
         )
 
-    @patch('telegram.Bot.send_photo')
-    def test_send_notification_with_image(self, mock_send_photo):
+    @patch('telegram.Bot.send_photo', new_callable=AsyncMock)
+    async def test_send_notification_with_image(self, mock_send_photo):
         """Test sending a notification with an image."""
         mock_send_photo.return_value = 'Message sent'
         chat_id = 'test_chat_id'
         message = 'Hello, Telegram!'
         image = np.zeros((100, 100, 3), dtype=np.uint8)
-        response = self.telegram_notifier.send_notification(
+        response = await self.telegram_notifier.send_notification(
             chat_id, message, image=image,
         )
         self.assertEqual(response, 'Message sent')