Skip to content

Commit f7ceee6

Browse files
author
Zhe Yu
committed
feat(cli): remove empty collection after removing files
1 parent 4621f5b commit f7ceee6

3 files changed

Lines changed: 34 additions & 1 deletion

File tree

src/vectorcode/subcommands/files/rm.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,9 @@ async def rm(configs: Config) -> int:
2222
)
2323
await collection.delete(where=cast(Where, {"path": {"$in": paths}}))
2424
print(f"Removed {len(paths)} file(s).")
25+
if await collection.count() == 0:
26+
logger.warning(
27+
f"The collection at {configs.project_root} is now empty and will be removed."
28+
)
29+
await client.delete_collection(collection.name)
2530
return 0

tests/subcommands/files/test_files_rm.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def collection():
2929
"content3",
3030
],
3131
}
32+
col.name = "test_collection"
3233
return col
3334

3435

@@ -55,6 +56,33 @@ async def test_rm(client, collection, capsys):
5556
collection.delete.assert_called_with(where={"path": {"$in": ["file1.py"]}})
5657

5758

59+
@pytest.mark.asyncio
60+
async def test_rm_empty_collection(client, collection, capsys):
61+
with (
62+
# patch("vectorcode.subcommands.files.rm.ClientManager") as MockClientManager,
63+
patch(
64+
"vectorcode.subcommands.files.rm.get_collection", return_value=collection
65+
),
66+
patch("vectorcode.common.try_server", return_value=True),
67+
patch("os.path.isfile", return_value=True),
68+
patch(
69+
"vectorcode.subcommands.files.rm.expand_path", side_effect=lambda x, y: x
70+
),
71+
):
72+
from vectorcode.subcommands.files.rm import ClientManager
73+
74+
ClientManager()._create_client = AsyncMock(return_value=client)
75+
config = Config(
76+
action=CliAction.files,
77+
files_action=FilesAction.rm,
78+
rm_paths=["file1.py"],
79+
)
80+
collection.count = AsyncMock(return_value=0)
81+
client.delete_collection = AsyncMock()
82+
await rm(config)
83+
client.delete_collection.assert_called_once_with(collection.name)
84+
85+
5886
@pytest.mark.asyncio
5987
async def test_rm_no_collection(client, collection, capsys):
6088
with (

tests/test_common.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -529,7 +529,7 @@ async def test_client_manager_get_client():
529529
patch("chromadb.AsyncHttpClient") as MockAsyncHttpClient,
530530
patch("vectorcode.common.try_server", return_value=True),
531531
):
532-
mock_client = MagicMock(spec=AsyncClientAPI)
532+
mock_client = MagicMock(spec=AsyncClientAPI, parent=AsyncClientAPI)
533533
MockAsyncHttpClient.return_value = mock_client
534534

535535
async with (

0 commit comments

Comments
 (0)