tests(cli): refactor tests and coverage for files rm

Author: Zhe Yu

src/vectorcode/subcommands/files/rm.py

@@ -2,20 +2,27 @@
 
 from vectorcode.cli_utils import Config
 from vectorcode.database import get_database_connector
+from vectorcode.database.errors import CollectionNotFoundError
 from vectorcode.database.types import ResultType
 
 logger = logging.getLogger(name=__name__)
 
 
 async def rm(configs: Config) -> int:
-    database = get_database_connector(configs)
-    remove_count = await database.delete()
+    try:
+        database = get_database_connector(configs)
+        remove_count = await database.delete()
 
-    if not configs.pipe:
-        print(f"Removed {remove_count} file(s).")
-    if await database.count(ResultType.chunk) == 0:
-        logger.warning(
-            f"The collection at {configs.project_root} is now empty, and will be removed."
-        )
-        await database.drop()
-    return 0
+        if not configs.pipe:
+            print(f"Removed {remove_count} file(s).")
+        if await database.count(ResultType.chunk) == 0:
+            logger.warning(
+                f"The collection at {configs.project_root} is now empty, and will be removed."
+            )
+            await database.drop()
+        return 0
+    except CollectionNotFoundError:
+        logger.error(f"There's no existing collection for `{configs.project_root}`.")
+        return 1
+    except Exception:  # pragma: nocover
+        raise

tests/subcommands/files/test_files_rm.py

@@ -1,101 +1,60 @@
+from pathlib import Path
 from unittest.mock import AsyncMock, patch
 
 import pytest
-from chromadb.api.models.AsyncCollection import AsyncCollection
 
-from vectorcode.cli_utils import CliAction, Config, FilesAction
+from vectorcode.cli_utils import Config
+from vectorcode.database.errors import CollectionNotFoundError
 from vectorcode.subcommands.files.rm import rm
 
 
 @pytest.fixture
-def client():
-    return AsyncMock()
+def mock_db():
+    db = AsyncMock()
 
+    def mock_delete():
+        count = 0
+        for f in db._configs.rm_paths:
+            if Path(f).name in {"file1.py", "file2.py", "file3.py"}:
+                count += 1
+        return count
 
-@pytest.fixture
-def collection():
-    col = AsyncMock(spec=AsyncCollection)
-    col.get.return_value = {
-        "ids": ["id1", "id2", "id3"],
-        "distances": [0.1, 0.2, 0.3],
-        "metadatas": [
-            {"path": "file1.py", "start": 1, "end": 1},
-            {"path": "file2.py", "start": 1, "end": 1},
-            {"path": "file3.py", "start": 1, "end": 1},
-        ],
-        "documents": [
-            "content1",
-            "content2",
-            "content3",
-        ],
-    }
-    col.name = "test_collection"
-    return col
+    db.delete = AsyncMock(side_effect=mock_delete)
+    return db
 
 
 @pytest.mark.asyncio
-async def test_rm(client, collection, capsys):
-    with (
-        patch("vectorcode.subcommands.files.rm.ClientManager") as MockClientManager,
-        patch(
-            "vectorcode.subcommands.files.rm.get_collection", return_value=collection
-        ),
-        patch("vectorcode.common.try_server", return_value=True),
-        patch("os.path.isfile", return_value=True),
-        patch(
-            "vectorcode.subcommands.files.rm.expand_path", side_effect=lambda x, y: x
-        ),
+async def test_rm(mock_db, capsys):
+    configs = Config(rm_paths=["file1.py", "file2.py"])
+    mock_db._configs = configs
+    with patch(
+        "vectorcode.subcommands.files.rm.get_database_connector", return_value=mock_db
     ):
-        MockClientManager.return_value._create_client.return_value = client
-        config = Config(
-            action=CliAction.files,
-            files_action=FilesAction.rm,
-            rm_paths=["file1.py"],
-        )
-        await rm(config)
-        collection.delete.assert_called_with(where={"path": {"$in": ["file1.py"]}})
+        assert await rm(configs) == 0
+        assert capsys.readouterr().out.strip() == "Removed 2 file(s)."
 
 
 @pytest.mark.asyncio
-async def test_rm_empty_collection(client, collection, capsys):
-    with (
-        patch(
-            "vectorcode.subcommands.files.rm.get_collection", return_value=collection
-        ),
-        patch("vectorcode.common.try_server", return_value=True),
-        patch("os.path.isfile", return_value=True),
-        patch(
-            "vectorcode.subcommands.files.rm.expand_path", side_effect=lambda x, y: x
-        ),
-        patch(
-            "vectorcode.subcommands.files.rm.ClientManager._create_client",
-            return_value=client,
-        ),
+async def test_rm_clean_after_rm(mock_db, capsys):
+    configs = Config(rm_paths=["file1.py", "file2.py"])
+    mock_db._configs = configs
+    mock_db.count = AsyncMock(return_value=0)
+    with patch(
+        "vectorcode.subcommands.files.rm.get_database_connector", return_value=mock_db
     ):
-        config = Config(
-            action=CliAction.files,
-            files_action=FilesAction.rm,
-            rm_paths=["file1.py"],
-        )
-        collection.count = AsyncMock(return_value=0)
-        client.delete_collection = AsyncMock()
-        await rm(config)
-        client.delete_collection.assert_called_once_with(collection.name)
+        assert await rm(configs) == 0
+        mock_db.drop.assert_called_once()
 
 
 @pytest.mark.asyncio
-async def test_rm_no_collection(client, collection, capsys):
-    with (
-        patch("vectorcode.subcommands.files.rm.ClientManager") as MockClientManager,
-        patch("vectorcode.subcommands.files.rm.get_collection", side_effect=ValueError),
+async def test_rm_no_collection(mock_db, capsys):
+    with patch(
+        "vectorcode.subcommands.files.rm.get_database_connector", return_value=mock_db
     ):
-        MockClientManager.return_value._create_client.return_value = client
+        mock_db.delete.side_effect = CollectionNotFoundError
         assert (
             await rm(
                 Config(
-                    action=CliAction.files,
-                    files_action=FilesAction.rm,
-                    pipe=True,
                     rm_paths=["file1.py"],
                 )
             )