rpi: исправить приём картинок и пересланных сообщений в Astra

2026-10-11 14:46:48 +03:00
parent 7d0a891793
commit bba3858894
6 changed files with 560 additions and 0 deletions
@@ -0,0 +1,76 @@
+import os
+os.environ['TG_TOKEN']='test'
+import tempfile
+import threading
+import unittest
+from pathlib import Path
+from unittest.mock import Mock, patch, MagicMock
+import bot
+
+class Incoming(unittest.TestCase):
+    def setUp(self):
+        self.b=bot.Bot.__new__(bot.Bot); self.b.message_lock=threading.RLock()
+        self.b.active={}; self.b.queued={}
+        self.b.thread_for=Mock(return_value='tid')
+        self.b.codex=Mock()
+        self.b.codex.call.return_value={'turn':{'id':'turn'}}
+        self.message={'chat':{'id':1,'type':'private'},'from':{'id':1},'photo':[{'file_id':'photo'}]}
+        for name,value in [('ALLOWED',{1}),('send',Mock()),('tg',Mock()),('download_image',Mock(return_value='/tmp/test.jpg'))]:
+            p=patch.object(bot,name,value);p.start();self.addCleanup(p.stop)
+
+    def test_photo_without_caption(self):
+        self.b.handle_message(self.message)
+        inputs=self.b.codex.call.call_args.kwargs['input']
+        self.assertEqual(inputs[1],{'type':'localImage','path':'/tmp/test.jpg'})
+
+    def test_forwarded_photo_caption(self):
+        self.message.update(caption='Подпись',forward_origin={'type':'hidden_user','sender_user_name':'Автор'})
+        self.b.handle_message(self.message)
+        text=self.b.codex.call.call_args.kwargs['input'][0]['text']
+        self.assertIn('Подпись',text); self.assertIn('Автор',text)
+
+    def test_busy_queues(self):
+        self.b.active['tid']='turn'
+        self.b.handle_message(self.message)
+        self.assertEqual(self.b.queued['tid'],[self.message])
+        bot.download_image.assert_not_called()
+
+    def test_unauthorized_does_not_download(self):
+        self.message['from']['id']=2
+        self.b.handle_message(self.message)
+        bot.download_image.assert_not_called()
+
+    def test_caption_is_not_command(self):
+        self.message['caption']='/stop'
+        self.b.handle_message(self.message)
+        self.assertEqual(self.b.codex.call.call_args.args[0],'turn/start')
+
+    def test_download_error_is_reported(self):
+        bot.download_image.side_effect=ValueError('download failed')
+        self.b.handle_message(self.message)
+        self.b.codex.call.assert_not_called()
+        self.assertIn('download failed',bot.send.call_args.args[1])
+
+class Download(unittest.TestCase):
+    def test_download_biggest_photo(self):
+        with tempfile.TemporaryDirectory() as directory:
+            response=MagicMock(); response.__enter__.return_value.read.return_value=b'image'
+            with patch.object(bot,'INCOMING',Path(directory)),patch.object(bot,'tg',return_value={'file_path':'photos/image.jpg'}) as tg,patch.object(bot.urllib.request,'urlopen',return_value=response):
+                path=bot.download_image({'photo':[{'file_id':'small','width':1,'height':1},{'file_id':'big','width':10,'height':10}]})
+                self.assertEqual(tg.call_args.kwargs['file_id'],'big')
+                self.assertEqual(Path(path).read_bytes(),b'image')
+                self.assertEqual(Path(path).stat().st_mode & 0o777,0o600)
+
+    def test_download_error_does_not_leak_token(self):
+        with patch.object(bot,'tg',return_value={'file_path':'photo.jpg'}),patch.object(bot.urllib.request,'urlopen',side_effect=RuntimeError('SECRET URL')):
+            with self.assertRaises(ValueError) as caught:
+                bot.download_image({'photo':[{'file_id':'file'}]})
+            self.assertNotIn('SECRET',str(caught.exception))
+
+    def test_large_file_rejected(self):
+        with patch.object(bot,'tg') as tg:
+            with self.assertRaises(ValueError):
+                bot.download_image({'photo':[{'file_id':'file','file_size':bot.MAX_DOWNLOAD+1}]})
+            tg.assert_not_called()
+
+if __name__=='__main__': unittest.main()