tests: Refactor database types and improve test coverage

Author: Davidyz

tests/subcommands/query/test_reranker.py

@@ -5,6 +5,7 @@
 import pytest
 
 from vectorcode.cli_utils import Config, QueryInclude
+from vectorcode.database.types import QueryResult
 from vectorcode.subcommands.query.reranker import (
     CrossEncoderReranker,
     NaiveReranker,
@@ -14,7 +15,6 @@
     get_available_rerankers,
     get_reranker,
 )
-from vectorcode.subcommands.query.types import QueryResult
 
 
 @pytest.fixture(scope="function")

tests/subcommands/query/test_types.py

@@ -2,7 +2,7 @@
 from tree_sitter import Point
 
 from vectorcode.chunking import Chunk
-from vectorcode.subcommands.query.types import QueryResult
+from vectorcode.database.types import QueryResult
 
 
 def make_dummy_chunk():

tests/subcommands/test_clean.py

@@ -1,82 +1,37 @@
 from unittest.mock import AsyncMock, patch
 
 import pytest
-from chromadb.api import AsyncClientAPI
 
 from vectorcode.cli_utils import Config
-from vectorcode.subcommands.clean import clean, run_clean_on_client
+from vectorcode.subcommands.clean import clean
 
 
 @pytest.mark.asyncio
-async def test_run_clean_on_client():
-    mock_client = AsyncMock(spec=AsyncClientAPI)
-    mock_collection1 = AsyncMock()
-    mock_collection1.name = "test_collection_1"
-    mock_collection1.metadata = {"path": "/test/path1"}
-    mock_collection1.count.return_value = 0  # Empty collection
-    mock_collection2 = AsyncMock()
-    mock_collection2.name = "test_collection_2"
-    mock_collection2.metadata = {"path": "/test/path2"}
-    mock_collection2.count.return_value = 1  # Non-empty collection
-
-    async def mock_get_collections(client):
-        yield mock_collection1
-        yield mock_collection2
-
-    with (
-        patch("vectorcode.subcommands.clean.get_collections", new=mock_get_collections),
-        patch("os.path.isdir", return_value=lambda x: x == "/test/path2"),
-    ):
-        await run_clean_on_client(mock_client, pipe_mode=False)
-
-    mock_client.delete_collection.assert_called_once_with(mock_collection1.name)
-
-
-@pytest.mark.asyncio
-async def test_run_clean_on_client_pipe_mode():
-    mock_client = AsyncMock(spec=AsyncClientAPI)
-    mock_collection1 = AsyncMock()
-    mock_collection1.name = "test_collection_1"
-    mock_collection1.metadata = {"path": "/test/path1"}
-    mock_collection1.count.return_value = 0  # Empty collection
-
-    async def mock_get_collections(client):
-        yield mock_collection1
+async def test_clean(capsys):
+    mock_db = AsyncMock()
+    mock_db.cleanup.return_value = ["/test/path1", "/test/path2"]
 
     with patch(
-        "vectorcode.subcommands.clean.get_collections", new=mock_get_collections
+        "vectorcode.subcommands.clean.get_database_connector", return_value=mock_db
     ):
-        await run_clean_on_client(mock_client, pipe_mode=True)
+        result = await clean(Config(pipe=False))
 
-    mock_client.delete_collection.assert_called_once_with(mock_collection1.name)
+    assert result == 0
+    captured = capsys.readouterr()
+    assert "Deleted collection: /test/path1" in captured.out
+    assert "Deleted collection: /test/path2" in captured.out
 
 
 @pytest.mark.asyncio
-async def test_run_clean_on_removed_dir():
-    mock_client = AsyncMock(spec=AsyncClientAPI)
-    mock_collection1 = AsyncMock()
-    mock_collection1.name = "test_collection_1"
-    mock_collection1.metadata = {"path": "/test/path1"}
-    mock_collection1.count.return_value = 10
-
-    async def mock_get_collections(client):
-        yield mock_collection1
+async def test_clean_pipe_mode(capsys):
+    mock_db = AsyncMock()
+    mock_db.cleanup.return_value = ["/test/path1", "/test/path2"]
 
-    with (
-        patch("vectorcode.subcommands.clean.get_collections", new=mock_get_collections),
-        patch("os.path.isdir", return_value=False),
+    with patch(
+        "vectorcode.subcommands.clean.get_database_connector", return_value=mock_db
     ):
-        await run_clean_on_client(mock_client, pipe_mode=True)
-
-    mock_client.delete_collection.assert_called_once_with(mock_collection1.name)
-
-
-@pytest.mark.asyncio
-async def test_clean():
-    AsyncMock(spec=AsyncClientAPI)
-    mock_config = Config(pipe=False)
-
-    with patch("vectorcode.subcommands.clean.ClientManager"):
-        result = await clean(mock_config)
+        result = await clean(Config(pipe=True))
 
     assert result == 0
+    captured = capsys.readouterr()
+    assert captured.out == ""

tests/subcommands/test_ls.py

@@ -1,161 +1,70 @@
 import json
-import socket
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 import tabulate
 
 from vectorcode.cli_utils import Config
-from vectorcode.subcommands.ls import get_collection_list, ls
-
-
-@pytest.mark.asyncio
-async def test_get_collection_list():
-    mock_client = AsyncMock()
-    mock_collection1 = AsyncMock()
-    mock_collection1.name = "test_collection_1"
-    mock_collection1.metadata = {
-        "path": "/test/path1",
-        "username": "test_user",
-        "embedding_function": "test_ef",
-    }
-    mock_collection1.count.return_value = 100
-    mock_collection1.get.return_value = {
-        "metadatas": [
-            {"path": "/test/path1/file1.txt"},
-            {"path": "/test/path1/file2.txt"},
-            None,
-        ]
-    }
-    mock_collection2 = AsyncMock()
-    mock_collection2.name = "test_collection_2"
-    mock_collection2.metadata = {
-        "path": "/test/path2",
-        "username": "test_user",
-        "embedding_function": "test_ef",
-    }
-    mock_collection2.count.return_value = 200
-    mock_collection2.get.return_value = {
-        "metadatas": [
-            {"path": "/test/path2/file1.txt"},
-            {"path": "/test/path2/file2.txt"},
-        ]
-    }
-
-    async def mock_get_collections(client):
-        yield mock_collection1
-        yield mock_collection2
-
-    with patch("vectorcode.subcommands.ls.get_collections", new=mock_get_collections):
-        result = await get_collection_list(mock_client)
-
-    assert len(result) == 2
-    assert result[0]["project-root"] == "/test/path1"
-    assert result[0]["user"] == "test_user"
-    assert result[0]["hostname"] == socket.gethostname()
-    assert result[0]["collection_name"] == "test_collection_1"
-    assert result[0]["size"] == 100
-    assert result[0]["embedding_function"] == "test_ef"
-    assert result[0]["num_files"] == 2
-    assert result[1]["num_files"] == 2
+from vectorcode.database.types import CollectionInfo
+from vectorcode.subcommands.ls import ls
+
+
+@pytest.fixture
+def mock_collections():
+    return [
+        CollectionInfo(
+            path="/test/path1",
+            id="test_collection_1",
+            chunk_count=100,
+            file_count=2,
+            embedding_function="test_ef",
+            database_backend="ChromaDB",
+        ),
+        CollectionInfo(
+            path="/test/path2",
+            id="test_collection_2",
+            chunk_count=200,
+            file_count=2,
+            embedding_function="test_ef",
+            database_backend="ChromaDB",
+        ),
+    ]
 
 
 @pytest.mark.asyncio
-async def test_ls_pipe_mode(capsys):
-    mock_client = AsyncMock()
-    mock_collection = AsyncMock()
-    mock_collection.name = "test_collection"
-    mock_collection.metadata = {
-        "path": "/test/path",
-        "username": "test_user",
-        "embedding_function": "test_ef",
-    }
-    mock_collection.count.return_value = 50
-    mock_collection.get.return_value = {"metadatas": [{"path": "/test/path/file.txt"}]}
-
-    async def mock_get_collections(client):
-        yield mock_collection
-
-    with (
-        patch("vectorcode.subcommands.ls.ClientManager") as MockClientManager,
-        patch(
-            "vectorcode.subcommands.ls.get_collection_list",
-            return_value=[
-                {
-                    "project-root": "/test/path",
-                    "size": 50,
-                    "num_files": 1,
-                    "embedding_function": "test_ef",
-                }
-            ],
-        ),
+async def test_ls_pipe_mode(capsys, mock_collections):
+    mock_db = AsyncMock()
+    mock_db.list_collections.return_value = mock_collections
+    with patch(
+        "vectorcode.subcommands.ls.get_database_connector", return_value=mock_db
     ):
-        mock_client = MagicMock()
-        mock_client_manager = MockClientManager.return_value
-        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
-
         config = Config(pipe=True)
         await ls(config)
         captured = capsys.readouterr()
-        expected_output = (
-            json.dumps(
-                [
-                    {
-                        "project-root": "/test/path",
-                        "size": 50,
-                        "num_files": 1,
-                        "embedding_function": "test_ef",
-                    }
-                ]
-            )
-            + "\n"
-        )
+        expected_output = json.dumps([c.to_dict() for c in mock_collections]) + "\n"
         assert captured.out == expected_output
 
 
 @pytest.mark.asyncio
-async def test_ls_table_mode(capsys, monkeypatch):
-    mock_client = AsyncMock()
-    mock_collection = AsyncMock()
-    mock_collection.name = "test_collection"
-    mock_collection.metadata = {
-        "path": "/test/path",
-        "username": "test_user",
-        "embedding_function": "test_ef",
-    }
-    mock_collection.count.return_value = 50
-    mock_collection.get.return_value = {"metadatas": [{"path": "/test/path/file.txt"}]}
-
-    async def mock_get_collections(client):
-        yield mock_collection
-
-    with (
-        patch("vectorcode.subcommands.ls.ClientManager") as MockClientManager,
-        patch(
-            "vectorcode.subcommands.ls.get_collection_list",
-            return_value=[
-                {
-                    "project-root": "/test/path",
-                    "size": 50,
-                    "num_files": 1,
-                    "embedding_function": "test_ef",
-                }
-            ],
-        ),
+async def test_ls_table_mode(capsys, mock_collections, monkeypatch):
+    mock_db = AsyncMock()
+    mock_db.list_collections.return_value = mock_collections
+    with patch(
+        "vectorcode.subcommands.ls.get_database_connector", return_value=mock_db
     ):
-        mock_client = MagicMock()
-        mock_client_manager = MockClientManager.return_value
-        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
-
         config = Config(pipe=False)
         await ls(config)
         captured = capsys.readouterr()
+        expected_table = [
+            ["/test/path1", 100, 2, "test_ef"],
+            ["/test/path2", 200, 2, "test_ef"],
+        ]
         expected_output = (
             tabulate.tabulate(
-                [["/test/path", 50, 1, "test_ef"]],
+                expected_table,
                 headers=[
                     "Project Root",
-                    "Collection Size",
+                    "Number of Embeddings",
                     "Number of Files",
                     "Embedding Function",
                 ],
@@ -166,32 +75,22 @@ async def mock_get_collections(client):
 
     # Test with HOME environment variable set
     monkeypatch.setenv("HOME", "/test")
-    with (
-        patch("vectorcode.subcommands.ls.ClientManager") as MockClientManager,
-        patch(
-            "vectorcode.subcommands.ls.get_collection_list",
-            return_value=[
-                {
-                    "project-root": "/test/path",
-                    "size": 50,
-                    "num_files": 1,
-                    "embedding_function": "test_ef",
-                }
-            ],
-        ),
+    with patch(
+        "vectorcode.subcommands.ls.get_database_connector", return_value=mock_db
     ):
-        mock_client = MagicMock()
-        mock_client_manager = MockClientManager.return_value
-        mock_client_manager._create_client = AsyncMock(return_value=mock_client)
         config = Config(pipe=False)
         await ls(config)
         captured = capsys.readouterr()
+        expected_table = [
+            ["~/path1", 100, 2, "test_ef"],
+            ["~/path2", 200, 2, "test_ef"],
+        ]
         expected_output = (
             tabulate.tabulate(
-                [["~/path", 50, 1, "test_ef"]],
+                expected_table,
                 headers=[
                     "Project Root",
-                    "Collection Size",
+                    "Number of Embeddings",
                     "Number of Files",
                     "Embedding Function",
                 ],