From 03934fdf11598d41d09cf9333fcc267783f0b0e9 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Thu, 13 Nov 2025 17:06:27 +0100 Subject: [PATCH] fix: make tests work with neo4j and fix get_graph_data in neo4j --- .../databases/graph/neo4j_driver/adapter.py | 17 ++++++++++++++++- cognee/tests/test_delete_default_graph.py | 4 +++- .../test_delete_default_graph_non_mocked.py | 4 +++- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py index ae2dd0edd..6ff8f7645 100644 --- a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py +++ b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py @@ -955,7 +955,22 @@ class Neo4jAdapter(GraphDBInterface): f"Retrieved {len(nodes)} nodes and {len(edges)} edges in {retrieval_time:.2f} seconds" ) - return (nodes, edges) + return ( + [ + ( + node_id, + { + **node_data, + "id": node_id, + "metadata": json.loads(node_data["metadata"]) + if "metadata" in node_data + else {}, + }, + ) + for (node_id, node_data) in nodes + ], + edges, + ) except Exception as e: logger.error(f"Error during graph data retrieval: {str(e)}") diff --git a/cognee/tests/test_delete_default_graph.py b/cognee/tests/test_delete_default_graph.py index a35b8a640..d6236500c 100644 --- a/cognee/tests/test_delete_default_graph.py +++ b/cognee/tests/test_delete_default_graph.py @@ -151,7 +151,9 @@ async def main(mock_create_structured_output: AsyncMock): nodes, edges = await graph_engine.get_graph_data() assert len(nodes) == 9 and len(edges) == 10, "Nodes and edges are not deleted." assert not any( - node[1]["name"] == "john" or node[1]["name"] == "food for hungry" for node in nodes + node[1]["name"] == "john" or node[1]["name"] == "food for hungry" + for node in nodes + if "name" in node[1] ), "Nodes are not deleted." after_first_delete_node_ids = set([node[0] for node in nodes]) diff --git a/cognee/tests/test_delete_default_graph_non_mocked.py b/cognee/tests/test_delete_default_graph_non_mocked.py index ce5087347..a1fe68f54 100644 --- a/cognee/tests/test_delete_default_graph_non_mocked.py +++ b/cognee/tests/test_delete_default_graph_non_mocked.py @@ -73,7 +73,9 @@ async def main(): "Nodes and edges are not deleted." ) assert not any( - node[1]["name"] == "john" or node[1]["name"] == "food for hungry" for node in nodes + node[1]["name"] == "john" or node[1]["name"] == "food for hungry" + for node in nodes + if "name" in node[1] ), "Nodes are not deleted." after_first_delete_node_ids = set([node[0] for node in nodes])