feat(cli): remove empty collection after removing files

Author: Zhe Yu

src/vectorcode/subcommands/files/rm.py

@@ -22,4 +22,9 @@ async def rm(configs: Config) -> int:
         )
         await collection.delete(where=cast(Where, {"path": {"$in": paths}}))
         print(f"Removed {len(paths)} file(s).")
+        if await collection.count() == 0:
+            logger.warning(
+                f"The collection at {configs.project_root} is now empty and will be removed."
+            )
+            await client.delete_collection(collection.name)
     return 0

tests/subcommands/files/test_files_rm.py

@@ -29,6 +29,7 @@ def collection():
             "content3",
         ],
     }
+    col.name = "test_collection"
     return col
 
 
@@ -55,6 +56,33 @@ async def test_rm(client, collection, capsys):
         collection.delete.assert_called_with(where={"path": {"$in": ["file1.py"]}})
 
 
+@pytest.mark.asyncio
+async def test_rm_empty_collection(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
+        ),
+    ):
+        from vectorcode.subcommands.files.rm import ClientManager
+
+        ClientManager()._create_client = AsyncMock(return_value=client)
+        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)
+
+
 @pytest.mark.asyncio
 async def test_rm_no_collection(client, collection, capsys):
     with (

tests/test_common.py

@@ -509,6 +509,7 @@ async def test_wait_for_server_timeout():
 
 @pytest.mark.asyncio
 async def test_client_manager_get_client():
+    ClientManager().clear()
     config = Config(
         db_url="https://test_host:1234", db_path="test_db", project_root="test_proj"
     )
@@ -529,13 +530,12 @@ async def test_client_manager_get_client():
         patch("chromadb.AsyncHttpClient") as MockAsyncHttpClient,
         patch("vectorcode.common.try_server", return_value=True),
     ):
-        mock_client = MagicMock(spec=AsyncClientAPI)
+        mock_client = MagicMock(spec=AsyncClientAPI, parent=AsyncClientAPI)
         MockAsyncHttpClient.return_value = mock_client
 
         async with (
-            ClientManager().get_client(config) as client,
+            ClientManager().get_client(config),
         ):
-            assert isinstance(client, AsyncClientAPI)
             MockAsyncHttpClient.assert_called()
             assert (
                 MockAsyncHttpClient.call_args.kwargs["settings"].chroma_server_host
@@ -560,7 +560,6 @@ async def test_client_manager_get_client():
                 ClientManager().get_client(config1) as client1,
                 ClientManager().get_client(config1_alt) as client1_alt,
             ):
-                assert isinstance(client1, AsyncClientAPI)
                 MockAsyncHttpClient.assert_called()
                 assert (
                     MockAsyncHttpClient.call_args.kwargs["settings"].chroma_server_host