File size: 3,847 Bytes
979853c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import pytest
from lightrag.kg.memgraph_impl import MemgraphStorage
pytestmark = pytest.mark.offline
class _FakeNode(dict):
def __init__(self, node_id: int, entity_id: str, **properties):
super().__init__(entity_id=entity_id, **properties)
self.id = node_id
class _FakeResult:
def __init__(self, record):
self._record = record
async def single(self):
return self._record
async def consume(self):
return None
class _FakeSession:
def __init__(self, record, calls):
self._record = record
self._calls = calls
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def run(self, query, parameters=None, **kwargs):
if parameters is None:
parameters = kwargs
self._calls.append((query, parameters))
return _FakeResult(self._record)
class _FakeDriver:
def __init__(self, record, calls):
self._record = record
self._calls = calls
def session(self, **kwargs):
return _FakeSession(self._record, self._calls)
def _make_storage(record):
calls = []
storage = MemgraphStorage(
namespace="chunk_entity_relation",
global_config={"max_graph_nodes": 1000},
embedding_func=None,
workspace="test",
)
storage._driver = _FakeDriver(record, calls)
storage._DATABASE = "memgraph"
return storage, calls
@pytest.mark.asyncio
async def test_get_knowledge_graph_preserves_isolated_start_node():
start_node = _FakeNode(1, "Start", description="isolated")
storage, calls = _make_storage(
{
"node_info": [{"node": start_node}],
"relationships": [],
"is_truncated": False,
}
)
result = await storage.get_knowledge_graph("Start", max_depth=0, max_nodes=1)
# Verify result data: isolated node must appear with correct labels and properties
assert len(result.nodes) == 1
assert result.nodes[0].labels == ["Start"]
assert result.nodes[0].properties["entity_id"] == "Start"
assert result.edges == []
assert result.is_truncated is False
# Verify query parameters: max_other_nodes must reserve a slot for the start node
assert len(calls) == 1
_, params = calls[0]
assert params["entity_id"] == "Start"
assert params["max_nodes"] == 1
assert (
params["max_other_nodes"] == 0
) # max_nodes - 1 = 0, start node occupies the only slot
@pytest.mark.asyncio
async def test_get_knowledge_graph_reserves_capacity_for_start_node_when_truncating():
start_node = _FakeNode(1, "Start")
storage, calls = _make_storage(
{
"node_info": [{"node": start_node}],
"relationships": [],
"is_truncated": True,
}
)
result = await storage.get_knowledge_graph("Start", max_depth=2, max_nodes=2)
# Verify truncation is reflected in result
assert result.is_truncated is True
assert len(result.nodes) == 1
assert result.edges == []
# Verify max_other_nodes leaves exactly one slot for the start node
assert len(calls) == 1
_, params = calls[0]
assert params["max_nodes"] == 2
assert (
params["max_other_nodes"] == 1
) # max_nodes - 1 = 1, start node always included
@pytest.mark.asyncio
async def test_get_knowledge_graph_max_nodes_zero_does_not_underflow():
"""max_other_nodes must not go negative when max_nodes=0."""
storage, calls = _make_storage(
{
"node_info": [],
"relationships": [],
"is_truncated": False,
}
)
await storage.get_knowledge_graph("Start", max_depth=1, max_nodes=0)
_, params = calls[0]
assert params["max_other_nodes"] == 0 # max(0 - 1, 0) = 0, no underflow
|