import base64 import json import unittest from types import SimpleNamespace from unittest.mock import Mock, patch from telebot import apihelper from app.agent.tools.impl.send_message import SendMessageInput from app.chain.message import MessageChain from app.core.config import settings from app.modules.slack import SlackModule from app.modules.telegram.telegram import Telegram from app.modules.telegram import TelegramModule from app.schemas import CommingMessage from app.schemas.types import MessageChannel class AgentImageSupportTest(unittest.TestCase): def test_telegram_extract_images_returns_prefixed_file_ids(self): images = TelegramModule._extract_images( { "photo": [{"file_id": "small"}, {"file_id": "large"}], "document": {"file_id": "doc-image", "mime_type": "image/png"}, } ) self.assertEqual( images, ["tg://file_id/large", "tg://file_id/doc-image"], ) def test_telegram_message_parser_accepts_double_encoded_body(self): module = TelegramModule() body = json.dumps( json.dumps( { "message": { "from": {"id": 10001, "username": "tester"}, "chat": {"id": 10001, "type": "private"}, "photo": [{"file_id": "small"}, {"file_id": "large"}], } } ) ) with patch.object( module, "get_config", return_value=SimpleNamespace(name="telegram-test", config={}), ), patch.object( module, "get_instance", return_value=SimpleNamespace(bot_username=None), ): message = module.message_parser( source="telegram-test", body=body, form={}, args={} ) self.assertIsNotNone(message) self.assertEqual(message.images, ["tg://file_id/large"]) def test_telegram_forward_payload_uses_dict_not_json_string(self): payload = Telegram._serialize_update_payload( SimpleNamespace( to_dict=lambda: { "text": "hi", "photo": [{"file_id": "image-1"}], } ) ) self.assertEqual( payload, {"text": "hi", "photo": [{"file_id": "image-1"}]}, ) def test_telegram_download_file_uses_configured_file_url(self): telegram = Telegram.__new__(Telegram) telegram._bot = Mock() telegram._telegram_token = "token-123" telegram._bot.get_file.return_value = SimpleNamespace(file_path="photos/a.jpg") old_file_url = apihelper.FILE_URL old_proxy = apihelper.proxy apihelper.FILE_URL = "https://tg-proxy.example/file/bot{0}/{1}" apihelper.proxy = {"https": "http://127.0.0.1:7890"} try: with patch( "app.modules.telegram.telegram.RequestUtils.get_res", return_value=SimpleNamespace(content=b"image-bytes"), ) as get_res: content = telegram.download_file("file-id-1") finally: apihelper.FILE_URL = old_file_url apihelper.proxy = old_proxy self.assertEqual(content, b"image-bytes") get_res.assert_called_once_with( "https://tg-proxy.example/file/bottoken-123/photos/a.jpg" ) def test_process_allows_image_only_message(self): chain = MessageChain() message = CommingMessage( channel=MessageChannel.Telegram, source="telegram-test", userid="10001", username="tester", images=["tg://file_id/image-1"], ) with patch.object(chain, "message_parser", return_value=message), patch.object( chain, "handle_message" ) as handle_message: chain.process(body="{}", form={}, args={"source": "telegram-test"}) handle_kwargs = handle_message.call_args.kwargs self.assertEqual(handle_kwargs["text"], "") self.assertEqual(handle_kwargs["images"], ["tg://file_id/image-1"]) def test_image_message_routes_to_agent_even_when_global_agent_is_disabled(self): chain = MessageChain() with patch.object(chain, "load_cache", return_value={}), patch.object( chain.messagehelper, "put" ), patch.object(chain.messageoper, "add"), patch.object( chain, "_handle_ai_message" ) as handle_ai_message, patch.object( settings, "AI_AGENT_ENABLE", True ), patch.object( settings, "AI_AGENT_GLOBAL", False ): chain.handle_message( channel=MessageChannel.Telegram, source="telegram-test", userid="10001", username="tester", text="", images=["tg://file_id/image-1"], ) handle_ai_message.assert_called_once() def test_slack_images_use_authenticated_data_url_download(self): chain = MessageChain() with patch.object( chain, "run_module", return_value="data:image/png;base64,abc123", ) as run_module: images = chain._download_images_to_base64( images=["https://files.slack.com/files-pri/T1-F1/test.png"], channel=MessageChannel.Slack, source="slack-test", ) self.assertEqual(images, ["data:image/png;base64,abc123"]) run_module.assert_called_once_with( "download_file_to_data_url", file_url="https://files.slack.com/files-pri/T1-F1/test.png", source="slack-test", ) def test_slack_module_download_file_to_data_url(self): module = SlackModule() client = Mock() client.download_file.return_value = (b"png-binary", "image/png") with patch.object( module, "get_config", return_value=SimpleNamespace(name="slack-test") ), patch.object(module, "get_instance", return_value=client): data_url = module.download_file_to_data_url( "https://files.slack.com/files-pri/T1-F1/test.png", "slack-test", ) self.assertEqual( data_url, f"data:image/png;base64,{base64.b64encode(b'png-binary').decode()}", ) def test_send_message_input_accepts_image_only_payload(self): payload = SendMessageInput( explanation="send poster image", image_url="https://example.com/poster.png", ) self.assertEqual(payload.image_url, "https://example.com/poster.png") if __name__ == "__main__": unittest.main()