11import os
22import tempfile
33from argparse import ArgumentParser
4- from contextlib import asynccontextmanager
54from unittest .mock import AsyncMock , MagicMock , mock_open , patch
65
76import pytest
2221async def test_list_collections_success ():
2322 with (
2423 patch ("vectorcode.mcp_main.get_collections" ) as mock_get_collections ,
25- patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
2624 ):
25+ from vectorcode .mcp_main import ClientManager
26+
2727 mock_client = AsyncMock ()
28- MockClientManager .return_value ._create_client = AsyncMock (
29- return_value = mock_client
30- )
28+ ClientManager ._create_client = AsyncMock (return_value = mock_client )
3129
3230 mock_collection1 = AsyncMock ()
3331 mock_collection1 .metadata = {"path" : "path1" }
@@ -48,12 +46,12 @@ async def async_generator():
4846async def test_list_collections_no_metadata ():
4947 with (
5048 patch ("vectorcode.mcp_main.get_collections" ) as mock_get_collections ,
51- patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
5249 ):
50+ from vectorcode .mcp_main import ClientManager
51+
5352 mock_client = AsyncMock ()
54- MockClientManager .return_value ._create_client = asynccontextmanager (
55- AsyncMock (return_value = mock_client )
56- )
53+ ClientManager ._create_client = AsyncMock (return_value = mock_client )
54+
5755 mock_collection1 = AsyncMock ()
5856 mock_collection1 .metadata = {"path" : "path1" }
5957 mock_collection2 = AsyncMock ()
@@ -97,15 +95,14 @@ async def test_query_tool_success():
9795 patch ("os.path.isfile" , return_value = True ),
9896 patch ("os.path.relpath" , return_value = "rel/path.py" ),
9997 patch ("vectorcode.cli_utils.load_config_file" ) as mock_load_config_file ,
100- patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
10198 ):
99+ from vectorcode .mcp_main import ClientManager
100+
102101 mock_config = Config (chunk_size = 100 , overlap_ratio = 0.1 , reranker = None )
103102 mock_load_config_file .return_value = mock_config
104103 mock_get_project_config .return_value = mock_config
105104 mock_client = AsyncMock ()
106- MockClientManager .return_value ._create_client = AsyncMock (
107- return_value = mock_client
108- )
105+ ClientManager ._create_client = AsyncMock (return_value = mock_client )
109106
110107 # Mock the collection's query method to return a valid QueryResult
111108 mock_collection = AsyncMock ()
@@ -139,13 +136,13 @@ async def test_query_tool_collection_access_failure():
139136 patch ("os.path.isdir" , return_value = True ),
140137 patch ("vectorcode.mcp_main.get_project_config" ),
141138 patch ("vectorcode.mcp_main.get_collection" ), # Still mock get_collection
142- patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
143139 ):
140+ from vectorcode .mcp_main import ClientManager
144141
145142 async def failing_get_client (* args , ** kwargs ):
146143 raise Exception ("Failed to connect" )
147144
148- MockClientManager . return_value . _create_client . side_effect = failing_get_client
145+ ClientManager . _create_client = AsyncMock ( side_effect = failing_get_client )
149146
150147 with pytest .raises (McpError ) as exc_info :
151148 await query_tool (
@@ -169,7 +166,8 @@ async def test_query_tool_no_collection():
169166 ) as mock_get_collection , # Still mock get_collection
170167 patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
171168 ):
172- MockClientManager .return_value ._create_client .return_value = AsyncMock ()
169+ mock_client = AsyncMock ()
170+ MockClientManager .return_value ._create_client .return_value = mock_client
173171 mock_get_collection .return_value = None
174172
175173 with pytest .raises (McpError ) as exc_info :
@@ -201,7 +199,6 @@ async def test_vectorise_files_success():
201199 f .write ("def func(): pass" )
202200
203201 with (
204- patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
205202 patch ("os.path.isdir" , return_value = True ),
206203 patch ("vectorcode.mcp_main.get_project_config" ) as mock_get_project_config ,
207204 patch ("vectorcode.mcp_main.get_collection" ) as mock_get_collection ,
@@ -210,14 +207,14 @@ async def test_vectorise_files_success():
210207 "vectorcode.subcommands.vectorise.hash_file" , return_value = "test_hash"
211208 ),
212209 ):
210+ from vectorcode .mcp_main import ClientManager
211+
213212 mock_config = Config (project_root = temp_dir )
214213 mock_get_project_config .return_value = mock_config
215214 mock_client = AsyncMock ()
216215
217216 # Ensure ClientManager's internal client creation method returns our mock.
218- MockClientManager .return_value ._create_client = AsyncMock (
219- return_value = mock_client
220- )
217+ ClientManager ._create_client = AsyncMock (return_value = mock_client )
221218
222219 mock_collection = AsyncMock ()
223220 mock_collection .get .return_value = {"ids" : [], "metadatas" : []}
@@ -233,21 +230,16 @@ async def test_vectorise_files_success():
233230
234231
235232@pytest .mark .asyncio
236- async def test_vectorise_files_collection_access_failure (): # Removed client_manager fixture
233+ async def test_vectorise_files_collection_access_failure ():
237234 with (
238235 patch ("os.path.isdir" , return_value = True ),
239236 patch ("vectorcode.mcp_main.get_project_config" ),
240- # patch("vectorcode.mcp_main.get_client", side_effect=Exception("Client error")), # Removed explicit patch
241- patch (
242- "vectorcode.common.ClientManager"
243- ) as MockClientManager , # Patch ClientManager class
237+ patch ("vectorcode.common.ClientManager" ), # Patch ClientManager class
244238 patch ("vectorcode.mcp_main.get_collection" ),
245239 ):
240+ from vectorcode .mcp_main import ClientManager
246241
247- async def failing_get_client (* args , ** kwargs ):
248- raise Exception ("Client error" )
249-
250- MockClientManager .return_value ._create_client = failing_get_client
242+ ClientManager ._create_client = AsyncMock (side_effect = Exception ("Client error" ))
251243
252244 with pytest .raises (McpError ) as exc_info :
253245 await vectorise_files (paths = ["file.py" ], project_root = "/valid/path" )
@@ -296,14 +288,13 @@ def mock_open_side_effect(filename, *args, **kwargs):
296288 "os.path.isfile" ,
297289 side_effect = lambda x : x in [file1 , excluded_file , exclude_spec_file ],
298290 ),
299- patch ("vectorcode.common.ClientManager" ) as MockClientManager ,
300291 ):
292+ from vectorcode .mcp_main import ClientManager
293+
301294 mock_config = Config (project_root = temp_dir )
302295 mock_get_project_config .return_value = mock_config
303296 mock_client = AsyncMock ()
304- MockClientManager .return_value ._create_client = AsyncMock (
305- return_value = mock_client
306- )
297+ ClientManager ._create_client = AsyncMock (return_value = mock_client )
307298
308299 mock_collection = AsyncMock ()
309300 mock_collection .get .return_value = {"ids" : [], "metadatas" : []}
@@ -330,13 +321,14 @@ async def test_mcp_server():
330321 # patch("vectorcode.mcp_main.get_client") as mock_get_client, # Removed
331322 patch ("vectorcode.mcp_main.get_collection" ) as mock_get_collection ,
332323 patch ("mcp.server.fastmcp.FastMCP.add_tool" ) as mock_add_tool ,
333- patch ("vectorcode.common.ClientManager" ) as MockClientManager , # Added
334324 ):
325+ from vectorcode .mcp_main import ClientManager
326+
335327 mock_find_project_config_dir .return_value = "/path/to/config"
336328 mock_load_config_file .return_value = Config (project_root = "/path/to/project" )
337329 mock_client = AsyncMock ()
338330
339- MockClientManager . return_value . get_client = AsyncMock (return_value = mock_client )
331+ ClientManager . _create_client = AsyncMock (return_value = mock_client )
340332 mock_collection = AsyncMock ()
341333 mock_get_collection .return_value = mock_collection
342334
@@ -347,30 +339,29 @@ async def test_mcp_server():
347339
348340@pytest .mark .asyncio
349341async def test_mcp_server_ls_on_start ():
342+ mock_collection = AsyncMock ()
343+
350344 with (
351345 patch (
352346 "vectorcode.mcp_main.find_project_config_dir"
353347 ) as mock_find_project_config_dir ,
354348 patch ("vectorcode.mcp_main.load_config_file" ) as mock_load_config_file ,
355- # patch("vectorcode.mcp_main.get_client") as mock_get_client, # Removed
356349 patch ("vectorcode.mcp_main.get_collection" ) as mock_get_collection ,
357350 patch (
358351 "vectorcode.mcp_main.get_collections" , spec = AsyncMock
359352 ) as mock_get_collections ,
360353 patch ("mcp.server.fastmcp.FastMCP.add_tool" ) as mock_add_tool ,
361- patch ("vectorcode.common.ClientManager" ) as MockClientManager , # Added
354+ patch ("vectorcode.common.try_server" , return_value = True ),
362355 ):
363- from vectorcode .mcp_main import mcp_config
356+ from vectorcode .mcp_main import ClientManager , mcp_config
364357
365358 mcp_config .ls_on_start = True
366359 mock_find_project_config_dir .return_value = "/path/to/config"
367360 mock_load_config_file .return_value = Config (project_root = "/path/to/project" )
368361 mock_client = AsyncMock ()
369362
370- MockClientManager .return_value ._create_client = AsyncMock (
371- return_value = mock_client
372- )
373- mock_collection = AsyncMock ()
363+ ClientManager ._create_client = AsyncMock (return_value = mock_client )
364+
374365 mock_collection .metadata = {"path" : "/path/to/project" }
375366 mock_get_collection .return_value = mock_collection
376367
0 commit comments