Enhance PostParser to handle service messages

This commit is contained in:
vvzvlad
2025-04-13 19:14:44 +03:00
parent 8516927ea2
commit defba4cdaf
3 changed files with 78 additions and 11 deletions
+13 -3
View File
@@ -134,9 +134,19 @@ class PostParser:
def _generate_title(self, message: Message) -> str:
"""Generate a title for a message, based on its content."""
if getattr(message, "channel_chat_created", False): #Channel created service message
return "✨ Channel created"
# Check for service messages first
if service := getattr(message, "service", None):
service_str = str(service)
if 'PINNED_MESSAGE' in service_str: return "📌 Pinned message"
elif 'NEW_CHAT_PHOTO' in service_str: return "🖼 New chat photo"
elif 'NEW_CHAT_TITLE' in service_str: return "✏️ New chat title"
elif 'VIDEO_CHAT_STARTED' in service_str: return "▶️ Video chat started"
elif 'VIDEO_CHAT_ENDED' in service_str: return "⏹ Video chat ended"
elif 'VIDEO_CHAT_SCHEDULED' in service_str: return "⏰ Video chat scheduled"
elif 'GROUP_CHAT_CREATED' in service_str: return "✨ Group chat created"
elif 'CHANNEL_CHAT_CREATED' in service_str: return "✨ Chat created"
elif 'DELETE_CHAT_PHOTO' in service_str: return "🗑️ Chat photo deleted"
elif 'NEW_CHAT_TITLE' in service_str: return "✏️ New chat title"
# Process media content
if message.media:
+11 -5
View File
@@ -83,11 +83,17 @@ async def _create_messages_groups(messages):
try:
# Skip service messages about pinned posts and new chat photos
if message.service:
if 'PINNED_MESSAGE' in str(message.service): continue
if 'NEW_CHAT_PHOTO' in str(message.service): continue
if 'NEW_CHAT_TITLE' in str(message.service): continue
if 'VIDEO_CHAT_STARTED' in str(message.service): continue
if 'PINNED_MESSAGE' in str(message.service): continue
if 'NEW_CHAT_PHOTO' in str(message.service): continue
if 'NEW_CHAT_TITLE' in str(message.service): continue
if 'VIDEO_CHAT_STARTED' in str(message.service): continue
if 'VIDEO_CHAT_ENDED' in str(message.service): continue
if 'VIDEO_CHAT_SCHEDULED' in str(message.service): continue
if 'GROUP_CHAT_CREATED' in str(message.service): continue
if 'CHANNEL_CHAT_CREATED' in str(message.service): continue
if 'DELETE_CHAT_PHOTO' in str(message.service): continue
if 'NEW_CHAT_TITLE' in str(message.service): continue
if message.media_group_id:
if message.media_group_id not in media_groups:
media_groups[message.media_group_id] = []
+54 -3
View File
@@ -23,7 +23,7 @@ class TestPostParserGenerateTitle(unittest.TestCase):
def tearDown(self):
pass
def _create_mock_message(self, media=None, text=None, caption=None, document_mime_type=None, channel_chat_created=False, web_page=None, poll=None):
def _create_mock_message(self, media=None, text=None, caption=None, document_mime_type=None, channel_chat_created=False, web_page=None, poll=None, service=None):
message = MagicMock(spec=Message)
message.media = media
message.text = text
@@ -31,6 +31,7 @@ class TestPostParserGenerateTitle(unittest.TestCase):
message.web_page = web_page
message.channel_chat_created = channel_chat_created
message.poll = poll
message.service = service
# Mock document properties if document exists
if document_mime_type:
@@ -79,8 +80,58 @@ class TestPostParserGenerateTitle(unittest.TestCase):
self.assertEqual(self.parser._generate_title(message), "📎 Document")
def test_generate_title_channel_created(self):
message = self._create_mock_message(channel_chat_created=True)
self.assertEqual(self.parser._generate_title(message), "✨ Channel created")
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.CHANNEL_CHAT_CREATED")
message = self._create_mock_message(service=mock_service, channel_chat_created=False) # channel_chat_created flag is now secondary
self.assertEqual(self.parser._generate_title(message), "✨ Chat created")
def test_generate_title_pinned_message(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.PINNED_MESSAGE")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "📌 Pinned message")
def test_generate_title_new_chat_photo(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.NEW_CHAT_PHOTO")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "🖼 New chat photo")
def test_generate_title_new_chat_title(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.NEW_CHAT_TITLE")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "✏️ New chat title")
def test_generate_title_video_chat_started(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.VIDEO_CHAT_STARTED")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "▶️ Video chat started")
def test_generate_title_video_chat_ended(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.VIDEO_CHAT_ENDED")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "⏹ Video chat ended")
def test_generate_title_video_chat_scheduled(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.VIDEO_CHAT_SCHEDULED")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "⏰ Video chat scheduled")
def test_generate_title_group_chat_created(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.GROUP_CHAT_CREATED")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "✨ Group chat created")
def test_generate_title_delete_chat_photo(self):
mock_service = MagicMock()
mock_service.__str__ = MagicMock(return_value="pyrogram.enums.MessageService.DELETE_CHAT_PHOTO")
message = self._create_mock_message(service=mock_service)
self.assertEqual(self.parser._generate_title(message), "🗑️ Chat photo deleted")
def test_generate_title_text_only(self):
message = self._create_mock_message(text="This is the first line.\nThis is the second line.")