setup pytest and intergation tests for tools

cognee-mcp/pyproject.toml

@@ -32,6 +32,11 @@ packages = ["src"]
 [dependency-groups]
 dev = [
     "debugpy>=1.8.12,<2.0.0",
+    "pytest>=8.0.0,<9.0.0",
+    "pytest-asyncio>=0.24.0,<1.0.0",
+    "pytest-mock>=3.14.0,<4.0.0",
+    "pytest-cov>=6.0.0,<7.0.0",
+    "ruff>=0.14.3",
 ]
 
 [tool.hatch.metadata]

cognee-mcp/pytest.ini

@@ -0,0 +1,8 @@
+[pytest]
+asyncio_mode = auto
+asyncio_default_fixture_loop_scope = function
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*
+

cognee-mcp/tests/conftest.py

@@ -0,0 +1,23 @@
+import pathlib
+import pytest
+import cognee
+from src import server
+from src.cognee_client import CogneeClient
+
+
+@pytest.fixture(autouse=True)
+async def setup_isolated_cognee(request):
+    test_name = request.node.name
+    test_base = pathlib.Path(request.fspath).parent
+
+    cognee.config.data_root_directory(str(test_base / f".data_storage/{test_name}"))
+    cognee.config.system_root_directory(str(test_base / f".cognee_system/{test_name}"))
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    server.cognee_client = CogneeClient(api_url=None, api_token=None)
+
+    yield
+
+    server.cognee_client = None

cognee-mcp/tests/tools/test_basic_tools.py

@@ -0,0 +1,31 @@
+"""
+Tests for basic MCP tools: prune, cognify_status
+
+These are integration tests that test the actual tool behavior.
+Run with: pytest tests/tools/test_basic_tools.py -v
+"""
+
+import pytest
+import mcp.types as types
+
+from src import server
+
+
+@pytest.mark.asyncio
+async def test_prune():
+    """Test prune tool - removes all data from knowledge graph"""
+    result = await server.prune()
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert "Pruned" in result[0].text or "not available" in result[0].text
+
+
+@pytest.mark.asyncio
+async def test_cognify_status():
+    """Test cognify_status tool - gets status of cognify pipeline"""
+    result = await server.cognify_status()
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert len(result[0].text) > 0

cognee-mcp/tests/tools/test_cognify.py

@@ -0,0 +1,21 @@
+"""
+Test for cognify tool
+
+These are integration tests that test the actual tool behavior.
+Run with: pytest tests/tools/test_cognify.py -v
+"""
+
+import pytest
+import mcp.types as types
+
+from src import server
+
+
+@pytest.mark.asyncio
+async def test_cognify():
+    """Test cognify tool - launches background task to process data"""
+    result = await server.cognify(data="Test data for cognify")
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert "Background process" in result[0].text or "launched" in result[0].text

cognee-mcp/tests/tools/test_list_data.py

@@ -0,0 +1,38 @@
+"""
+Tests for data management tools: list_data, delete
+
+These are integration tests that test the actual tool behavior.
+Run with: pytest tests/tools/test_data_tools.py -v
+"""
+
+import pytest
+import mcp.types as types
+import cognee
+
+from src import server
+
+
+@pytest.mark.asyncio
+async def test_list_data_when_empty():
+    """Test list_data tool - lists all datasets"""
+    result = await server.list_data()
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert (
+        "❌ Failed to list data: DatabaseNotCreatedError: The database has not been created yet"
+        in result[0].text
+    )
+
+
+@pytest.mark.asyncio
+async def test_list_data_when_cognified():
+    """Test list_data tool - lists all datasets"""
+    await cognee.add("Every node tells a story. Most of mine are unit tests.")
+    await cognee.cognify()
+
+    result = await server.list_data()
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert all(s in result[0].text for s in ["📂 Available Datasets", "Dataset ID: "])

cognee-mcp/tests/tools/test_search.py

@@ -0,0 +1,92 @@
+"""
+Test for search tool
+
+These are integration tests that test the actual tool behavior.
+Run with: pytest tests/tools/test_search.py -v
+"""
+
+import cognee
+from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
+import pytest
+import mcp.types as types
+from cognee import SearchType
+
+from src import server
+
+
+@pytest.mark.asyncio
+async def test_search_handles_database_not_ready():
+    """Test search tool handles database not ready scenario gracefully"""
+    with pytest.raises(DatabaseNotCreatedError):
+        await server.search(search_query="test query", search_type="GRAPH_COMPLETION")
+
+
+@pytest.mark.asyncio
+async def test_search_graph_completion():
+    """Test search with GRAPH_COMPLETION type"""
+    await cognee.add("Artificial intelligence and machine learning are transforming technology.")
+    await cognee.cognify()
+
+    result = await server.search(
+        search_query="What is AI?", search_type=SearchType.GRAPH_COMPLETION.value
+    )
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert result[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_search_rag_completion():
+    """Test search with RAG_COMPLETION type"""
+    await cognee.add("Python is a programming language that emphasizes readability.")
+    await cognee.cognify()
+
+    result = await server.search(
+        search_query="What is Python?", search_type=SearchType.RAG_COMPLETION.value
+    )
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert result[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_search_chunks():
+    """Test search with CHUNKS type"""
+    await cognee.add("JavaScript is the language of the web browser.")
+    await cognee.cognify()
+
+    result = await server.search(search_query="web browser", search_type=SearchType.CHUNKS.value)
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert result[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_search_summaries():
+    """Test search with SUMMARIES type"""
+    await cognee.add("Database systems manage and store structured data efficiently.")
+    await cognee.cognify()
+
+    result = await server.search(search_query="database", search_type=SearchType.SUMMARIES.value)
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert result[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_search_feeling_lucky():
+    """Test search with FEELING_LUCKY type"""
+    await cognee.add("Machine learning models learn patterns from data.")
+    await cognee.cognify()
+
+    result = await server.search(
+        search_query="learning", search_type=SearchType.FEELING_LUCKY.value
+    )
+
+    assert len(result) == 1
+    assert isinstance(result[0], types.TextContent)
+    assert result[0].text is not None