tests(cli): fix broken tests

Author: Zhe Yu

tests/test_mcp.py

@@ -1,7 +1,6 @@
 import os
 import tempfile
 from argparse import ArgumentParser
-from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, MagicMock, mock_open, patch
 
 import pytest
@@ -25,9 +24,7 @@ async def test_list_collections_success():
         patch("vectorcode.common.ClientManager") as MockClientManager,
     ):
         mock_client = AsyncMock()
-        MockClientManager.return_value._create_client = AsyncMock(
-            return_value=mock_client
-        )
+        MockClientManager.return_value._create_client.return_value = mock_client
 
         mock_collection1 = AsyncMock()
         mock_collection1.metadata = {"path": "path1"}
@@ -51,9 +48,8 @@ async def test_list_collections_no_metadata():
         patch("vectorcode.common.ClientManager") as MockClientManager,
     ):
         mock_client = AsyncMock()
-        MockClientManager.return_value._create_client = asynccontextmanager(
-            AsyncMock(return_value=mock_client)
-        )
+        MockClientManager.return_value._create_client.return_value = mock_client
+
         mock_collection1 = AsyncMock()
         mock_collection1.metadata = {"path": "path1"}
         mock_collection2 = AsyncMock()
@@ -103,9 +99,7 @@ async def test_query_tool_success():
         mock_load_config_file.return_value = mock_config
         mock_get_project_config.return_value = mock_config
         mock_client = AsyncMock()
-        MockClientManager.return_value._create_client = AsyncMock(
-            return_value=mock_client
-        )
+        MockClientManager.return_value._create_client.return_value = mock_client
 
         # Mock the collection's query method to return a valid QueryResult
         mock_collection = AsyncMock()
@@ -169,7 +163,8 @@ async def test_query_tool_no_collection():
         ) as mock_get_collection,  # Still mock get_collection
         patch("vectorcode.common.ClientManager") as MockClientManager,
     ):
-        MockClientManager.return_value._create_client.return_value = AsyncMock()
+        mock_client = AsyncMock()
+        MockClientManager.return_value._create_client.return_value = mock_client
         mock_get_collection.return_value = None
 
         with pytest.raises(McpError) as exc_info:
@@ -215,9 +210,7 @@ async def test_vectorise_files_success():
             mock_client = AsyncMock()
 
             # Ensure ClientManager's internal client creation method returns our mock.
-            MockClientManager.return_value._create_client = AsyncMock(
-                return_value=mock_client
-            )
+            MockClientManager.return_value._create_client.return_value = mock_client
 
             mock_collection = AsyncMock()
             mock_collection.get.return_value = {"ids": [], "metadatas": []}
@@ -243,11 +236,9 @@ async def test_vectorise_files_collection_access_failure():  # Removed client_ma
         ) as MockClientManager,  # Patch ClientManager class
         patch("vectorcode.mcp_main.get_collection"),
     ):
-
-        async def failing_get_client(*args, **kwargs):
-            raise Exception("Client error")
-
-        MockClientManager.return_value._create_client = failing_get_client
+        MockClientManager.return_value._create_client.side_effect = Exception(
+            "Client error"
+        )
 
         with pytest.raises(McpError) as exc_info:
             await vectorise_files(paths=["file.py"], project_root="/valid/path")
@@ -301,9 +292,7 @@ def mock_open_side_effect(filename, *args, **kwargs):
             mock_config = Config(project_root=temp_dir)
             mock_get_project_config.return_value = mock_config
             mock_client = AsyncMock()
-            MockClientManager.return_value._create_client = AsyncMock(
-                return_value=mock_client
-            )
+            MockClientManager.return_value._create_client.return_value = mock_client
 
             mock_collection = AsyncMock()
             mock_collection.get.return_value = {"ids": [], "metadatas": []}
@@ -336,7 +325,7 @@ async def test_mcp_server():
         mock_load_config_file.return_value = Config(project_root="/path/to/project")
         mock_client = AsyncMock()
 
-        MockClientManager.return_value.get_client = AsyncMock(return_value=mock_client)
+        MockClientManager.return_value._create_client.return_value = mock_client
         mock_collection = AsyncMock()
         mock_get_collection.return_value = mock_collection
 
@@ -347,30 +336,29 @@ async def test_mcp_server():
 
 @pytest.mark.asyncio
 async def test_mcp_server_ls_on_start():
+    mock_collection = AsyncMock()
+
     with (
         patch(
             "vectorcode.mcp_main.find_project_config_dir"
         ) as mock_find_project_config_dir,
         patch("vectorcode.mcp_main.load_config_file") as mock_load_config_file,
-        # patch("vectorcode.mcp_main.get_client") as mock_get_client, # Removed
         patch("vectorcode.mcp_main.get_collection") as mock_get_collection,
         patch(
             "vectorcode.mcp_main.get_collections", spec=AsyncMock
         ) as mock_get_collections,
         patch("mcp.server.fastmcp.FastMCP.add_tool") as mock_add_tool,
-        patch("vectorcode.common.ClientManager") as MockClientManager,  # Added
+        patch("vectorcode.common.try_server", return_value=True),
     ):
-        from vectorcode.mcp_main import mcp_config
+        from vectorcode.mcp_main import ClientManager, mcp_config
 
         mcp_config.ls_on_start = True
         mock_find_project_config_dir.return_value = "/path/to/config"
         mock_load_config_file.return_value = Config(project_root="/path/to/project")
         mock_client = AsyncMock()
 
-        MockClientManager.return_value._create_client = AsyncMock(
-            return_value=mock_client
-        )
-        mock_collection = AsyncMock()
+        ClientManager._create_client = AsyncMock(return_value=mock_client)
+
         mock_collection.metadata = {"path": "/path/to/project"}
         mock_get_collection.return_value = mock_collection