Oleksandr Kozachuk
213 lines
10 KiB
Python
213 lines
10 KiB
Python
"""Unit coverage for SPEC-004 input pipeline (IMG-10..16)."""
|
|
|
|
import base64
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
from fjerkroa_bot.ai_responder import AIMessage, AIResponse
|
|
from fjerkroa_bot.images import ImageCache, sniff_ext
|
|
from fjerkroa_bot.openai_responder import OpenAIResponder
|
|
from fjerkroa_bot.persistence import PersistentStore
|
|
|
|
from .test_bdd_envelope import FakeModelResponder
|
|
from .test_spec_ops import OpsBase
|
|
|
|
PNG = b"\x89PNG\r\n\x1a\n" + b"x" * 64
|
|
|
|
|
|
def make_cache(tmp, config=None):
|
|
store = PersistentStore(Path(tmp) / "bot.db")
|
|
cache = ImageCache(store, Path(tmp) / "images", lambda: config or {})
|
|
return store, cache
|
|
|
|
|
|
class TestIngest(unittest.TestCase):
|
|
def test_sniffed_types_only(self):
|
|
"""IMG-10: magic bytes decide; garbage and foreign types are rejected."""
|
|
self.assertEqual(sniff_ext(PNG), "png")
|
|
self.assertEqual(sniff_ext(b"\xff\xd8\xff\xe0rest"), "jpg")
|
|
self.assertIsNone(sniff_ext(b"MZ\x90\x00 definitely-an-exe"))
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store, cache = make_cache(tmp)
|
|
self.assertIsNone(cache.ingest_bytes(b"not an image", "chat", "alice", "1"))
|
|
sha = cache.ingest_bytes(PNG, "chat", "alice", "1")
|
|
self.assertIsNotNone(sha)
|
|
self.assertTrue((Path(tmp) / "images" / f"{sha}.png").exists())
|
|
self.assertEqual(store.images_recent("chat", 5)[0]["sha256"], sha)
|
|
|
|
def test_size_cap(self):
|
|
"""IMG-10: oversized uploads are dropped."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
_, cache = make_cache(tmp, {"image-max-bytes": 32})
|
|
self.assertIsNone(cache.ingest_bytes(PNG, "chat", "alice", "1"))
|
|
|
|
|
|
class TestVisionDataUrls(OpsBase):
|
|
async def test_attachment_becomes_data_url(self):
|
|
"""IMG-11: the model sees a data: URL, never the CDN link."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
_, cache = make_cache(tmp)
|
|
self.bot.airesponder.image_cache = cache
|
|
self.bot.respond = AsyncMock()
|
|
message = self.public_msg("look at this")
|
|
attachment = Mock()
|
|
attachment.url = "https://cdn.discordapp.com/attachments/1/2/cat.png?ex=deadbeef"
|
|
message.attachments = [attachment]
|
|
message.id = 42
|
|
with patch.object(ImageCache, "_download", new_callable=AsyncMock, return_value=PNG):
|
|
await self.bot.on_message(message)
|
|
sent_msg = self.bot.respond.await_args.args[0]
|
|
self.assertTrue(sent_msg.urls[0].startswith("data:image/png;base64,"))
|
|
self.assertNotIn("cdn.discordapp.com", sent_msg.urls[0])
|
|
|
|
|
|
class TestEviction(unittest.TestCase):
|
|
def test_lru_cap(self):
|
|
"""IMG-12: byte cap evicts oldest first, file + row together."""
|
|
big = b"\x89PNG\r\n\x1a\n" + b"a" * (700 * 1024)
|
|
big2 = b"\x89PNG\r\n\x1a\n" + b"b" * (700 * 1024)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store, cache = make_cache(tmp, {"image-cache-mb": 1})
|
|
first = cache.ingest_bytes(big, "chat", "alice", "1")
|
|
second = cache.ingest_bytes(big2, "chat", "alice", "2")
|
|
shas = [row["sha256"] for row in store.images_recent("chat", 5)]
|
|
self.assertNotIn(first, shas)
|
|
self.assertIn(second, shas)
|
|
self.assertFalse((Path(tmp) / "images" / f"{first}.png").exists())
|
|
|
|
def test_ttl(self):
|
|
"""IMG-12: entries past image-cache-ttl-days age out."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store, cache = make_cache(tmp, {"image-cache-ttl-days": 30})
|
|
sha = cache.ingest_bytes(PNG, "chat", "alice", "1")
|
|
with sqlite3.connect(store.db_path) as conn:
|
|
conn.execute("UPDATE images SET created_at = datetime('now', '-60 days') WHERE sha256 = ?", (sha,))
|
|
cache.evict()
|
|
self.assertEqual(store.images_recent("chat", 5), [])
|
|
self.assertFalse((Path(tmp) / "images" / f"{sha}.png").exists())
|
|
|
|
|
|
class TestEditPath(OpsBase):
|
|
async def prepare(self, with_images):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
_, cache = make_cache(self.tmp.name)
|
|
self.bot.airesponder.image_cache = cache
|
|
if with_images:
|
|
cache.ingest_bytes(PNG, "chat", "alice", "1")
|
|
self.bot.airesponder.edit_openai = AsyncMock(return_value=[__import__("io").BytesIO(PNG)])
|
|
self.bot.airesponder.draw = AsyncMock(return_value=[__import__("io").BytesIO(PNG)])
|
|
response = AIResponse("her", True, "chat", None, "als wikinger", True, False)
|
|
channel = MagicMock()
|
|
channel.name = "chat"
|
|
channel.send = AsyncMock()
|
|
channel.typing = MagicMock(return_value=AsyncMock(__aenter__=AsyncMock(), __aexit__=AsyncMock()))
|
|
await self.bot.send_answer_with_typing(response, channel, self.bot.airesponder, factual=True)
|
|
|
|
async def test_edit_uses_cached_sources(self):
|
|
"""IMG-13: picture_edit + cached images -> images.edit path."""
|
|
await self.prepare(with_images=True)
|
|
self.bot.airesponder.edit_openai.assert_awaited_once()
|
|
self.bot.airesponder.draw.assert_not_awaited()
|
|
|
|
async def test_empty_cache_falls_back_to_generate(self):
|
|
"""IMG-13: empty cache -> plain generation, the flag never fails a reply."""
|
|
await self.prepare(with_images=False)
|
|
self.bot.airesponder.edit_openai.assert_not_awaited()
|
|
self.bot.airesponder.draw.assert_awaited_once()
|
|
|
|
|
|
class TestPurges(OpsBase):
|
|
async def test_message_delete_and_forgetme_purge_images(self):
|
|
"""IMG-14: message deletion and !forgetme remove files + rows."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store, cache = make_cache(tmp)
|
|
self.bot.airesponder.image_cache = cache
|
|
cache.ingest_bytes(PNG, "chat", "alice", "99")
|
|
deleted = MagicMock()
|
|
deleted.id = 99
|
|
deleted.content = "pic"
|
|
deleted.author.name = "alice"
|
|
deleted.channel = MagicMock()
|
|
await self.bot.on_message_delete(deleted)
|
|
self.assertEqual(store.images_recent("chat", 5), [])
|
|
cache.ingest_bytes(b"\x89PNG\r\n\x1a\n" + b"z" * 32, "chat", "alice", "100")
|
|
message = self.public_msg("!forgetme")
|
|
message.author.name = "alice"
|
|
await self.bot.on_message(message)
|
|
self.assertEqual(store.images_recent("chat", 5), [])
|
|
|
|
|
|
class TestGeneratedImagesCached(OpsBase):
|
|
async def test_bot_output_joins_cache(self):
|
|
"""IMG-15: generated images are ingested as user 'assistant'."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store, cache = make_cache(tmp)
|
|
self.bot.airesponder.image_cache = cache
|
|
self.bot.airesponder.draw = AsyncMock(return_value=[__import__("io").BytesIO(PNG)])
|
|
response = AIResponse("her", True, "chat", None, "en katt", False, False)
|
|
channel = MagicMock()
|
|
channel.name = "chat"
|
|
channel.send = AsyncMock()
|
|
await self.bot.send_answer_with_typing(response, channel, self.bot.airesponder, factual=True)
|
|
rows = store.images_recent("chat", 5)
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0]["user"], "assistant")
|
|
|
|
|
|
class TestImageOnlyMessages(OpsBase):
|
|
async def test_image_only_post_cached_no_reply(self):
|
|
"""IMG-17: attachment without text -> cached + observed, no reply."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store, cache = make_cache(tmp)
|
|
self.bot.airesponder.image_cache = cache
|
|
self.bot.airesponder.observe_event = AsyncMock()
|
|
self.bot.respond = AsyncMock()
|
|
message = self.public_msg("")
|
|
message.content = ""
|
|
message.channel.name = "chat"
|
|
attachment = Mock()
|
|
attachment.url = "https://cdn.discordapp.com/attachments/1/2/silent.png"
|
|
message.attachments = [attachment]
|
|
message.id = 77
|
|
with patch.object(ImageCache, "_download", new_callable=AsyncMock, return_value=PNG):
|
|
await self.bot.on_message(message)
|
|
self.assertEqual(len(store.images_recent("chat", 5)), 1)
|
|
self.bot.airesponder.observe_event.assert_awaited_once()
|
|
self.bot.respond.assert_not_awaited()
|
|
|
|
|
|
class TestContextAnnouncesImages(unittest.IsolatedAsyncioTestCase):
|
|
def test_suffix_mentions_picture_edit(self):
|
|
"""IMG-16: cached channel images are announced in the context suffix."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = {"system": "s", "history-limit": 5, "history-directory": tmp}
|
|
responder = FakeModelResponder(config, "chat")
|
|
responder.image_cache.ingest_bytes(PNG, "chat", "alice", "1")
|
|
system = responder.message(AIMessage("alice", "hei", "chat"))[0]["content"]
|
|
self.assertIn("picture_edit", system)
|
|
self.assertIn("recent images in this channel: 1", system)
|
|
|
|
|
|
class TestEditOpenai(unittest.IsolatedAsyncioTestCase):
|
|
async def test_edit_call_shape_and_metering(self):
|
|
"""IMG-13: images.edit gets the file handles, n clamped, ledger counts."""
|
|
responder = OpenAIResponder({"openai-token": "t", "model": "m", "system": "s", "history-limit": 5}, "chat")
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
paths = []
|
|
for index in range(2):
|
|
path = Path(tmp) / f"in{index}.png"
|
|
path.write_bytes(PNG)
|
|
paths.append(path)
|
|
api_result = Mock(data=[Mock(b64_json=base64.b64encode(b"out").decode())])
|
|
with patch("fjerkroa_bot.openai_responder.openai_image_edit", new_callable=AsyncMock) as edit_mock:
|
|
edit_mock.return_value = api_result
|
|
buffers = await responder.edit_openai("wikinger", paths, 9)
|
|
self.assertEqual(buffers[0].read(), b"out")
|
|
self.assertEqual(edit_mock.await_args.kwargs["n"], 4)
|
|
self.assertEqual(len(edit_mock.await_args.kwargs["image"]), 2)
|
|
self.assertEqual(responder.ledger.images_today(), 1)
|