77 lines
3.7 KiB
Python
77 lines
3.7 KiB
Python
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()
|