diff --git a/tests/retrieval/services/test_embedding_service.py b/tests/retrieval/services/test_embedding_service.py new file mode 100644 index 0000000..7e908d7 --- /dev/null +++ b/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] diff --git a/tests/retrieval/services/test_retrieval_service.py b/tests/retrieval/services/test_retrieval_service.py new file mode 100644 index 0000000..3c535b4 --- /dev/null +++ b/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