tests(cli): fix failed tests due to ClientManager refactor.

Author: Davidyz

tests/subcommands/query/test_query.py

@@ -355,7 +355,7 @@ async def test_query_success(mock_config):
     mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.query.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.query.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.query.get_collection", return_value=mock_collection
         ),
@@ -367,6 +367,7 @@ async def test_query_success(mock_config):
         patch("os.path.relpath", return_value="rel/path.py"),
         patch("os.path.abspath", return_value="/abs/path.py"),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         # Set up the mock file paths and contents
         mock_get_files.return_value = ["file1.py", "file2.py"]
         mock_file_handle = MagicMock()
@@ -396,7 +397,7 @@ async def test_query_pipe_mode(mock_config):
     mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.query.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.query.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.query.get_collection", return_value=mock_collection
         ),
@@ -408,6 +409,7 @@ async def test_query_pipe_mode(mock_config):
         patch("os.path.relpath", return_value="rel/path.py"),
         patch("os.path.abspath", return_value="/abs/path.py"),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         # Set up the mock file paths and contents
         mock_get_files.return_value = ["file1.py", "file2.py"]
         mock_file_handle = MagicMock()
@@ -434,7 +436,7 @@ async def test_query_absolute_path(mock_config):
     mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.query.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.query.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.query.get_collection", return_value=mock_collection
         ),
@@ -445,6 +447,7 @@ async def test_query_absolute_path(mock_config):
         patch("os.path.relpath", return_value="rel/path.py"),
         patch("os.path.abspath", return_value="/abs/path.py"),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         # Set up the mock file paths and contents
         mock_get_files.return_value = ["file1.py"]
         mock_file_handle = MagicMock()
@@ -463,7 +466,7 @@ async def test_query_collection_not_found():
     config = Config(project_root="/test/project")
 
     with (
-        patch("vectorcode.subcommands.query.get_client"),
+        patch("vectorcode.subcommands.query.ClientManager"),
         patch("vectorcode.subcommands.query.get_collection") as mock_get_collection,
         patch("sys.stderr"),
     ):
@@ -482,7 +485,7 @@ async def test_query_invalid_collection():
     config = Config(project_root="/test/project")
 
     with (
-        patch("vectorcode.subcommands.query.get_client"),
+        patch("vectorcode.subcommands.query.ClientManager"),
         patch("vectorcode.subcommands.query.get_collection") as mock_get_collection,
         patch("sys.stderr"),
     ):
@@ -503,7 +506,7 @@ async def test_query_invalid_dimension():
     config = Config(project_root="/test/project")
 
     with (
-        patch("vectorcode.subcommands.query.get_client"),
+        patch("vectorcode.subcommands.query.ClientManager"),
         patch("vectorcode.subcommands.query.get_collection") as mock_get_collection,
         patch("sys.stderr"),
     ):
@@ -524,14 +527,15 @@ async def test_query_invalid_file(mock_config):
     mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.query.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.query.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.query.get_collection", return_value=mock_collection
         ),
         patch("vectorcode.subcommands.query.verify_ef", return_value=True),
         patch("vectorcode.subcommands.query.get_query_result_files") as mock_get_files,
         patch("os.path.isfile", return_value=False),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         # Set up the mock file paths
         mock_get_files.return_value = ["invalid_file.py"]
 
@@ -549,12 +553,13 @@ async def test_query_invalid_ef(mock_config):
     mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.query.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.query.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.query.get_collection", return_value=mock_collection
         ),
         patch("vectorcode.subcommands.query.verify_ef", return_value=False),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         # Call the function
         result = await query(mock_config)
 
@@ -580,13 +585,14 @@ async def test_query_chunk_mode_no_metadata_fallback(mock_config):
     mock_collection.get.return_value = {"ids": []}
 
     with (
-        patch("vectorcode.subcommands.query.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.query.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.query.get_collection", return_value=mock_collection
         ),
         patch("vectorcode.subcommands.query.verify_ef", return_value=True),
         patch("vectorcode.subcommands.query.build_query_results") as mock_build_results,
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         mock_build_results.return_value = []  # Return empty results for simplicity
 
         result = await query(mock_config)

tests/subcommands/test_clean.py

@@ -73,10 +73,10 @@ async def mock_get_collections(client):
 
 @pytest.mark.asyncio
 async def test_clean():
-    mock_client = AsyncMock(spec=AsyncClientAPI)
+    AsyncMock(spec=AsyncClientAPI)
     mock_config = Config(pipe=False)
 
-    with patch("vectorcode.subcommands.clean.get_client", return_value=mock_client):
+    with patch("vectorcode.subcommands.clean.ClientManager"):
         result = await clean(mock_config)
 
     assert result == 0

tests/subcommands/test_drop.py

@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -31,19 +32,31 @@ def mock_collection():
 async def test_drop_success(mock_config, mock_client, mock_collection):
     mock_client.get_collection.return_value = mock_collection
     mock_client.delete_collection = AsyncMock()
-    with patch("vectorcode.subcommands.drop.get_client", return_value=mock_client):
-        with patch(
+    with (
+        patch("vectorcode.subcommands.drop.ClientManager") as MockClientManager,
+        patch(
             "vectorcode.subcommands.drop.get_collection", return_value=mock_collection
-        ):
-            result = await drop(mock_config)
-            assert result == 0
-            mock_client.delete_collection.assert_called_once_with(mock_collection.name)
+        ),
+    ):
+        mock_client = AsyncMock()
+
+        @asynccontextmanager
+        async def _get_client(self, config=None, need_lock=True):
+            yield mock_client
+
+        mock_client_manager = MockClientManager.return_value
+        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
+        mock_client_manager.get_client = _get_client
+
+        result = await drop(mock_config)
+        assert result == 0
+        mock_client.delete_collection.assert_called_once_with(mock_collection.name)
 
 
 @pytest.mark.asyncio
 async def test_drop_collection_not_found(mock_config, mock_client):
     mock_client.get_collection.side_effect = ValueError("Collection not found")
-    with patch("vectorcode.subcommands.drop.get_client", return_value=mock_client):
+    with patch("vectorcode.subcommands.drop.ClientManager"):
         with patch(
             "vectorcode.subcommands.drop.get_collection",
             side_effect=ValueError("Collection not found"),

tests/subcommands/test_ls.py

@@ -1,6 +1,6 @@
 import json
 import socket
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 import tabulate
@@ -77,7 +77,7 @@ async def mock_get_collections(client):
         yield mock_collection
 
     with (
-        patch("vectorcode.subcommands.ls.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.ls.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.ls.get_collection_list",
             return_value=[
@@ -90,6 +90,10 @@ async def mock_get_collections(client):
             ],
         ),
     ):
+        mock_client = MagicMock()
+        mock_client_manager = MockClientManager.return_value
+        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
+
         config = Config(pipe=True)
         await ls(config)
         captured = capsys.readouterr()
@@ -126,7 +130,7 @@ async def mock_get_collections(client):
         yield mock_collection
 
     with (
-        patch("vectorcode.subcommands.ls.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.ls.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.ls.get_collection_list",
             return_value=[
@@ -139,6 +143,10 @@ async def mock_get_collections(client):
             ],
         ),
     ):
+        mock_client = MagicMock()
+        mock_client_manager = MockClientManager.return_value
+        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
+
         config = Config(pipe=False)
         await ls(config)
         captured = capsys.readouterr()
@@ -159,7 +167,7 @@ async def mock_get_collections(client):
     # Test with HOME environment variable set
     monkeypatch.setenv("HOME", "/test")
     with (
-        patch("vectorcode.subcommands.ls.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.ls.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.ls.get_collection_list",
             return_value=[
@@ -172,6 +180,9 @@ async def mock_get_collections(client):
             ],
         ),
     ):
+        mock_client = MagicMock()
+        mock_client_manager = MockClientManager.return_value
+        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
         config = Config(pipe=False)
         await ls(config)
         captured = capsys.readouterr()

tests/subcommands/test_update.py

@@ -19,7 +19,7 @@ async def test_update_success():
     mock_client.get_max_batch_size.return_value = 100
 
     with (
-        patch("vectorcode.subcommands.update.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.update.ClientManager"),
         patch(
             "vectorcode.subcommands.update.get_collection", return_value=mock_collection
         ),
@@ -50,7 +50,7 @@ async def test_update_with_orphans():
     mock_client.get_max_batch_size.return_value = 100
 
     with (
-        patch("vectorcode.subcommands.update.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.update.ClientManager"),
         patch(
             "vectorcode.subcommands.update.get_collection", return_value=mock_collection
         ),
@@ -78,10 +78,11 @@ async def test_update_index_error():
     # mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.update.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.update.ClientManager") as MockClientManager,
         patch("vectorcode.subcommands.update.get_collection", side_effect=IndexError),
         patch("sys.stderr"),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         config = Config(project_root="/test/project", pipe=False)
         result = await update(config)
 
@@ -94,10 +95,11 @@ async def test_update_value_error():
     # mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.update.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.update.ClientManager") as MockClientManager,
         patch("vectorcode.subcommands.update.get_collection", side_effect=ValueError),
         patch("sys.stderr"),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         config = Config(project_root="/test/project", pipe=False)
         result = await update(config)
 
@@ -110,13 +112,14 @@ async def test_update_invalid_collection_exception():
     # mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.update.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.update.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.update.get_collection",
             side_effect=InvalidCollectionException,
         ),
         patch("sys.stderr"),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         config = Config(project_root="/test/project", pipe=False)
         result = await update(config)