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

Author: Zhe Yu

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)
 

tests/subcommands/test_vectorise.py

@@ -370,9 +370,7 @@ async def test_vectorise(capsys):
 
     with ExitStack() as stack:
         stack.enter_context(
-            patch(
-                "vectorcode.subcommands.vectorise.get_client", return_value=mock_client
-            )
+            patch("vectorcode.subcommands.vectorise.ClientManager"),
         )
         stack.enter_context(patch("os.path.isfile", return_value=False))
         stack.enter_context(
@@ -427,7 +425,7 @@ async def mock_chunked_add(*args, **kwargs):
             "vectorcode.subcommands.vectorise.chunked_add", side_effect=mock_chunked_add
         ) as mock_add,
         patch("sys.stderr") as mock_stderr,
-        patch("vectorcode.subcommands.vectorise.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.vectorise.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.vectorise.get_collection",
             return_value=mock_collection,
@@ -438,6 +436,7 @@ async def mock_chunked_add(*args, **kwargs):
             lambda x: not (x.endswith("gitignore") or x.endswith("vectorcode.exclude")),
         ),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         result = await vectorise(configs)
         assert result == 1
         mock_add.assert_called_once()
@@ -458,7 +457,7 @@ async def test_vectorise_orphaned_files():
         pipe=False,
     )
 
-    mock_client = AsyncMock()
+    AsyncMock()
     mock_collection = AsyncMock()
 
     # Define a mock response for collection.get in vectorise
@@ -494,7 +493,7 @@ def is_file_side_effect(path):
             "vectorcode.subcommands.vectorise.TreeSitterChunker",
             return_value=mock_chunker,
         ),
-        patch("vectorcode.subcommands.vectorise.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.vectorise.ClientManager"),
         patch(
             "vectorcode.subcommands.vectorise.get_collection",
             return_value=mock_collection,
@@ -532,10 +531,11 @@ async def test_vectorise_collection_index_error():
     mock_client = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.vectorise.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.vectorise.ClientManager") as MockClientManager,
         patch("vectorcode.subcommands.vectorise.get_collection") as mock_get_collection,
         patch("os.path.isfile", return_value=False),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         mock_get_collection.side_effect = IndexError("Collection not found")
         result = await vectorise(configs)
         assert result == 1
@@ -558,14 +558,15 @@ async def test_vectorise_verify_ef_false():
     mock_collection = AsyncMock()
 
     with (
-        patch("vectorcode.subcommands.vectorise.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.vectorise.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.vectorise.get_collection",
             return_value=mock_collection,
         ),
         patch("vectorcode.subcommands.vectorise.verify_ef", return_value=False),
         patch("os.path.isfile", return_value=False),
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         result = await vectorise(configs)
         assert result == 1
 
@@ -588,7 +589,7 @@ async def test_vectorise_gitignore():
     mock_collection.get.return_value = {"metadatas": []}
 
     with (
-        patch("vectorcode.subcommands.vectorise.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.vectorise.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.vectorise.get_collection",
             return_value=mock_collection,
@@ -608,6 +609,7 @@ async def test_vectorise_gitignore():
             "vectorcode.subcommands.vectorise.exclude_paths_by_spec"
         ) as mock_exclude_paths,
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         await vectorise(configs)
         mock_exclude_paths.assert_called_once()
 
@@ -635,7 +637,7 @@ async def test_vectorise_exclude_file(tmpdir):
     mock_collection.get.return_value = {"ids": []}
 
     with (
-        patch("vectorcode.subcommands.vectorise.get_client", return_value=mock_client),
+        patch("vectorcode.subcommands.vectorise.ClientManager") as MockClientManager,
         patch(
             "vectorcode.subcommands.vectorise.get_collection",
             return_value=mock_collection,
@@ -652,6 +654,7 @@ async def test_vectorise_exclude_file(tmpdir):
         ),
         patch("vectorcode.subcommands.vectorise.chunked_add") as mock_chunked_add,
     ):
+        MockClientManager.return_value._create_client.return_value = mock_client
         await vectorise(configs)
         # Assert that chunked_add is only called for test_file.py, not excluded_file.py
         call_args = [call[0][0] for call in mock_chunked_add.call_args_list]
@@ -664,7 +667,6 @@ async def test_vectorise_exclude_file(tmpdir):
 
 
 @pytest.mark.asyncio
-@patch("vectorcode.subcommands.vectorise.get_client", new_callable=AsyncMock)
 @patch("vectorcode.subcommands.vectorise.get_collection", new_callable=AsyncMock)
 @patch("vectorcode.subcommands.vectorise.expand_globs", new_callable=AsyncMock)
 @patch("vectorcode.subcommands.vectorise.chunked_add", new_callable=AsyncMock)
@@ -681,7 +683,6 @@ async def test_vectorise_uses_global_exclude_when_local_missing(
     mock_chunked_add,
     mock_expand_globs,
     mock_get_collection,
-    mock_get_client,
     tmp_path,
 ):
     """
@@ -712,14 +713,20 @@ def isfile_side_effect(p):
 
     global_exclude_content = "*.bin"
     m_open = mock_open(read_data=global_exclude_content)
-    with patch("builtins.open", m_open):
+    with (
+        patch("builtins.open", m_open),
+        patch("vectorcode.subcommands.vectorise.ClientManager") as MockClientManager,
+    ):
         mock_spec_instance = MagicMock()
         mock_spec_instance.match_file = lambda path: str(path).endswith(".bin")
         mock_gitignore_spec.from_lines.return_value = mock_spec_instance
 
         mock_client_instance = AsyncMock()
         mock_client_instance.get_max_batch_size = AsyncMock(return_value=100)
-        mock_get_client.return_value = mock_client_instance
+
+        MockClientManager.return_value._create_client.return_value = (
+            mock_client_instance
+        )
 
         mock_collection_instance = AsyncMock()
         mock_collection_instance.get = AsyncMock(