feat(chroma0): Raise CollectionNotFoundError on missing collection

Author: Davidyz

src/vectorcode/database/chroma0.py

@@ -507,11 +507,15 @@ async def delete(self) -> int:
         return len(rm_paths)
 
     async def drop(self, *, collection_id=None, collection_path=None):
-        collection_id = collection_id or get_collection_id(
-            str(collection_path or self._configs.project_root)
-        )
+        project_root = str(collection_path or self._configs.project_root)
+        collection_id = collection_id or get_collection_id(project_root)
         async with _Chroma0ClientManager().get_client(self._configs) as client:
-            await client.delete_collection(collection_id)
+            try:
+                await client.delete_collection(collection_id)
+            except ValueError as e:
+                raise CollectionNotFoundError(
+                    f"Collection at {project_root} is not found."
+                ) from e
 
     async def get_chunks(self, file_path) -> list[Chunk]:
         file_path = os.path.abspath(file_path)

tests/database/test_chroma0.py

@@ -220,6 +220,26 @@ async def test_drop(mock_config):
         mock_client.delete_collection.assert_called_once_with("collection_id")
 
 
+@pytest.mark.asyncio
+async def test_drop_invalid_collection(mock_config):
+    """Test the drop method."""
+    connector = ChromaDB0Connector(mock_config)
+    with (
+        patch(
+            "vectorcode.database.chroma0._Chroma0ClientManager.get_client"
+        ) as mock_get_client,
+        patch(
+            "vectorcode.database.chroma0.get_collection_id",
+            return_value="collection_id",
+        ),
+    ):
+        mock_client = AsyncMock()
+        mock_get_client.return_value.__aenter__.return_value = mock_client
+        mock_client.delete_collection.side_effect = ValueError
+        with pytest.raises(CollectionNotFoundError):
+            await connector.drop()
+
+
 @pytest.mark.asyncio
 async def test_get_chunks(mock_config):
     """Test the get_chunks method."""