Oleksandr Kozachuk
1a5da0ae7c
- Add extensive tests for igdblib.py (0% -> 100% coverage expected) - Add tests for leonardo_draw.py AI image generation - Add tests for openai_responder.py with GPT integration - Add tests for discord_bot.py bot functionality - Add extended tests for ai_responder.py edge cases - Fix critical bugs in igdblib.py: * Fix platforms() method treating name as string instead of list * Fix game_info() method missing endpoint parameter * Add safe dictionary access with .get() methods Coverage improvements target areas with lowest coverage to maximize impact. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
247 lines
11 KiB
Python
247 lines
11 KiB
Python
import asyncio
|
|
import unittest
|
|
from io import BytesIO
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import aiohttp
|
|
|
|
from fjerkroa_bot.leonardo_draw import LeonardoAIDrawMixIn
|
|
|
|
|
|
class MockLeonardoDrawer(LeonardoAIDrawMixIn):
|
|
"""Mock class to test the mixin."""
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
|
|
class TestLeonardoAIDrawMixIn(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
self.config = {"leonardo-token": "test_token"}
|
|
self.drawer = MockLeonardoDrawer(self.config)
|
|
|
|
async def test_draw_leonardo_success(self):
|
|
"""Test successful image generation with Leonardo AI."""
|
|
# Mock image data
|
|
fake_image_data = b"fake_image_data"
|
|
|
|
# Mock responses
|
|
generation_response = {
|
|
"sdGenerationJob": {"generationId": "test_generation_id"}
|
|
}
|
|
|
|
status_response = {
|
|
"generations_by_pk": {
|
|
"generated_images": [{"url": "http://example.com/image.jpg"}]
|
|
}
|
|
}
|
|
|
|
with patch("aiohttp.ClientSession") as mock_session_class:
|
|
# Create mock session
|
|
mock_session = AsyncMock()
|
|
mock_session_class.return_value.__aenter__.return_value = mock_session
|
|
mock_session_class.return_value.__aexit__.return_value = None
|
|
|
|
# Mock POST request (generation)
|
|
mock_post_response = AsyncMock()
|
|
mock_post_response.json.return_value = generation_response
|
|
mock_session.post.return_value.__aenter__.return_value = mock_post_response
|
|
mock_session.post.return_value.__aexit__.return_value = None
|
|
|
|
# Mock GET requests (status check and image download)
|
|
mock_get_response1 = AsyncMock()
|
|
mock_get_response1.json.return_value = status_response
|
|
|
|
mock_get_response2 = AsyncMock()
|
|
mock_get_response2.read.return_value = fake_image_data
|
|
|
|
mock_session.get.side_effect = [
|
|
mock_session.get.return_value, # Status check
|
|
mock_session.get.return_value # Image download
|
|
]
|
|
mock_session.get.return_value.__aenter__.side_effect = [
|
|
mock_get_response1, # Status check
|
|
mock_get_response2 # Image download
|
|
]
|
|
mock_session.get.return_value.__aexit__.return_value = None
|
|
|
|
# Mock DELETE request
|
|
mock_delete_response = AsyncMock()
|
|
mock_delete_response.json.return_value = {}
|
|
mock_session.delete.return_value.__aenter__.return_value = mock_delete_response
|
|
mock_session.delete.return_value.__aexit__.return_value = None
|
|
|
|
result = await self.drawer.draw_leonardo("A beautiful landscape")
|
|
|
|
# Verify the result
|
|
self.assertIsInstance(result, BytesIO)
|
|
self.assertEqual(result.read(), fake_image_data)
|
|
|
|
async def test_draw_leonardo_no_generation_job(self):
|
|
"""Test when generation job is not returned."""
|
|
generation_response = {} # No sdGenerationJob
|
|
|
|
with patch("aiohttp.ClientSession") as mock_session_class:
|
|
mock_session = AsyncMock()
|
|
mock_session_class.return_value.__aenter__.return_value = mock_session
|
|
mock_session_class.return_value.__aexit__.return_value = None
|
|
|
|
mock_post_response = AsyncMock()
|
|
mock_post_response.json.return_value = generation_response
|
|
mock_session.post.return_value.__aenter__.return_value = mock_post_response
|
|
mock_session.post.return_value.__aexit__.return_value = None
|
|
|
|
with patch("asyncio.sleep") as mock_sleep:
|
|
with patch("fjerkroa_bot.leonardo_draw.exponential_backoff") as mock_backoff:
|
|
mock_backoff.return_value = iter([1, 2, 4]) # Limited attempts
|
|
|
|
with self.assertRaises(StopIteration):
|
|
await self.drawer.draw_leonardo("test description")
|
|
|
|
async def test_draw_leonardo_no_generations_by_pk(self):
|
|
"""Test when generations_by_pk is not in response."""
|
|
generation_response = {"sdGenerationJob": {"generationId": "test_id"}}
|
|
status_response = {} # No generations_by_pk
|
|
|
|
with patch("aiohttp.ClientSession") as mock_session_class:
|
|
mock_session = AsyncMock()
|
|
mock_session_class.return_value.__aenter__.return_value = mock_session
|
|
mock_session_class.return_value.__aexit__.return_value = None
|
|
|
|
# Mock POST (successful)
|
|
mock_post_response = AsyncMock()
|
|
mock_post_response.json.return_value = generation_response
|
|
mock_session.post.return_value.__aenter__.return_value = mock_post_response
|
|
mock_session.post.return_value.__aexit__.return_value = None
|
|
|
|
# Mock GET (status check - no generations)
|
|
mock_get_response = AsyncMock()
|
|
mock_get_response.json.return_value = status_response
|
|
mock_session.get.return_value.__aenter__.return_value = mock_get_response
|
|
mock_session.get.return_value.__aexit__.return_value = None
|
|
|
|
with patch("asyncio.sleep") as mock_sleep:
|
|
with patch("fjerkroa_bot.leonardo_draw.exponential_backoff") as mock_backoff:
|
|
mock_backoff.return_value = iter([1, 2]) # Limited attempts
|
|
|
|
with self.assertRaises(StopIteration):
|
|
await self.drawer.draw_leonardo("test description")
|
|
|
|
async def test_draw_leonardo_no_generated_images(self):
|
|
"""Test when no generated images are available yet."""
|
|
generation_response = {"sdGenerationJob": {"generationId": "test_id"}}
|
|
status_response = {"generations_by_pk": {"generated_images": []}}
|
|
|
|
with patch("aiohttp.ClientSession") as mock_session_class:
|
|
mock_session = AsyncMock()
|
|
mock_session_class.return_value.__aenter__.return_value = mock_session
|
|
mock_session_class.return_value.__aexit__.return_value = None
|
|
|
|
# Mock POST (successful)
|
|
mock_post_response = AsyncMock()
|
|
mock_post_response.json.return_value = generation_response
|
|
mock_session.post.return_value.__aenter__.return_value = mock_post_response
|
|
mock_session.post.return_value.__aexit__.return_value = None
|
|
|
|
# Mock GET (status check - empty images)
|
|
mock_get_response = AsyncMock()
|
|
mock_get_response.json.return_value = status_response
|
|
mock_session.get.return_value.__aenter__.return_value = mock_get_response
|
|
mock_session.get.return_value.__aexit__.return_value = None
|
|
|
|
with patch("asyncio.sleep") as mock_sleep:
|
|
with patch("fjerkroa_bot.leonardo_draw.exponential_backoff") as mock_backoff:
|
|
mock_backoff.return_value = iter([1, 2]) # Limited attempts
|
|
|
|
with self.assertRaises(StopIteration):
|
|
await self.drawer.draw_leonardo("test description")
|
|
|
|
async def test_draw_leonardo_exception_handling(self):
|
|
"""Test exception handling during image generation."""
|
|
with patch("aiohttp.ClientSession") as mock_session_class:
|
|
mock_session = AsyncMock()
|
|
mock_session_class.return_value.__aenter__.return_value = mock_session
|
|
mock_session_class.return_value.__aexit__.return_value = None
|
|
|
|
# Make POST request raise an exception
|
|
mock_session.post.side_effect = Exception("Network error")
|
|
|
|
with patch("asyncio.sleep") as mock_sleep:
|
|
with patch("fjerkroa_bot.leonardo_draw.exponential_backoff") as mock_backoff:
|
|
mock_backoff.return_value = iter([1, 2]) # Limited attempts
|
|
|
|
with self.assertRaises(StopIteration):
|
|
await self.drawer.draw_leonardo("test description")
|
|
|
|
async def test_draw_leonardo_request_parameters(self):
|
|
"""Test that correct parameters are sent to Leonardo API."""
|
|
fake_image_data = b"fake_image_data"
|
|
generation_response = {"sdGenerationJob": {"generationId": "test_id"}}
|
|
status_response = {
|
|
"generations_by_pk": {
|
|
"generated_images": [{"url": "http://example.com/image.jpg"}]
|
|
}
|
|
}
|
|
|
|
with patch("aiohttp.ClientSession") as mock_session_class:
|
|
mock_session = AsyncMock()
|
|
mock_session_class.return_value.__aenter__.return_value = mock_session
|
|
mock_session_class.return_value.__aexit__.return_value = None
|
|
|
|
# Mock all responses
|
|
mock_post_response = AsyncMock()
|
|
mock_post_response.json.return_value = generation_response
|
|
mock_session.post.return_value.__aenter__.return_value = mock_post_response
|
|
mock_session.post.return_value.__aexit__.return_value = None
|
|
|
|
mock_get_response1 = AsyncMock()
|
|
mock_get_response1.json.return_value = status_response
|
|
mock_get_response2 = AsyncMock()
|
|
mock_get_response2.read.return_value = fake_image_data
|
|
|
|
mock_session.get.side_effect = [
|
|
mock_session.get.return_value,
|
|
mock_session.get.return_value
|
|
]
|
|
mock_session.get.return_value.__aenter__.side_effect = [
|
|
mock_get_response1,
|
|
mock_get_response2
|
|
]
|
|
mock_session.get.return_value.__aexit__.return_value = None
|
|
|
|
mock_delete_response = AsyncMock()
|
|
mock_delete_response.json.return_value = {}
|
|
mock_session.delete.return_value.__aenter__.return_value = mock_delete_response
|
|
mock_session.delete.return_value.__aexit__.return_value = None
|
|
|
|
description = "A beautiful sunset"
|
|
await self.drawer.draw_leonardo(description)
|
|
|
|
# Verify POST request parameters
|
|
mock_session.post.assert_called_once_with(
|
|
"https://cloud.leonardo.ai/api/rest/v1/generations",
|
|
json={
|
|
"prompt": description,
|
|
"modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3",
|
|
"num_images": 1,
|
|
"sd_version": "v2",
|
|
"promptMagic": True,
|
|
"unzoomAmount": 1,
|
|
"width": 512,
|
|
"height": 512,
|
|
},
|
|
headers={
|
|
"Authorization": f"Bearer {self.config['leonardo-token']}",
|
|
"Accept": "application/json",
|
|
"Content-Type": "application/json",
|
|
},
|
|
)
|
|
|
|
# Verify DELETE request was called
|
|
mock_session.delete.assert_called_once_with(
|
|
"https://cloud.leonardo.ai/api/rest/v1/generations/test_id",
|
|
headers={"Authorization": f"Bearer {self.config['leonardo-token']}"},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |