| from __future__ import annotations |
|
|
| import unittest |
| from types import SimpleNamespace |
| from unittest import mock |
|
|
| from services.protocol import conversation |
|
|
|
|
| class ImageParallelGenerationTests(unittest.TestCase): |
| def test_multi_image_request_uses_independent_pool_jobs_when_enabled(self) -> None: |
| calls: list[tuple[int, int, int | None]] = [] |
|
|
| def fake_serial(request, *, index_offset=0, total_count=None, **_kwargs): |
| calls.append((request.n, index_offset, total_count)) |
| yield conversation.ImageOutput( |
| kind="result", |
| model=request.model, |
| index=index_offset + 1, |
| total=int(total_count or request.n), |
| data=[{"url": f"https://example.test/{index_offset + 1}.png"}], |
| ) |
|
|
| with ( |
| mock.patch.object(conversation, "config", SimpleNamespace(image_parallel_generation=True)), |
| mock.patch.object(conversation, "_stream_image_outputs_with_pool", side_effect=fake_serial), |
| ): |
| outputs = list(conversation.stream_image_outputs_with_pool( |
| conversation.ConversationRequest(model="gpt-image-2", prompt="cat", n=3) |
| )) |
|
|
| self.assertEqual(sorted(calls), [(1, 0, 3), (1, 1, 3), (1, 2, 3)]) |
| self.assertEqual(sorted(output.index for output in outputs), [1, 2, 3]) |
| self.assertTrue(all(output.total == 3 for output in outputs)) |
|
|
| def test_multi_image_request_stays_serial_when_disabled(self) -> None: |
| calls: list[int] = [] |
|
|
| def fake_serial(request, **_kwargs): |
| calls.append(request.n) |
| for index in range(1, request.n + 1): |
| yield conversation.ImageOutput( |
| kind="result", |
| model=request.model, |
| index=index, |
| total=request.n, |
| data=[{"url": f"https://example.test/{index}.png"}], |
| ) |
|
|
| with ( |
| mock.patch.object(conversation, "config", SimpleNamespace(image_parallel_generation=False)), |
| mock.patch.object(conversation, "_stream_image_outputs_with_pool", side_effect=fake_serial), |
| ): |
| outputs = list(conversation.stream_image_outputs_with_pool( |
| conversation.ConversationRequest(model="gpt-image-2", prompt="cat", n=2) |
| )) |
|
|
| self.assertEqual(calls, [2]) |
| self.assertEqual([output.index for output in outputs], [1, 2]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
