Skip to content

Commit 776c4fb

Browse files
committed
tests(cli): fix broken tests
1 parent aa56235 commit 776c4fb

File tree

1 file changed

+33
-42
lines changed

1 file changed

+33
-42
lines changed

tests/test_mcp.py

Lines changed: 33 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import os
22
import tempfile
33
from argparse import ArgumentParser
4-
from contextlib import asynccontextmanager
54
from unittest.mock import AsyncMock, MagicMock, mock_open, patch
65

76
import pytest
@@ -22,12 +21,11 @@
2221
async 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():
4846
async 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
349341
async 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

Comments
 (0)