diff --git a/src/line_notifier.py b/src/line_notifier.py index e5e6437..217b46c 100644 --- a/src/line_notifier.py +++ b/src/line_notifier.py @@ -2,18 +2,17 @@ import os from io import BytesIO -from typing import TypedDict, Optional +from typing import TypedDict import numpy as np import requests from dotenv import load_dotenv from PIL import Image -from typing import TypedDict class NotificationData(TypedDict): message: str - image: Optional[np.ndarray] + image: np.ndarray | None class LineNotifier: @@ -37,27 +36,31 @@ def __init__(self, line_token: str | None = None): 'LINE_NOTIFY_TOKEN not provided or in environment variables.', ) - def send_notification(self, data: NotificationData) -> int: + def send_notification( + self, + message: str, + image: np.ndarray | None = None, + ) -> int: """ Sends a notification via LINE Notify, optionally including an image. Args: - data (NotificationData): The notification data including message - and optional image. + message (str): The message to send. + label (Optional[str]): The label of the image_name. + image (Optional[np.ndarray]): The image to send with the message. + Defaults to None. Returns: response.status_code (int): The status code of the response. """ headers = {'Authorization': f"Bearer {self.line_token}"} - payload = {'message': data['message']} + payload = {'message': message} files = {} - if data['image'] is not None: - if isinstance(data['image'], bytes): + if image is not None: + if isinstance(image, bytes): # Convert bytes to NumPy array - image = np.array(Image.open(BytesIO(data['image']))) - else: - image = data['image'] + image = np.array(Image.open(BytesIO(image))) image_pil = Image.fromarray(image) buffer = BytesIO() image_pil.save(buffer, format='PNG') @@ -85,11 +88,7 @@ def main(): message = 'Hello, LINE Notify!' # Create a dummy image for testing image = np.zeros((100, 100, 3), dtype=np.uint8) - data: NotificationData = { - 'message': message, - 'image': image, - } - response_code = notifier.send_notification(data) + response_code = notifier.send_notification(message, image=image) print(f"Response code: {response_code}") diff --git a/tests/line_notifier_test.py b/tests/line_notifier_test.py new file mode 100644 index 0000000..928c4b2 --- /dev/null +++ b/tests/line_notifier_test.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import unittest +from io import BytesIO +from unittest.mock import MagicMock +from unittest.mock import patch + +import numpy as np +from PIL import Image + +from src.line_notifier import LineNotifier + + +class TestLineNotifier(unittest.TestCase): + def setUp(self): + self.line_token = 'test_token' + self.message = 'Test Message' + self.image = np.zeros((100, 100, 3), dtype=np.uint8) + self.notifier = LineNotifier(line_token=self.line_token) + + @patch.dict('os.environ', {'LINE_NOTIFY_TOKEN': 'test_env_token'}) + def test_init_with_env_token(self): + notifier = LineNotifier() + self.assertEqual(notifier.line_token, 'test_env_token') + + def test_init_with_provided_token(self): + notifier = LineNotifier(line_token='provided_token') + self.assertEqual(notifier.line_token, 'provided_token') + + def test_init_without_token(self): + with self.assertRaises(ValueError): + LineNotifier(line_token=None) + + @patch('src.line_notifier.requests.post') + def test_send_notification_without_image(self, mock_post): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_post.return_value = mock_response + + status_code = self.notifier.send_notification(self.message) + self.assertEqual(status_code, 200) + mock_post.assert_called_once_with( + 'https://notify-api.line.me/api/notify', + headers={'Authorization': f'Bearer {self.line_token}'}, + params={'message': self.message}, + ) + + @patch('src.line_notifier.requests.post') + def test_send_notification_with_image(self, mock_post): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_post.return_value = mock_response + + status_code = self.notifier.send_notification(self.message, self.image) + self.assertEqual(status_code, 200) + mock_post.assert_called_once() + args, kwargs = mock_post.call_args + self.assertIn('files', kwargs) + self.assertIn('imageFile', kwargs['files']) + + # Check if the image is correctly converted and sent + image_file = kwargs['files']['imageFile'] + self.assertEqual(image_file[0], 'image.png') + self.assertEqual(image_file[2], 'image/png') + image = Image.open(image_file[1]) + self.assertTrue(np.array_equal(np.array(image), self.image)) + + @patch('src.line_notifier.requests.post') + def test_send_notification_with_bytes_image(self, mock_post): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_post.return_value = mock_response + + buffer = BytesIO() + Image.fromarray(self.image).save(buffer, format='PNG') + buffer.seek(0) + image_bytes = buffer.read() + + status_code = self.notifier.send_notification( + self.message, image_bytes, + ) + self.assertEqual(status_code, 200) + mock_post.assert_called_once() + args, kwargs = mock_post.call_args + self.assertIn('files', kwargs) + self.assertIn('imageFile', kwargs['files']) + + # Check if the image is correctly converted and sent + image_file = kwargs['files']['imageFile'] + self.assertEqual(image_file[0], 'image.png') + self.assertEqual(image_file[2], 'image/png') + image = Image.open(image_file[1]) + self.assertTrue(np.array_equal(np.array(image), self.image)) + + +if __name__ == '__main__': + unittest.main()