Add unit tests for embedding_service and retrieval_service

pyproject.toml

@@ -26,6 +26,7 @@ dev = [
     "black",
     "psycopg2-binary",
     "pre-commit>=4.6.0",
+    "pytest-cov>=7.1.0",
 ]
 
 [tool.ruff]
@@ -47,3 +48,8 @@ asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "session"
 asyncio_default_test_loop_scope = "session"
 testpaths = ["tests"]
+addopts = "--cov=src --cov-report=term-missing"
+
+[tool.coverage.run]
+source = ["src"]
+omit = ["src/infrastructure/db/versions/*"]

tests/retrieval/services/test_embedding_service.py

@@ -0,0 +1,25 @@
+from unittest.mock import AsyncMock, patch
+
+from src.retrieval.services.embedding_service import embed_document, embed_query
+
+
+# for question
+async def test_embed_query_calls_llm_embed_with_is_query_true():
+    with patch("src.infrastructure.llm.embed", new_callable=AsyncMock) as mock_embed:
+        mock_embed.return_value = [0.1, 0.2, 0.3]
+
+        result = await embed_query("what electives should I take?")
+
+        mock_embed.assert_awaited_once_with("what electives should I take?", is_query=True)
+        assert result == [0.1, 0.2, 0.3]
+
+
+# for document
+async def test_embed_document_calls_llm_embed_with_is_query_is_false():
+    with patch("src.infrastructure.llm.embed", new_callable=AsyncMock) as mock_embed:
+        mock_embed.return_value = [0.4, 0.5, 0.6]
+
+        result = await embed_document("some course description text")
+
+        mock_embed.assert_awaited_once_with("some course description text", is_query=False)
+        assert result == [0.4, 0.5, 0.6]

tests/retrieval/services/test_retrieval_service.py

@@ -0,0 +1,36 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from src.config import settings
+from src.infrastructure.db.repository import Repository
+from src.retrieval.services.retrieval_service import get_relevant_chunks
+
+
+async def test_get_relevant_chunks_wires_embedding_and_repository_correctly():
+    fake_embedding = [0.1, 0.2, 0.3]
+    fake_chunks = ["sentinel_chunk_1", "sentinel_chunk_2"]
+
+    fake_session = MagicMock()
+    fake_context_manager = MagicMock()
+    fake_context_manager.__aenter__ = AsyncMock(return_value=fake_session)
+    fake_context_manager.__aexit__ = AsyncMock(return_value=None)
+    mock_session_factory = MagicMock(return_value=fake_context_manager)
+
+    with (
+        patch(
+            "src.retrieval.services.embedding_service.embed_query",
+            new_callable=AsyncMock,
+            return_value=fake_embedding,
+        ) as mock_embed_query,
+        patch.object(
+            Repository, "top_k_chunks", new_callable=AsyncMock, return_value=fake_chunks
+        ) as mock_top_k_chunks,
+        patch(
+            "src.retrieval.services.retrieval_service.async_session_factory",
+            mock_session_factory,
+        ),
+    ):
+        result = await get_relevant_chunks("what electives should I take?")
+
+        mock_embed_query.assert_awaited_once_with("what electives should I take?")
+        mock_top_k_chunks.assert_awaited_once_with(fake_session, fake_embedding, k=settings.top_k)
+        assert result == fake_chunks