diff --git a/tests/helpers_test.py b/tests/helpers_test.py index 006a9cd3..e0f8b99c 100644 --- a/tests/helpers_test.py +++ b/tests/helpers_test.py @@ -16,11 +16,19 @@ limitations under the License. import pytest -from graphiti_core.helpers import lucene_sanitize - -def test_lucene_sanitize(): - # Call the function with test data +def test_neo4j_sanitize(): + """Test Neo4j sanitization by importing the driver and using its sanitize method.""" + from graphiti_core.driver.neo4j_driver import Neo4jDriver + + # Create a driver instance - if it fails, we'll handle it gracefully + try: + driver = Neo4jDriver(uri="bolt://localhost:7687", user="test", password="test") + except Exception: + # If we can't create a real driver, skip this test + pytest.skip("Neo4j driver connection failed - skipping sanitize test") + + # Test the sanitize method queries = [ ( 'This has every escape character + - && || ! ( ) { } [ ] ^ " ~ * ? : \\ /', @@ -29,9 +37,51 @@ def test_lucene_sanitize(): ('this has no escape characters', 'this has no escape characters'), ] - for query, assert_result in queries: - result = lucene_sanitize(query) - assert assert_result == result + for query, expected_result in queries: + result = driver.sanitize(query) + assert expected_result == result + + +def test_falkordb_sanitize(): + """Test FalkorDB sanitization by importing the driver and using its sanitize method.""" + try: + from graphiti_core.driver.falkordb_driver import FalkorDriver + driver = FalkorDriver() + except ImportError: + pytest.skip("FalkorDB not installed - skipping sanitize test") + except Exception: + # If we can't create a real driver, skip this test + pytest.skip("FalkorDB driver connection failed - skipping sanitize test") + + # Test the sanitize method - FalkorDB replaces special chars with spaces + queries = [ + ( + 'This has special characters: ,.<>{}[]"\':;!@#$%^&*()-+=~', + 'This has special characters', # All special chars replaced with spaces, then collapsed + ), + ('this has no special characters', 'this has no special characters'), + ('word1,word2.word3', 'word1 word2 word3'), # Separators become spaces + ('keep_underscores_intact', 'keep_underscores_intact'), # Underscores are NOT replaced + ] + + for query, expected_result in queries: + result = driver.sanitize(query) + assert expected_result == result + + +# Keep the drivers and get_driver functions for other tests that import from this file +drivers = ['neo4j', 'falkordb'] + +def get_driver(driver_name: str): + """Helper function to get a driver instance for testing""" + if driver_name == 'neo4j': + from graphiti_core.driver.neo4j_driver import Neo4jDriver + return Neo4jDriver(uri="bolt://localhost:7687", user="neo4j", password="test") + elif driver_name == 'falkordb': + from graphiti_core.driver.falkordb_driver import FalkorDriver + return FalkorDriver() + else: + raise ValueError(f"Unknown driver: {driver_name}") if __name__ == '__main__':