gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Update tests with edge change

Browse source
This commit is contained in:
Igor Ilic 2026-01-19 16:03:40 +01:00
parent 099d78ccfc
commit 6b8ff648ad
1 changed files with 5 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cognee/tests/unit/modules/retrieval/test_brute_force_triplet_search.py
View file

@ -6,6 +6,7 @@ from cognee.modules.retrieval.utils.brute_force_triplet_search import (
get_memory_fragment, get_memory_fragment,
format_triplets, format_triplets,
) )
from cognee.modules.engine.utils.generate_edge_id import generate_edge_id
from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
from cognee.modules.graph.exceptions.exceptions import EntityNotFoundError from cognee.modules.graph.exceptions.exceptions import EntityNotFoundError
from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError
@ -1036,9 +1037,11 @@ async def test_cognee_graph_mapping_batch_shapes():
] ]
} }
edge_1_text = "relates_to"
edge_2_text = "relates_to"
edge_distances_batch = [ edge_distances_batch = [
[MockScoredResult("edge1", 0.92, payload={"text": "relates_to"})], [MockScoredResult(generate_edge_id(edge_1_text), 0.92, payload={"text": edge_1_text})],
[MockScoredResult("edge2", 0.88, payload={"text": "relates_to"})], [MockScoredResult(generate_edge_id(edge_2_text), 0.88, payload={"text": edge_2_text})],
] ]
await graph.map_vector_distances_to_graph_nodes( await graph.map_vector_distances_to_graph_nodes(