97 lines
4.2 KiB
Python
97 lines
4.2 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, AsyncMock, patch
|
|
import os
|
|
import sys
|
|
|
|
# Add project root to path
|
|
sys.path.append(os.getcwd())
|
|
|
|
from src.knowledge.graph import GraphDatabase
|
|
from src import config
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_txt_add_vector_entity_parsing():
|
|
# Mock driver and session
|
|
mock_driver = MagicMock()
|
|
mock_session = MagicMock()
|
|
mock_driver.session.return_value.__enter__.return_value = mock_session
|
|
|
|
# Setup mock transaction
|
|
mock_tx = MagicMock()
|
|
def side_effect_execute_write(func, *args, **kwargs):
|
|
return func(mock_tx, *args, **kwargs)
|
|
|
|
# Mock execute_read to return empty list (no missing embeddings for this test)
|
|
# The code calls _get_nodes_without_embedding which returns [record['name']]
|
|
# If we return [], it means all nodes have embeddings or none found.
|
|
# Actually, the code checks:
|
|
# nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities)
|
|
# Let's mock it to return empty list so we skip embedding generation loop which simplifies test
|
|
mock_session.execute_read.return_value = []
|
|
|
|
mock_session.execute_write.side_effect = side_effect_execute_write
|
|
|
|
# Mock embedding model
|
|
with patch('src.knowledge.graph.select_embedding_model') as mock_select_model:
|
|
mock_embed_model = MagicMock()
|
|
mock_select_model.return_value = mock_embed_model
|
|
|
|
# Instantiate GraphDatabase with mocked driver
|
|
# We also need to patch GD.driver in the init
|
|
with patch('src.knowledge.graph.GD.driver', return_value=mock_driver):
|
|
gd = GraphDatabase()
|
|
# Manually set driver and status just in case init didn't work as expected due to other mocks
|
|
gd.driver = mock_driver
|
|
gd.status = "open"
|
|
gd.embed_model_name = "test_model" # avoid config check issues if possible
|
|
|
|
# Mock config to match
|
|
with patch('src.knowledge.graph.config.embed_model', "test_model"):
|
|
with patch('src.knowledge.graph.config.embed_model_names', {"test_model": MagicMock(dimension=1024)}):
|
|
|
|
# Test data: Mixed format
|
|
triples = [
|
|
# Legacy format
|
|
{"h": "A", "r": "KNOWS", "t": "B"},
|
|
# Extended format
|
|
{
|
|
"h": {"name": "C", "age": 30},
|
|
"r": {"type": "LIKES", "weight": 0.8},
|
|
"t": {"name": "D", "role": "User"}
|
|
}
|
|
]
|
|
|
|
# Run the method
|
|
await gd.txt_add_vector_entity(triples)
|
|
|
|
# Verify calls to mock_tx.run
|
|
merge_calls = []
|
|
for call in mock_tx.run.call_args_list:
|
|
args, kwargs = call
|
|
query = args[0] if args else kwargs.get('query', '')
|
|
if "MERGE (h:Entity:Upload" in query:
|
|
# The args are passed as kwargs to run: h_name=..., etc.
|
|
merge_calls.append(kwargs)
|
|
|
|
assert len(merge_calls) == 2, f"Expected 2 merge calls, got {len(merge_calls)}"
|
|
|
|
# Call 1 (Legacy)
|
|
call1 = merge_calls[0]
|
|
assert call1['h_name'] == "A"
|
|
assert call1['h_props'] == {}
|
|
assert call1['t_name'] == "B"
|
|
assert call1['t_props'] == {}
|
|
assert call1['r_type'] == "KNOWS"
|
|
assert call1['r_props'] == {}
|
|
|
|
# Call 2 (Extended)
|
|
call2 = merge_calls[1]
|
|
assert call2['h_name'] == "C"
|
|
assert call2['h_props'] == {'age': 30}
|
|
assert call2['t_name'] == "D"
|
|
assert call2['t_props'] == {'role': 'User'}
|
|
assert call2['r_type'] == "LIKES"
|
|
assert call2['r_props'] == {'weight': 0.8}
|
|
|
|
print("Verification passed!")
|