File size: 2,534 Bytes
85270f9 | 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 | 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()
|