Update test mocks to be json serializable

Author: MathyouMB

tests/completions/services/test_completion_service.py

@@ -3,13 +3,19 @@
 from completions import CompletionService
 
 
+class FakeChunk:
+    def __init__(self, id, content):
+        self.id = id
+        self.content = content
+
+
 @pytest.fixture
 def fake_repository():
     repo = MagicMock()
     repo.get_top_k_chunks_by_similarity.return_value = [
-        MagicMock(content="Chunk 1"),
-        MagicMock(content="Chunk 2"),
-        MagicMock(content="Chunk 3"),
+        FakeChunk(id=1, content="Chunk 1"),
+        FakeChunk(id=2, content="Chunk 2"),
+        FakeChunk(id=3, content="Chunk 3"),
     ]
     return repo
 
@@ -30,7 +36,7 @@ def service(fake_repository):
 def test_create(mock_create_prompt, mock_create_completion, service, capsys):
     """Should call the LLM with the correct prompt and print the answer."""
 
-    service.create("What is AI?", k=3)
+    result = service.create("What is AI?", k=3)
 
     service.repository.get_top_k_chunks_by_similarity.assert_called_once_with(
         "What is AI?", 3
@@ -40,5 +46,4 @@ def test_create(mock_create_prompt, mock_create_completion, service, capsys):
     )
     mock_create_completion.assert_called_once_with("Generated Prompt")
 
-    captured = capsys.readouterr()
-    assert "LLM Answer" in captured.out
+    assert result == "LLM Answer"

tests/completions/services/test_context_generation_service.py

@@ -1,17 +1,36 @@
 import pytest
-from unittest.mock import MagicMock
+from dataclasses import dataclass
 from completions.services.context_generation_service import ContextGenerationService
 
 
+@dataclass
+class FakeChunk:
+    id: int
+    content: str
+
+
 @pytest.fixture
 def fake_repository():
-    repo = MagicMock()
-    repo.get_top_k_chunks_by_similarity.return_value = [
-        MagicMock(content="Chunk 1"),
-        MagicMock(content="Chunk 2"),
-        MagicMock(content="Chunk 3"),
-    ]
-    return repo
+    class FakeRepository:
+        def __init__(self):
+            self._return_value = [
+                FakeChunk(id=1, content="Chunk 1"),
+                FakeChunk(id=2, content="Chunk 2"),
+                FakeChunk(id=3, content="Chunk 3"),
+            ]
+            self.called_with = None
+
+        def get_top_k_chunks_by_similarity(self, query, k):
+            self.called_with = (query, k)
+            return self._return_value
+
+        def set_return_value(self, chunks):
+            self._return_value = chunks
+
+        def assert_called_once_with(self, query, k):
+            assert self.called_with == (query, k)
+
+    return FakeRepository()
 
 
 def test_generate_context_returns_joined_content(fake_repository):
@@ -20,35 +39,31 @@ def test_generate_context_returns_joined_content(fake_repository):
     )
     context = service.process(k=3)
 
-    fake_repository.get_top_k_chunks_by_similarity.assert_called_once_with(
-        "example query", 3
-    )
+    fake_repository.assert_called_once_with("example query", 3)
     assert context == "Chunk 1\n\nChunk 2\n\nChunk 3"
 
 
 def test_generate_context_with_different_k(fake_repository):
-    # Change fake_repository to return fewer chunks
-    fake_repository.get_top_k_chunks_by_similarity.return_value = [
-        MagicMock(content="Only Chunk 1"),
-    ]
+    fake_repository.set_return_value(
+        [
+            FakeChunk(id=1, content="Only Chunk 1"),
+        ]
+    )
+
     service = ContextGenerationService(
         query="different query", repository=fake_repository
     )
     context = service.process(k=1)
 
-    fake_repository.get_top_k_chunks_by_similarity.assert_called_once_with(
-        "different query", 1
-    )
+    fake_repository.assert_called_once_with("different query", 1)
     assert context == "Only Chunk 1"
 
 
 def test_generate_context_when_no_chunks(fake_repository):
-    # Make it return an empty list
-    fake_repository.get_top_k_chunks_by_similarity.return_value = []
+    fake_repository.set_return_value([])
+
     service = ContextGenerationService(query="empty query", repository=fake_repository)
     context = service.process(k=3)
 
-    fake_repository.get_top_k_chunks_by_similarity.assert_called_once_with(
-        "empty query", 3
-    )
+    fake_repository.assert_called_once_with("empty query", 3)
     assert context == ""