From 3a1d25c80a7c7c56bb54e11adfacac983269164b Mon Sep 17 00:00:00 2001 From: Gregory Eremeev Date: Thu, 24 Oct 2024 15:53:25 -0300 Subject: [PATCH 1/2] Added caption extraction when send a msg with a file (#4499) Closes #4499 --- telethon/client/messages.py | 2 ++ tests/telethon/client/test_messages.py | 37 ++++++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/telethon/client/messages.py b/telethon/client/messages.py index 0c61a9080..4632e3936 100644 --- a/telethon/client/messages.py +++ b/telethon/client/messages.py @@ -826,6 +826,8 @@ async def callback(event): await client.send_message(chat, 'Hi, future!', schedule=timedelta(minutes=5)) """ if file is not None: + if isinstance(message, types.Message): + message = message.message return await self.send_file( entity, file, caption=message, reply_to=reply_to, attributes=attributes, parse_mode=parse_mode, diff --git a/tests/telethon/client/test_messages.py b/tests/telethon/client/test_messages.py index d47fcf459..c9b86afc8 100644 --- a/tests/telethon/client/test_messages.py +++ b/tests/telethon/client/test_messages.py @@ -1,8 +1,12 @@ import inspect +from unittest import mock +from unittest.mock import MagicMock import pytest from telethon import TelegramClient +from telethon.client import MessageMethods +from telethon.tl.types import PeerChat, MessageMediaDocument, Message @pytest.mark.asyncio @@ -38,3 +42,36 @@ async def send_file(self, entity, file, **kwargs): client = MockedClient() assert (await client.send_message('a', file='b', **arguments)) == sentinel + + +class TestMessageMethods: + @pytest.mark.asyncio + async def test_send_msg_and_file(self): + async def async_func(result): # AsyncMock was added only in 3.8 + return result + msg_methods = MessageMethods() + expected_result = Message( + id=0, peer_id=PeerChat(chat_id=0), message='', date=None, + ) + entity = 'test_entity' + message = Message( + id=1, peer_id=PeerChat(chat_id=0), message='expected_caption', date=None, + ) + media_file = MessageMediaDocument() + + with mock.patch.object( + target=MessageMethods, attribute='send_file', + new=MagicMock(return_value=async_func(expected_result)), create=True, + ) as mock_obj: + result = await msg_methods.send_message( + entity=entity, message=message, file=media_file, + ) + mock_obj.assert_called_once_with( + entity, media_file, caption=message.message, + reply_to=None, attributes=None, parse_mode=(), + force_document=False, thumb=None, buttons=None, + clear_draft=False, silent=None, schedule=None, + supports_streaming=False, formatting_entities=None, + comment_to=None, background=None, nosound_video=None, + ) + assert result == expected_result From 255dcbafcedae27400c1388ae69feaeb07df1c29 Mon Sep 17 00:00:00 2001 From: Gregory Eremeev Date: Fri, 25 Oct 2024 17:37:39 -0300 Subject: [PATCH 2/2] Fixes after review (#4499) --- telethon/client/messages.py | 1 + tests/telethon/client/test_messages.py | 16 +++++++++++----- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/telethon/client/messages.py b/telethon/client/messages.py index 4632e3936..5dcd2ee63 100644 --- a/telethon/client/messages.py +++ b/telethon/client/messages.py @@ -827,6 +827,7 @@ async def callback(event): """ if file is not None: if isinstance(message, types.Message): + formatting_entities = formatting_entities or message.entities message = message.message return await self.send_file( entity, file, caption=message, reply_to=reply_to, diff --git a/tests/telethon/client/test_messages.py b/tests/telethon/client/test_messages.py index c9b86afc8..2bdbb9e78 100644 --- a/tests/telethon/client/test_messages.py +++ b/tests/telethon/client/test_messages.py @@ -6,7 +6,7 @@ from telethon import TelegramClient from telethon.client import MessageMethods -from telethon.tl.types import PeerChat, MessageMediaDocument, Message +from telethon.tl.types import PeerChat, MessageMediaDocument, Message, MessageEntityBold @pytest.mark.asyncio @@ -46,7 +46,11 @@ async def send_file(self, entity, file, **kwargs): class TestMessageMethods: @pytest.mark.asyncio - async def test_send_msg_and_file(self): + @pytest.mark.parametrize( + 'formatting_entities', + ([MessageEntityBold(offset=0, length=0)], None) + ) + async def test_send_msg_and_file(self, formatting_entities): async def async_func(result): # AsyncMock was added only in 3.8 return result msg_methods = MessageMethods() @@ -56,6 +60,7 @@ async def async_func(result): # AsyncMock was added only in 3.8 entity = 'test_entity' message = Message( id=1, peer_id=PeerChat(chat_id=0), message='expected_caption', date=None, + entities=[MessageEntityBold(offset=9, length=9)], ) media_file = MessageMediaDocument() @@ -65,13 +70,14 @@ async def async_func(result): # AsyncMock was added only in 3.8 ) as mock_obj: result = await msg_methods.send_message( entity=entity, message=message, file=media_file, + formatting_entities=formatting_entities, ) mock_obj.assert_called_once_with( entity, media_file, caption=message.message, - reply_to=None, attributes=None, parse_mode=(), + formatting_entities=formatting_entities or message.entities, + reply_to=None, silent=None, attributes=None, parse_mode=(), force_document=False, thumb=None, buttons=None, - clear_draft=False, silent=None, schedule=None, - supports_streaming=False, formatting_entities=None, + clear_draft=False, schedule=None, supports_streaming=False, comment_to=None, background=None, nosound_video=None, ) assert result == expected_result