Oleksandr Kozachuk
88 lines
4.4 KiB
Python
88 lines
4.4 KiB
Python
"""Unit coverage for SPEC-004 image generation (IMG-01..05)."""
|
|
|
|
import base64
|
|
import unittest
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
from discord import TextChannel
|
|
|
|
from fjerkroa_bot.ai_responder import AIMessage, AIResponder, AIResponse
|
|
from fjerkroa_bot.openai_responder import OpenAIResponder
|
|
|
|
from .test_bdd_envelope import FakeModelResponder, envelope
|
|
from .test_spec_ops import OpsBase
|
|
|
|
RESPONDER_CONFIG = {"openai-token": "t", "model": "m", "system": "s", "history-limit": 5}
|
|
|
|
|
|
def image_api_result(count):
|
|
return Mock(data=[Mock(b64_json=base64.b64encode(f"png{i}".encode()).decode()) for i in range(count)])
|
|
|
|
|
|
class TestBase64Generation(unittest.IsolatedAsyncioTestCase):
|
|
async def test_draw_returns_decoded_buffers_and_meters(self):
|
|
"""IMG-01: count images decoded from b64_json, each metered in the ledger."""
|
|
responder = OpenAIResponder(RESPONDER_CONFIG, "chat")
|
|
with patch("fjerkroa_bot.openai_responder.openai_image", new_callable=AsyncMock) as image_mock:
|
|
image_mock.return_value = image_api_result(2)
|
|
buffers = await responder.draw_openai("en katt på brygga", 2)
|
|
self.assertEqual([buf.read() for buf in buffers], [b"png0", b"png1"])
|
|
self.assertEqual(responder.ledger.images_today(), 2)
|
|
self.assertEqual(image_mock.await_args.kwargs["n"], 2)
|
|
self.assertEqual(image_mock.await_args.kwargs["model"], "gpt-image-2")
|
|
self.assertNotIn("response_format", image_mock.await_args.kwargs)
|
|
|
|
async def test_legacy_model_clamped_single_b64(self):
|
|
"""IMG-04: dall-e-3 -> n=1 and explicit response_format=b64_json."""
|
|
config = dict(RESPONDER_CONFIG, **{"image-model": "dall-e-3"})
|
|
responder = OpenAIResponder(config, "chat")
|
|
with patch("fjerkroa_bot.openai_responder.openai_image", new_callable=AsyncMock) as image_mock:
|
|
image_mock.return_value = image_api_result(1)
|
|
buffers = await responder.draw_openai("a cat", 3)
|
|
self.assertEqual(len(buffers), 1)
|
|
self.assertEqual(image_mock.await_args.kwargs["n"], 1)
|
|
self.assertEqual(image_mock.await_args.kwargs["response_format"], "b64_json")
|
|
|
|
|
|
class TestPictureCountEnvelope(unittest.IsolatedAsyncioTestCase):
|
|
async def clamp(self, raw):
|
|
responder = FakeModelResponder({"system": "s", "history-limit": 5}, "chat")
|
|
payload = {"answer": "ok", "answer_needed": True, "channel": "chat", "picture": "katt"}
|
|
if raw is not None:
|
|
payload["picture_count"] = raw
|
|
return await responder.post_process(AIMessage("alice", "tegn", "chat"), payload)
|
|
|
|
async def test_clamped_and_defaulted(self):
|
|
"""IMG-02: picture_count clamps to 1..4, defaults to 1 when absent."""
|
|
self.assertEqual((await self.clamp(3)).picture_count, 3)
|
|
self.assertEqual((await self.clamp(9)).picture_count, 4)
|
|
self.assertEqual((await self.clamp(0)).picture_count, 1)
|
|
self.assertEqual((await self.clamp(None)).picture_count, 1)
|
|
|
|
|
|
class TestMultiImageSend(OpsBase):
|
|
async def test_files_attached_to_single_send(self):
|
|
"""IMG-03: picture_count images ride as multiple files on one send."""
|
|
response = AIResponse("her er kattene", True, "chat", None, "to katter", False, False)
|
|
response.picture_count = 2
|
|
import io
|
|
|
|
self.bot.airesponder.draw = AsyncMock(return_value=[io.BytesIO(b"a"), io.BytesIO(b"b")])
|
|
channel = MagicMock(spec=TextChannel)
|
|
channel.send = AsyncMock()
|
|
await self.bot.send_answer_with_typing(response, channel, self.bot.airesponder, factual=True)
|
|
self.bot.airesponder.draw.assert_awaited_once_with("to katter", 2)
|
|
files = channel.send.await_args.kwargs["files"]
|
|
self.assertEqual(len(files), 2)
|
|
|
|
|
|
class TestNoTranslateStep(unittest.IsolatedAsyncioTestCase):
|
|
async def test_translate_is_gone_prompt_untouched(self):
|
|
"""IMG-05: no translate() anywhere; the picture prompt survives verbatim."""
|
|
responder = FakeModelResponder({"system": "s", "history-limit": 5}, "chat")
|
|
self.assertFalse(hasattr(responder, "translate"))
|
|
self.assertFalse(hasattr(AIResponder, "translate"))
|
|
responder.scripted.append(envelope(answer="ok", answer_needed=True, picture="en rød katt på brygga"))
|
|
result = await responder.send(AIMessage("alice", "tegn en katt", "chat"))
|
|
self.assertEqual(result.picture, "en rød katt på brygga")
|