slidegent / tests /test_gemma_client.py
jomasego's picture
Upload folder using huggingface_hub
ad0ebea verified
Raw
History Blame Contribute Delete
1.56 kB
import unittest
from unittest.mock import patch
from src.llm.gemma_client import GemmaClient, generate_slides, load_model, safe_parse_generated_json
class TestGemmaClient(unittest.TestCase):
def test_model_loading(self):
model = load_model(backend="transformers")
self.assertIsNotNone(model, "Model client should initialize successfully.")
@patch.object(GemmaClient, "generate_json")
def test_generate_slides_with_mock(self, mock_generate_json):
mock_generate_json.return_value = (
'[{"title": "AI Overview", "bullets": ["Why AI matters"], '
'"layout_type": "title_slide"}]'
)
result = generate_slides("Create a presentation about AI.")
self.assertIsInstance(result, list, "Output should be a list of slide objects.")
self.assertGreater(len(result), 0, "There should be at least one slide generated.")
for slide in result:
self.assertIn('title', slide, "Each slide should have a title.")
self.assertIn('bullets', slide, "Each slide should have bullets.")
self.assertIn('layout_type', slide, "Each slide should include layout_type.")
def test_safe_parse_generated_json(self):
raw = """```json
[{"title": "Intro", "bullets": ["Point"], "layout_type": "content_slide"}]
```"""
parsed = safe_parse_generated_json(raw)
self.assertEqual(parsed[0]["title"], "Intro")
self.assertEqual(parsed[0]["layout_type"], "content_slide")
if __name__ == '__main__':
unittest.main()