Initial upload

yihong1120 · Jul 16, 2024 · 978d516 · 978d516
1 parent 0643eaa
commit 978d516
Show file tree

Hide file tree

Showing 6 changed files with 153 additions and 74 deletions.
diff --git a/src/messenger_notifier.py b/src/messenger_notifier.py
@@ -1,21 +1,36 @@
+from __future__ import annotations
+
 import os
+from io import BytesIO
+
+import numpy as np
 import requests
 from dotenv import load_dotenv
-import numpy as np
 from PIL import Image
-from io import BytesIO
 
 
 class MessengerNotifier:
     def __init__(self, page_access_token: str | None = None):
         load_dotenv()
-        self.page_access_token = page_access_token or os.getenv('FACEBOOK_PAGE_ACCESS_TOKEN')
+        self.page_access_token = page_access_token or os.getenv(
+            'FACEBOOK_PAGE_ACCESS_TOKEN',
+        )
         if not self.page_access_token:
-            raise ValueError('FACEBOOK_PAGE_ACCESS_TOKEN not provided or in environment variables.')
+            raise ValueError(
+                'FACEBOOK_PAGE_ACCESS_TOKEN missing.',
+            )
 
-    def send_notification(self, recipient_id: str, message: str, image: np.ndarray | None = None) -> int:
+    def send_notification(
+        self,
+        recipient_id: str,
+        message: str,
+        image: np.ndarray | None = None,
+    ) -> int:
         headers = {'Authorization': f"Bearer {self.page_access_token}"}
-        payload = {'message': {'text': message}, 'recipient': {'id': recipient_id}}
+        payload = {
+            'message': {'text': message},
+            'recipient': {'id': recipient_id},
+        }
 
         if image is not None:
             image_pil = Image.fromarray(image)
@@ -27,13 +42,16 @@ def send_notification(self, recipient_id: str, message: str, image: np.ndarray |
                 f"https://graph.facebook.com/v11.0/me/messages?access_token={self.page_access_token}",
                 headers=headers,
                 files=files,
-                data={'recipient': f'{{"id":"{recipient_id}"}}', 'message': '{"attachment":{"type":"image","payload":{}}}'}
+                data={
+                    'recipient': f'{{"id":"{recipient_id}"}}',
+                    'message': '{"attachment":{"type":"image","payload":{}}}',
+                },
             )
         else:
             response = requests.post(
                 f"https://graph.facebook.com/v11.0/me/messages?access_token={self.page_access_token}",
                 headers=headers,
-                json=payload
+                json=payload,
             )
         return response.status_code
 
@@ -44,7 +62,9 @@ def main():
     recipient_id = 'your_recipient_id_here'
     message = 'Hello, Messenger!'
     image = np.zeros((100, 100, 3), dtype=np.uint8)
-    response_code = notifier.send_notification(recipient_id, message, image=image)
+    response_code = notifier.send_notification(
+        recipient_id, message, image=image,
+    )
     print(f"Response code: {response_code}")
 
 

diff --git a/src/telegram_notifier.py b/src/telegram_notifier.py
@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 import os
-from dotenv import load_dotenv
+from io import BytesIO
+
 import numpy as np
+from dotenv import load_dotenv
 from PIL import Image
-from io import BytesIO
 from telegram import Bot
 
 
@@ -12,7 +15,11 @@ 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(self, chat_id: str, message: str, image: np.ndarray | None = None) -> str:
+    def send_notification(
+        self, chat_id: str,
+        message: str,
+        image: np.ndarray | None = None,
+    ) -> str:
         if image is not None:
             image_pil = Image.fromarray(image)
             buffer = BytesIO()
@@ -21,7 +28,7 @@ def send_notification(self, chat_id: str, message: str, image: np.ndarray | None
             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"
+        return 'Message sent'
 
 
 # Example usage

diff --git a/src/wechat_notifier.py b/src/wechat_notifier.py
@@ -1,13 +1,20 @@
+from __future__ import annotations
+
 import os
+from io import BytesIO
+
+import numpy as np
 import requests
 from dotenv import load_dotenv
-import numpy as np
 from PIL import Image
-from io import BytesIO
 
 
 class WeChatNotifier:
-    def __init__(self, corp_id: str | None = None, corp_secret: str | None = None, agent_id: int | None = None):
+    def __init__(
+        self, corp_id: str | None = None,
+        corp_secret: str | None = None,
+        agent_id: int | None = None,
+    ):
         load_dotenv()
         self.corp_id = corp_id or os.getenv('WECHAT_CORP_ID')
         self.corp_secret = corp_secret or os.getenv('WECHAT_CORP_SECRET')
@@ -19,28 +26,33 @@ def get_access_token(self) -> str:
         response = requests.get(url)
         return response.json().get('access_token')
 
-    def send_notification(self, user_id: str, message: str, image: np.ndarray | None = None) -> str:
+    def send_notification(
+        self,
+        user_id: str,
+        message: str,
+        image: np.ndarray | None = None,
+    ) -> str:
         url = f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={self.access_token}"
         payload = {
-            "touser": user_id,
-            "msgtype": "text",
-            "agentid": self.agent_id,
-            "text": {
-                "content": message
+            'touser': user_id,
+            'msgtype': 'text',
+            'agentid': self.agent_id,
+            'text': {
+                'content': message,
             },
-            "safe": 0
+            'safe': 0,
         }
 
         if image is not None:
             media_id = self.upload_media(image)
             payload = {
-                "touser": user_id,
-                "msgtype": "image",
-                "agentid": self.agent_id,
-                "image": {
-                    "media_id": media_id
+                'touser': user_id,
+                'msgtype': 'image',
+                'agentid': self.agent_id,
+                'image': {
+                    'media_id': media_id,
                 },
-                "safe": 0
+                'safe': 0,
             }
 
         response = requests.post(url, json=payload)

diff --git a/tests/messenger_notifier_test.py b/tests/messenger_notifier_test.py
@@ -1,32 +1,47 @@
-import os
+from __future__ import annotations
+
 import unittest
-import numpy as np
 from unittest.mock import patch
-from io import BytesIO
-from PIL import Image
+
+import numpy as np
+
 from src.messenger_notifier import MessengerNotifier
 
+
 class TestMessengerNotifier(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        cls.messenger_notifier = MessengerNotifier(page_access_token='test_page_access_token')
+        cls.messenger_notifier = MessengerNotifier(
+            page_access_token='test_page_access_token',
+        )
 
     def test_init(self):
         """Test if the MessengerNotifier instance is initialised correctly."""
-        self.assertEqual(self.messenger_notifier.page_access_token, 'test_page_access_token')
+        self.assertEqual(
+            self.messenger_notifier.page_access_token,
+            'test_page_access_token',
+        )
 
     @patch('requests.post')
     def test_send_notification_no_image(self, mock_post):
         """Test sending a notification without an image."""
         mock_post.return_value.status_code = 200
         recipient_id = 'test_recipient_id'
         message = 'Hello, Messenger!'
-        response_code = self.messenger_notifier.send_notification(recipient_id, message)
+        response_code = self.messenger_notifier.send_notification(
+            recipient_id, message,
+        )
         self.assertEqual(response_code, 200)
         self.assertTrue(mock_post.called)
         args, kwargs = mock_post.call_args
-        self.assertEqual(kwargs['json'], {'message': {'text': message}, 'recipient': {'id': recipient_id}})
+        self.assertEqual(
+            kwargs['json'], {
+                'message': {
+                    'text': message,
+                }, 'recipient': {'id': recipient_id},
+            },
+        )
 
     @patch('requests.post')
     def test_send_notification_with_image(self, mock_post):
@@ -35,7 +50,9 @@ def test_send_notification_with_image(self, mock_post):
         recipient_id = 'test_recipient_id'
         message = 'Hello, Messenger!'
         image = np.zeros((100, 100, 3), dtype=np.uint8)
-        response_code = self.messenger_notifier.send_notification(recipient_id, message, image=image)
+        response_code = self.messenger_notifier.send_notification(
+            recipient_id, message, image=image,
+        )
         self.assertEqual(response_code, 200)
         self.assertTrue(mock_post.called)
         args, kwargs = mock_post.call_args
@@ -46,8 +63,12 @@ def test_send_notification_with_image(self, mock_post):
 
     def test_missing_page_access_token(self):
         """Test initialisation without a page access token."""
-        with self.assertRaises(ValueError, msg='FACEBOOK_PAGE_ACCESS_TOKEN not provided or in environment variables.'):
+        with self.assertRaises(
+            ValueError,
+            msg='FACEBOOK_PAGE_ACCESS_TOKEN missing.',
+        ):
             MessengerNotifier(page_access_token=None)
 
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/telegram_notifier_test.py b/tests/telegram_notifier_test.py
@@ -1,10 +1,13 @@
-import os
+from __future__ import annotations
+
 import unittest
+from io import BytesIO
 from unittest.mock import patch
+
 import numpy as np
-from io import BytesIO
+
 from src.telegram_notifier import TelegramNotifier
-from PIL import Image
+
 
 class TestTelegramNotifier(unittest.TestCase):
 
@@ -19,27 +22,32 @@ def test_init(self):
     @patch('telegram.Bot.send_message')
     def test_send_notification_no_image(self, mock_send_message):
         """Test sending a notification without an image."""
-        mock_send_message.return_value = "Message sent"
+        mock_send_message.return_value = 'Message sent'
         chat_id = 'test_chat_id'
         message = 'Hello, Telegram!'
         response = 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)
+        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):
         """Test sending a notification with an image."""
-        mock_send_photo.return_value = "Message sent"
+        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(chat_id, message, image=image)
-        self.assertEqual(response, "Message sent")
+        response = self.telegram_notifier.send_notification(
+            chat_id, message, image=image,
+        )
+        self.assertEqual(response, 'Message sent')
         mock_send_photo.assert_called_once()
         args, kwargs = mock_send_photo.call_args
         self.assertEqual(kwargs['chat_id'], chat_id)
         self.assertEqual(kwargs['caption'], message)
         self.assertIsInstance(kwargs['photo'], BytesIO)
 
+
 if __name__ == '__main__':
     unittest.main()