2d085f61bb
- Fixed hardcoded {1,3} path lengths in both edge_bfs_search and node_bfs_search
- Replaced with {1,} to use the bfs_max_depth parameter correctly
- Added tests to verify the fix works
- Resolves issue where bfs_max_depth=1 would still traverse 3 hops
Closes #772
94 lines
2.8 KiB
Python
94 lines
2.8 KiB
Python
"""Test for BFS max_depth parameter fix."""
|
|
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from graphiti_core.search.search_filters import SearchFilters
|
|
from graphiti_core.search.search_utils import edge_bfs_search, node_bfs_search
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_edge_bfs_search_uses_depth_parameter():
|
|
"""Test that edge_bfs_search uses the bfs_max_depth parameter in the query."""
|
|
# Mock driver
|
|
mock_driver = AsyncMock()
|
|
mock_driver.execute_query.return_value = ([], None, None)
|
|
|
|
# Mock search filter
|
|
search_filter = SearchFilters()
|
|
|
|
# Call edge_bfs_search with depth=2
|
|
await edge_bfs_search(
|
|
driver=mock_driver,
|
|
bfs_origin_node_uuids=['test-uuid'],
|
|
bfs_max_depth=2,
|
|
search_filter=search_filter,
|
|
group_ids=['test-group'],
|
|
limit=10,
|
|
)
|
|
|
|
# Verify the query was called
|
|
assert mock_driver.execute_query.called
|
|
call_args = mock_driver.execute_query.call_args
|
|
|
|
# Check that depth parameter is passed
|
|
assert 'depth' in call_args.kwargs
|
|
assert call_args.kwargs['depth'] == 2
|
|
|
|
# Check that the query contains the variable depth pattern
|
|
query = call_args.args[0]
|
|
assert '{1,$depth}' in query, f"Query should contain '{{1,$depth}}' but got: {query}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_bfs_search_uses_depth_parameter():
|
|
"""Test that node_bfs_search uses the bfs_max_depth parameter in the query."""
|
|
# Mock driver
|
|
mock_driver = AsyncMock()
|
|
mock_driver.execute_query.return_value = ([], None, None)
|
|
|
|
# Mock search filter
|
|
search_filter = SearchFilters()
|
|
|
|
# Call node_bfs_search with depth=1
|
|
await node_bfs_search(
|
|
driver=mock_driver,
|
|
bfs_origin_node_uuids=['test-uuid'],
|
|
search_filter=search_filter,
|
|
bfs_max_depth=1,
|
|
group_ids=['test-group'],
|
|
limit=10,
|
|
)
|
|
|
|
# Verify the query was called
|
|
assert mock_driver.execute_query.called
|
|
call_args = mock_driver.execute_query.call_args
|
|
|
|
# Check that depth parameter is passed
|
|
assert 'depth' in call_args.kwargs
|
|
assert call_args.kwargs['depth'] == 1
|
|
|
|
# Check that the query contains the variable depth pattern
|
|
query = call_args.args[0]
|
|
assert '{1,$depth}' in query, f"Query should contain '{{1,$depth}}' but got: {query}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_depth_values():
|
|
"""Test that different bfs_max_depth values are correctly passed."""
|
|
mock_driver = AsyncMock()
|
|
mock_driver.execute_query.return_value = ([], None, None)
|
|
search_filter = SearchFilters()
|
|
|
|
# Test depth=5
|
|
await edge_bfs_search(
|
|
driver=mock_driver,
|
|
bfs_origin_node_uuids=['test-uuid'],
|
|
bfs_max_depth=5,
|
|
search_filter=search_filter,
|
|
group_ids=['test-group'],
|
|
)
|
|
|
|
call_args = mock_driver.execute_query.call_args
|
|
assert call_args.kwargs['depth'] == 5
|