File size: 4,873 Bytes
ee2d10d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import asyncio
import base64
import json
from pathlib import Path
from tempfile import TemporaryDirectory
import unittest
from unittest import mock
from fastapi import FastAPI
from fastapi.testclient import TestClient
import api.image_tasks as image_tasks_module
from services.image_task_service import ImageTaskService, _collect_image_urls, _inline_image
class InlineTaskDeliveryTests(unittest.TestCase):
def make_service(self, root):
service = ImageTaskService(root / 'tasks.json', retention_days_getter=lambda: 0)
service._tasks['owner:result'] = {
'id': 'result', 'owner_id': 'owner', 'status': 'running',
'base_url': 'https://site.test',
}
return service
def finish(self, service, payload=b'\x89PNG\r\n\x1a\n' + b'x' * 200000):
encoded = base64.b64encode(payload).decode('ascii')
with mock.patch('services.image_task_service.proxied_image_url', side_effect=lambda url: url):
service._update_task('owner:result', status='success', data=[
{'url': 'data:image/png;base64,' + encoded, 'revised_prompt': 'test'},
])
public = service.list_tasks({'id': 'owner'}, ['result'])['items'][0]
return payload, public, service._tasks['owner:result']['inline_result_token']
def test_poll_is_small_and_complete_image_survives_database_reload(self):
with TemporaryDirectory() as directory:
service = self.make_service(Path(directory))
payload, public, token = self.finish(service)
self.assertLess(len(json.dumps(public)), 1500)
self.assertNotIn('b64_json', public['data'][0])
self.assertEqual(public['data'][0]['url'], f'https://site.test/image-task-results/{token}/0')
self.assertEqual(service.get_inline_result(token, 0), (payload, 'image/png'))
restored = ImageTaskService(service.path, retention_days_getter=lambda: 0)
self.assertEqual(restored.get_inline_result(token, 0), (payload, 'image/png'))
def test_foreign_users_cannot_discover_capability_and_invalid_requests_fail(self):
with TemporaryDirectory() as directory:
service = self.make_service(Path(directory))
_, _, token = self.finish(service)
self.assertEqual(service.list_tasks({'id': 'foreign'}, ['result'])['items'], [])
self.assertIsNone(service.get_inline_result('x' * 43, 0))
self.assertIsNone(service.get_inline_result(token, -1))
self.assertIsNone(service.get_inline_result(token, 2))
def test_capability_rotates_when_image_is_replaced(self):
with TemporaryDirectory() as directory:
service = self.make_service(Path(directory))
_, _, old = self.finish(service, b'old')
payload, _, new = self.finish(service, b'new')
self.assertNotEqual(old, new)
self.assertIsNone(service.get_inline_result(old, 0))
self.assertEqual(service.get_inline_result(new, 0)[0], payload)
def test_image_endpoint_does_not_use_network_mount_or_event_loop(self):
with TemporaryDirectory() as directory:
service = self.make_service(Path(directory))
payload, _, token = self.finish(service)
original = service.get_inline_result
def off_loop(*args):
with self.assertRaises(RuntimeError):
asyncio.get_running_loop()
return original(*args)
with mock.patch.object(image_tasks_module, 'image_task_service', service), \
mock.patch.object(service, 'get_inline_result', side_effect=off_loop):
app = FastAPI()
app.include_router(image_tasks_module.create_router())
with TestClient(app) as client:
response = client.get(f'/image-task-results/{token}/0')
self.assertEqual(response.status_code, 200)
self.assertEqual(response.content, payload)
self.assertEqual(response.headers['content-type'], 'image/png')
self.assertEqual(response.headers['x-content-type-options'], 'nosniff')
missing = client.get('/image-task-results/' + 'x' * 43 + '/0')
self.assertEqual(missing.status_code, 404)
self.assertEqual(missing.headers['cache-control'], 'no-store')
def test_inline_image_is_not_copied_to_usage_log_urls(self):
self.assertEqual(_collect_image_urls([{'url': 'data:image/png;base64,eA=='}]), [])
self.assertEqual(_collect_image_urls([{'url': 'https://site/image.png'}]), ['https://site/image.png'])
def test_active_svg_content_is_not_served_from_public_image_endpoint(self):
self.assertIsNone(_inline_image({'url': 'data:image/svg+xml;base64,eA=='}))
|