make falkor tests optional
This commit is contained in:
prestonrasmussen
parent
885d14c9f8
commit
7a8283dbac
2 changed files with 95 additions and 50 deletions
|
|
@ -15,18 +15,27 @@ limitations under the License.
|
|||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from graphiti_core.driver.falkordb_driver import FalkorDriver, FalkorDriverSession
|
||||
from graphiti_core.helpers import DEFAULT_DATABASE
|
||||
|
||||
try:
|
||||
from graphiti_core.driver.falkordb_driver import FalkorDriver, FalkorDriverSession
|
||||
|
||||
HAS_FALKORDB = True
|
||||
except ImportError:
|
||||
FalkorDriver = None
|
||||
HAS_FALKORDB = False
|
||||
|
||||
|
||||
class TestFalkorDriver:
|
||||
"""Comprehensive test suite for FalkorDB driver."""
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.mock_client = MagicMock()
|
||||
|
|
@ -34,6 +43,7 @@ class TestFalkorDriver:
|
|||
self.driver = FalkorDriver()
|
||||
self.driver.client = self.mock_client
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_init_with_connection_params(self):
|
||||
"""Test initialization with connection parameters."""
|
||||
with patch('graphiti_core.driver.falkordb_driver.FalkorDB') as mock_falkor_db:
|
||||
|
|
@ -51,6 +61,7 @@ class TestFalkorDriver:
|
|||
password='test-pass'
|
||||
)
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_init_with_falkor_db_instance(self):
|
||||
"""Test initialization with a FalkorDB instance."""
|
||||
with patch('graphiti_core.driver.falkordb_driver.FalkorDB') as mock_falkor_db_class:
|
||||
|
|
@ -60,31 +71,35 @@ class TestFalkorDriver:
|
|||
assert driver.client is mock_falkor_db
|
||||
mock_falkor_db_class.assert_not_called()
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_provider(self):
|
||||
"""Test driver provider identification."""
|
||||
assert self.driver.provider == 'falkordb'
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_get_graph_with_name(self):
|
||||
"""Test _get_graph with specific graph name."""
|
||||
mock_graph = MagicMock()
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
result = self.driver._get_graph('test_graph')
|
||||
|
||||
|
||||
self.mock_client.select_graph.assert_called_once_with('test_graph')
|
||||
assert result is mock_graph
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_get_graph_with_none_defaults_to_default_database(self):
|
||||
"""Test _get_graph with None defaults to DEFAULT_DATABASE."""
|
||||
mock_graph = MagicMock()
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
result = self.driver._get_graph(None)
|
||||
|
||||
|
||||
self.mock_client.select_graph.assert_called_once_with(DEFAULT_DATABASE)
|
||||
assert result is mock_graph
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_execute_query_success(self):
|
||||
"""Test successful query execution."""
|
||||
mock_graph = MagicMock()
|
||||
|
|
@ -93,51 +108,54 @@ class TestFalkorDriver:
|
|||
mock_result.result_set = [['row1col1', 'row1col2']]
|
||||
mock_graph.query = AsyncMock(return_value=mock_result)
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
result = await self.driver.execute_query(
|
||||
'MATCH (n) RETURN n',
|
||||
param1='value1',
|
||||
database_='test_db'
|
||||
)
|
||||
|
||||
|
||||
self.mock_client.select_graph.assert_called_once_with('test_db')
|
||||
mock_graph.query.assert_called_once_with(
|
||||
'MATCH (n) RETURN n',
|
||||
{'param1': 'value1'}
|
||||
)
|
||||
|
||||
|
||||
result_set, header, summary = result
|
||||
assert result_set == [{'column1': 'row1col1', 'column2': 'row1col2'}]
|
||||
assert header == ['column1', 'column2']
|
||||
assert summary is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_execute_query_handles_index_already_exists_error(self):
|
||||
"""Test handling of 'already indexed' error."""
|
||||
mock_graph = MagicMock()
|
||||
mock_graph.query = AsyncMock(side_effect=Exception('Index already indexed'))
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
with patch('graphiti_core.driver.falkordb_driver.logger') as mock_logger:
|
||||
result = await self.driver.execute_query('CREATE INDEX ...')
|
||||
|
||||
|
||||
mock_logger.info.assert_called_once()
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_execute_query_propagates_other_exceptions(self):
|
||||
"""Test that other exceptions are properly propagated."""
|
||||
mock_graph = MagicMock()
|
||||
mock_graph.query = AsyncMock(side_effect=Exception('Other error'))
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
with patch('graphiti_core.driver.falkordb_driver.logger') as mock_logger:
|
||||
with pytest.raises(Exception, match='Other error'):
|
||||
await self.driver.execute_query('INVALID QUERY')
|
||||
|
||||
|
||||
mock_logger.error.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_execute_query_converts_datetime_parameters(self):
|
||||
"""Test that datetime objects in kwargs are converted to ISO strings."""
|
||||
mock_graph = MagicMock()
|
||||
|
|
@ -146,66 +164,70 @@ class TestFalkorDriver:
|
|||
mock_result.result_set = []
|
||||
mock_graph.query = AsyncMock(return_value=mock_result)
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
await self.driver.execute_query(
|
||||
'CREATE (n:Node) SET n.created_at = $created_at',
|
||||
created_at=test_datetime
|
||||
)
|
||||
|
||||
|
||||
call_args = mock_graph.query.call_args[0]
|
||||
assert call_args[1]['created_at'] == test_datetime.isoformat()
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_session_creation(self):
|
||||
"""Test session creation with specific database."""
|
||||
mock_graph = MagicMock()
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
session = self.driver.session('test_db')
|
||||
|
||||
|
||||
assert isinstance(session, FalkorDriverSession)
|
||||
assert session.graph is mock_graph
|
||||
self.mock_client.select_graph.assert_called_once_with('test_db')
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_session_creation_with_none_uses_default_database(self):
|
||||
"""Test session creation with None uses default database."""
|
||||
mock_graph = MagicMock()
|
||||
self.mock_client.select_graph.return_value = mock_graph
|
||||
|
||||
|
||||
session = self.driver.session(None)
|
||||
|
||||
|
||||
assert isinstance(session, FalkorDriverSession)
|
||||
self.mock_client.select_graph.assert_called_once_with(DEFAULT_DATABASE)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_close_calls_connection_close(self):
|
||||
"""Test driver close method calls connection close."""
|
||||
mock_connection = MagicMock()
|
||||
mock_connection.close = AsyncMock()
|
||||
self.mock_client.connection = mock_connection
|
||||
|
||||
|
||||
# Ensure hasattr checks work correctly
|
||||
del self.mock_client.aclose # Remove aclose if it exists
|
||||
|
||||
|
||||
with patch('builtins.hasattr') as mock_hasattr:
|
||||
# hasattr(self.client, 'aclose') returns False
|
||||
# hasattr(self.client.connection, 'aclose') returns False
|
||||
# hasattr(self.client.connection, 'close') returns True
|
||||
mock_hasattr.side_effect = lambda obj, attr: (
|
||||
attr == 'close' and obj is mock_connection
|
||||
attr == 'close' and obj is mock_connection
|
||||
)
|
||||
|
||||
|
||||
await self.driver.close()
|
||||
|
||||
|
||||
mock_connection.close.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_delete_all_indexes(self):
|
||||
"""Test delete_all_indexes method."""
|
||||
with patch.object(self.driver, 'execute_query', new_callable=AsyncMock) as mock_execute:
|
||||
await self.driver.delete_all_indexes('test_db')
|
||||
|
||||
|
||||
mock_execute.assert_called_once_with(
|
||||
'CALL db.indexes() YIELD name DROP INDEX name',
|
||||
database_='test_db'
|
||||
|
|
@ -215,25 +237,30 @@ class TestFalkorDriver:
|
|||
class TestFalkorDriverSession:
|
||||
"""Test FalkorDB driver session functionality."""
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.mock_graph = MagicMock()
|
||||
self.session = FalkorDriverSession(self.mock_graph)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_session_async_context_manager(self):
|
||||
"""Test session can be used as async context manager."""
|
||||
async with self.session as s:
|
||||
assert s is self.session
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_close_method(self):
|
||||
"""Test session close method doesn't raise exceptions."""
|
||||
await self.session.close() # Should not raise
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_execute_write_passes_session_and_args(self):
|
||||
"""Test execute_write method passes session and arguments correctly."""
|
||||
|
||||
async def test_func(session, *args, **kwargs):
|
||||
assert session is self.session
|
||||
assert args == ('arg1', 'arg2')
|
||||
|
|
@ -246,49 +273,52 @@ class TestFalkorDriverSession:
|
|||
assert result == 'result'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_run_single_query_with_parameters(self):
|
||||
"""Test running a single query with parameters."""
|
||||
self.mock_graph.query = AsyncMock()
|
||||
|
||||
|
||||
await self.session.run(
|
||||
'MATCH (n) RETURN n',
|
||||
param1='value1',
|
||||
param2='value2'
|
||||
)
|
||||
|
||||
|
||||
self.mock_graph.query.assert_called_once_with(
|
||||
'MATCH (n) RETURN n',
|
||||
{'param1': 'value1', 'param2': 'value2'}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_run_multiple_queries_as_list(self):
|
||||
"""Test running multiple queries passed as list."""
|
||||
self.mock_graph.query = AsyncMock()
|
||||
|
||||
|
||||
queries = [
|
||||
('MATCH (n) RETURN n', {'param1': 'value1'}),
|
||||
('CREATE (n:Node)', {'param2': 'value2'})
|
||||
]
|
||||
|
||||
|
||||
await self.session.run(queries)
|
||||
|
||||
|
||||
assert self.mock_graph.query.call_count == 2
|
||||
calls = self.mock_graph.query.call_args_list
|
||||
assert calls[0][0] == ('MATCH (n) RETURN n', {'param1': 'value1'})
|
||||
assert calls[1][0] == ('CREATE (n:Node)', {'param2': 'value2'})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_run_converts_datetime_objects_to_iso_strings(self):
|
||||
"""Test that datetime objects are converted to ISO strings."""
|
||||
self.mock_graph.query = AsyncMock()
|
||||
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
await self.session.run(
|
||||
'CREATE (n:Node) SET n.created_at = $created_at',
|
||||
created_at=test_datetime
|
||||
)
|
||||
|
||||
|
||||
self.mock_graph.query.assert_called_once()
|
||||
call_args = self.mock_graph.query.call_args[0]
|
||||
assert call_args[1]['created_at'] == test_datetime.isoformat()
|
||||
|
|
@ -297,10 +327,11 @@ class TestFalkorDriverSession:
|
|||
class TestDatetimeConversion:
|
||||
"""Test datetime conversion utility function."""
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_convert_datetime_dict(self):
|
||||
"""Test datetime conversion in nested dictionary."""
|
||||
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
|
||||
|
||||
|
||||
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
input_dict = {
|
||||
'string_val': 'test',
|
||||
|
|
@ -310,27 +341,28 @@ class TestDatetimeConversion:
|
|||
'nested_string': 'nested_test'
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
result = convert_datetimes_to_strings(input_dict)
|
||||
|
||||
|
||||
assert result['string_val'] == 'test'
|
||||
assert result['datetime_val'] == test_datetime.isoformat()
|
||||
assert result['nested_dict']['nested_datetime'] == test_datetime.isoformat()
|
||||
assert result['nested_dict']['nested_string'] == 'nested_test'
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_convert_datetime_list_and_tuple(self):
|
||||
"""Test datetime conversion in lists and tuples."""
|
||||
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
|
||||
|
||||
|
||||
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
# Test list
|
||||
input_list = ['test', test_datetime, ['nested', test_datetime]]
|
||||
result_list = convert_datetimes_to_strings(input_list)
|
||||
assert result_list[0] == 'test'
|
||||
assert result_list[1] == test_datetime.isoformat()
|
||||
assert result_list[2][1] == test_datetime.isoformat()
|
||||
|
||||
|
||||
# Test tuple
|
||||
input_tuple = ('test', test_datetime)
|
||||
result_tuple = convert_datetimes_to_strings(input_tuple)
|
||||
|
|
@ -338,18 +370,20 @@ class TestDatetimeConversion:
|
|||
assert result_tuple[0] == 'test'
|
||||
assert result_tuple[1] == test_datetime.isoformat()
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_convert_single_datetime(self):
|
||||
"""Test datetime conversion for single datetime object."""
|
||||
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
|
||||
|
||||
|
||||
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
result = convert_datetimes_to_strings(test_datetime)
|
||||
assert result == test_datetime.isoformat()
|
||||
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
def test_convert_other_types_unchanged(self):
|
||||
"""Test that non-datetime types are returned unchanged."""
|
||||
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
|
||||
|
||||
|
||||
assert convert_datetimes_to_strings('string') == 'string'
|
||||
assert convert_datetimes_to_strings(123) == 123
|
||||
assert convert_datetimes_to_strings(None) is None
|
||||
|
|
@ -362,25 +396,26 @@ class TestFalkorDriverIntegration:
|
|||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_basic_integration_with_real_falkordb(self):
|
||||
"""Basic integration test with real FalkorDB instance."""
|
||||
pytest.importorskip('falkordb')
|
||||
|
||||
|
||||
falkor_host = os.getenv('FALKORDB_HOST', 'localhost')
|
||||
falkor_port = os.getenv('FALKORDB_PORT', '6379')
|
||||
|
||||
|
||||
try:
|
||||
driver = FalkorDriver(host=falkor_host, port=falkor_port)
|
||||
|
||||
|
||||
# Test basic query execution
|
||||
result = await driver.execute_query('RETURN 1 as test')
|
||||
assert result is not None
|
||||
|
||||
|
||||
result_set, header, summary = result
|
||||
assert header == ['test']
|
||||
assert result_set == [{'test': 1}]
|
||||
|
||||
|
||||
await driver.close()
|
||||
|
||||
|
||||
except Exception as e:
|
||||
pytest.skip(f"FalkorDB not available for integration test: {e}")
|
||||
|
|
|
|||
|
|
@ -17,18 +17,26 @@ limitations under the License.
|
|||
import logging
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from graphiti_core.driver.falkordb_driver import FalkorDriver
|
||||
from graphiti_core.edges import EntityEdge, EpisodicEdge
|
||||
from graphiti_core.graphiti import Graphiti
|
||||
from graphiti_core.helpers import semaphore_gather
|
||||
from graphiti_core.nodes import EntityNode, EpisodicNode
|
||||
from graphiti_core.search.search_helpers import search_results_to_context_string
|
||||
|
||||
try:
|
||||
from graphiti_core.driver.falkordb_driver import FalkorDriver
|
||||
|
||||
HAS_FALKORDB = True
|
||||
except ImportError:
|
||||
FalkorDriver = None
|
||||
HAS_FALKORDB = False
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
|
|
@ -63,16 +71,17 @@ def setup_logging():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_graphiti_falkordb_init():
|
||||
logger = setup_logging()
|
||||
|
||||
|
||||
falkor_driver = FalkorDriver(
|
||||
host=FALKORDB_HOST,
|
||||
port=FALKORDB_PORT,
|
||||
username=FALKORDB_USER,
|
||||
password=FALKORDB_PASSWORD
|
||||
)
|
||||
|
||||
|
||||
graphiti = Graphiti(graph_driver=falkor_driver)
|
||||
|
||||
results = await graphiti.search_(query='Who is the user?')
|
||||
|
|
@ -85,6 +94,7 @@ async def test_graphiti_falkordb_init():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
|
||||
async def test_graph_falkordb_integration():
|
||||
falkor_driver = FalkorDriver(
|
||||
host=FALKORDB_HOST,
|
||||
|
|
@ -92,7 +102,7 @@ async def test_graph_falkordb_integration():
|
|||
username=FALKORDB_USER,
|
||||
password=FALKORDB_PASSWORD
|
||||
)
|
||||
|
||||
|
||||
client = Graphiti(graph_driver=falkor_driver)
|
||||
embedder = client.embedder
|
||||
driver = client.driver
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue