22 lines
1.3 KiB
Python
22 lines
1.3 KiB
Python
import io
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from image_upload import read_image_upload
|
|
from image_test import decode_references
|
|
|
|
class UploadTests(unittest.TestCase):
|
|
def test_multiple_images_keep_order_and_reject_oversized_body(self):
|
|
boundary='synthetic-boundary'
|
|
parts=[]
|
|
for name,value,mime in [('model',b'athena-image','text/plain'),('prompt',b'change color','text/plain'),('image[]',b'\x89PNG\r\n\x1a\nfirst','image/png'),('image[]',b'\xff\xd8\xffsecond','image/jpeg')]:
|
|
parts.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\nContent-Type: {mime}\r\n\r\n'.encode()+value+b'\r\n')
|
|
body=b''.join(parts)+f'--{boundary}--\r\n'.encode()
|
|
handler=SimpleNamespace(connection=SimpleNamespace(settimeout=lambda _:None),headers={'Content-Type':'multipart/form-data; boundary='+boundary,'Content-Length':str(len(body))},rfile=io.BytesIO(body))
|
|
fields,images=read_image_upload(handler)
|
|
self.assertEqual(fields,{'model':'athena-image','prompt':'change color'})
|
|
self.assertEqual([raw for raw,_ in decode_references(images,4)],[b'\x89PNG\r\n\x1a\nfirst',b'\xff\xd8\xffsecond'])
|
|
handler.headers['Content-Length']=str(50*1024*1024)
|
|
with self.assertRaises(ValueError):read_image_upload(handler)
|
|
|
|
if __name__=='__main__':unittest.main()
|