| 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=='})) |
|
|
