tests(cli): Refactor drop subcommand and add exception handling

Author: Zhe Yu

src/vectorcode/subcommands/drop.py

@@ -8,12 +8,14 @@
 
 
 async def drop(config: Config) -> int:
-    database = get_database_connector(config)
     try:
+        database = get_database_connector(config)
         await database.drop()
         if not config.pipe:
             print(f"Collection for {config.project_root} has been deleted.")
         return 0
     except CollectionNotFoundError:
         logger.warning(f"Collection for {config.project_root} doesn't exist.")
         return 1
+    except Exception:  # pragma: nocover
+        raise

tests/subcommands/test_drop.py

@@ -1,65 +1,27 @@
-from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, patch
 
 import pytest
 
 from vectorcode.cli_utils import Config
+from vectorcode.database.errors import CollectionNotFoundError
 from vectorcode.subcommands.drop import drop
 
 
-@pytest.fixture
-def mock_config():
-    config = Config(
-        project_root="/path/to/project",
-    )  # Removed positional args
-    return config
-
-
-@pytest.fixture
-def mock_client():
-    return AsyncMock()
-
-
-@pytest.fixture
-def mock_collection():
-    collection = AsyncMock()
-    collection.name = "test_collection"
-    collection.metadata = {"path": "/path/to/project"}
-    return collection
-
-
 @pytest.mark.asyncio
-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.ClientManager") as MockClientManager,
-        patch(
-            "vectorcode.subcommands.drop.get_collection", return_value=mock_collection
-        ),
+async def test_drop_success():
+    mock_db = AsyncMock()
+    with patch(
+        "vectorcode.subcommands.drop.get_database_connector", return_value=mock_db
     ):
-        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)
+        await drop(config=Config(project_root="DummyDir"))
+        mock_db.drop.assert_called_once()
 
 
 @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.ClientManager"):
-        with patch(
-            "vectorcode.subcommands.drop.get_collection",
-            side_effect=ValueError("Collection not found"),
-        ):
-            result = await drop(mock_config)
-            assert result == 1
+async def test_drop_collection_not_found():
+    mock_db = AsyncMock()
+    mock_db.drop = AsyncMock(side_effect=CollectionNotFoundError)
+    with patch(
+        "vectorcode.subcommands.drop.get_database_connector", return_value=mock_db
+    ):
+        assert await drop(config=Config(project_root="DummyDir")) != 0