tests(cli): Mock embedding function in tests

Author: Zhe Yu

src/vectorcode/common.py

@@ -134,7 +134,7 @@ def get_embedding_function(configs: Config) -> chromadb.EmbeddingFunction:
         ef = getattr(embedding_functions, configs.embedding_function)(
             **configs.embedding_params
         )
-        if ef is None:
+        if ef is None:  # pragma: nocover
             raise AttributeError()
         return ef
     except AttributeError:

tests/subcommands/query/test_query.py

@@ -64,8 +64,14 @@ def mock_config():
 
 @pytest.mark.asyncio
 async def test_get_query_result_files(mock_collection, mock_config):
-    # Mock the reranker
-    with patch("vectorcode.subcommands.query.get_reranker") as mock_get_reranker:
+    mock_embedding_function = MagicMock()
+    with (
+        patch("vectorcode.subcommands.query.get_reranker") as mock_get_reranker,
+        patch(
+            "vectorcode.subcommands.query.get_embedding_function",
+            return_value=mock_embedding_function,
+        ),
+    ):
         mock_reranker_instance = MagicMock()
         mock_reranker_instance.rerank = AsyncMock(
             return_value=[
@@ -82,9 +88,7 @@ async def test_get_query_result_files(mock_collection, mock_config):
         # Check that query was called with the right parameters
         mock_collection.query.assert_called_once()
         args, kwargs = mock_collection.query.call_args
-        assert kwargs["query_texts"] == [
-            "test query"
-        ]  # Assuming chunking produces this
+        mock_embedding_function.assert_called_once_with(["test query"])
         assert kwargs["n_results"] == 6  # n_result(3) * query_multiplier(2)
         assert IncludeEnum.metadatas in kwargs["include"]
         assert IncludeEnum.distances in kwargs["include"]
@@ -285,10 +289,14 @@ async def test_get_query_result_files_chunking(mock_collection, mock_config):
     mock_config.query = [
         "this is a longer query that should be chunked into multiple parts"
     ]
-
+    mock_embedding_function = MagicMock()
     with (
         patch("vectorcode.subcommands.query.StringChunker") as MockChunker,
         patch("vectorcode.subcommands.query.reranker.NaiveReranker") as MockReranker,
+        patch(
+            "vectorcode.subcommands.query.get_embedding_function",
+            return_value=mock_embedding_function,
+        ),
     ):
         # Set up MockChunker to chunk the query
         mock_chunker_instance = MagicMock()
@@ -309,7 +317,7 @@ async def test_get_query_result_files_chunking(mock_collection, mock_config):
         # Check query was called with chunked query
         mock_collection.query.assert_called_once()
         _, kwargs = mock_collection.query.call_args
-        assert kwargs["query_texts"] == ["chunk1", "chunk2", "chunk3"]
+        mock_embedding_function.assert_called_once_with(["chunk1", "chunk2", "chunk3"])
 
         # Check the result
         assert result == ["file1.py", "file2.py"]
@@ -319,10 +327,14 @@ async def test_get_query_result_files_chunking(mock_collection, mock_config):
 async def test_get_query_result_files_multiple_queries(mock_collection, mock_config):
     # Set multiple query terms
     mock_config.query = ["term1", "term2", "term3"]
-
+    mock_embedding_function = MagicMock()
     with (
         patch("vectorcode.subcommands.query.StringChunker") as MockChunker,
         patch("vectorcode.subcommands.query.reranker.NaiveReranker") as MockReranker,
+        patch(
+            "vectorcode.subcommands.query.get_embedding_function",
+            return_value=mock_embedding_function,
+        ),
     ):
         # Set up MockChunker to return the query terms as is
         mock_chunker_instance = MagicMock()
@@ -342,7 +354,7 @@ async def test_get_query_result_files_multiple_queries(mock_collection, mock_con
         # Check query was called with all query terms
         mock_collection.query.assert_called_once()
         _, kwargs = mock_collection.query.call_args
-        assert set(kwargs["query_texts"]) == set(["term1", "term2", "term3"])
+        mock_embedding_function.assert_called_once_with(["term1", "term2", "term3"])
 
         # Check the result
         assert result == ["file1.py", "file2.py"]

tests/subcommands/test_vectorise.py

@@ -6,6 +6,7 @@
 from contextlib import ExitStack
 from unittest.mock import AsyncMock, MagicMock, mock_open, patch
 
+import numpy
 import pytest
 from chromadb.api.models.AsyncCollection import AsyncCollection
 from tree_sitter import Point
@@ -511,6 +512,7 @@ def is_file_side_effect(path):
         else:
             return True
 
+    mock_embedding_function = MagicMock(return_value=numpy.random.random((100,)))
     with (
         patch("os.path.isfile", side_effect=is_file_side_effect),
         patch(
@@ -522,6 +524,10 @@ def is_file_side_effect(path):
             "vectorcode.subcommands.vectorise.get_collection",
             return_value=mock_collection,
         ),
+        patch(
+            "vectorcode.subcommands.vectorise.get_embedding_function",
+            return_value=mock_embedding_function,
+        ),
         patch("vectorcode.subcommands.vectorise.verify_ef", return_value=True),
         patch(
             "vectorcode.subcommands.vectorise.expand_globs",

tests/test_mcp.py

@@ -3,6 +3,7 @@
 from argparse import ArgumentParser
 from unittest.mock import AsyncMock, MagicMock, mock_open, patch
 
+import numpy
 import pytest
 from mcp import McpError
 
@@ -202,6 +203,7 @@ async def test_vectorise_files_success():
             f.write("def func(): pass")
         mock_client = AsyncMock()
 
+        mock_embedding_function = AsyncMock(return_value=numpy.random.random((100,)))
         with (
             patch("os.path.isdir", return_value=True),
             patch("vectorcode.mcp_main.get_project_config") as mock_get_project_config,
@@ -210,6 +212,10 @@ async def test_vectorise_files_success():
                 "vectorcode.mcp_main.ClientManager._create_client",
                 return_value=mock_client,
             ),
+            patch(
+                "vectorcode.subcommands.vectorise.get_embedding_function",
+                return_value=mock_embedding_function,
+            ),
             patch("vectorcode.subcommands.vectorise.chunked_add"),
             patch(
                 "vectorcode.subcommands.vectorise.hash_file", return_value="test_hash"