| 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() |