From 0ea011ccd7bec863822d27bad4d1f6df8ab5d1d7 Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Fri, 15 Nov 2024 10:27:27 +0100 Subject: [PATCH] Adapt graph interfaces tests to debugged get_graph_from_model --- .../tests/unit/interfaces/graph/conftest.py | 18 +----- .../graph/get_graph_from_model_test.py | 59 +++++++++++++++---- .../get_model_instance_from_graph_test.py | 10 +++- cognee/tests/unit/interfaces/graph/util.py | 2 + 4 files changed, 60 insertions(+), 29 deletions(-) diff --git a/cognee/tests/unit/interfaces/graph/conftest.py b/cognee/tests/unit/interfaces/graph/conftest.py index 9a784bb53..45f977bd6 100644 --- a/cognee/tests/unit/interfaces/graph/conftest.py +++ b/cognee/tests/unit/interfaces/graph/conftest.py @@ -1,14 +1,9 @@ -from datetime import datetime, timezone from enum import Enum from typing import Optional import pytest from cognee.infrastructure.engine import DataPoint -from cognee.modules.graph.utils import ( - get_graph_from_model, - get_model_instance_from_graph, -) class CarTypeName(Enum): @@ -47,8 +42,8 @@ class Person(DataPoint): _metadata: dict = dict(index_fields=["name"]) -@pytest.fixture(scope="session") -def graph_outputs(): +@pytest.fixture(scope="function") +def boris(): boris = Person( id="boris", name="Boris", @@ -70,11 +65,4 @@ def graph_outputs(): "expires_on": "2025-11-06", }, ) - nodes, edges = get_graph_from_model(boris) - - car, person = nodes[0], nodes[1] - edge = edges[0] - - parsed_person = get_model_instance_from_graph(nodes, edges, "boris") - - return (car, person, edge, parsed_person) + return boris diff --git a/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py b/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py index 17dd69a0e..e56a2dff2 100644 --- a/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py +++ b/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py @@ -1,6 +1,19 @@ +from cognee.modules.graph.utils import get_graph_from_model from cognee.tests.unit.interfaces.graph.util import run_test_against_ground_truth -EDGE_GROUND_TRUTH = ( +CAR_SEDAN_EDGE = ( + "car1", + "sedan", + "is_type", + { + "source_node_id": "car1", + "target_node_id": "sedan", + "relationship_name": "is_type", + }, +) + + +BORIS_CAR_EDGE_GROUND_TRUTH = ( "boris", "car1", "owns_car", @@ -12,6 +25,8 @@ EDGE_GROUND_TRUTH = ( }, ) +CAR_TYPE_GROUND_TRUTH = {"id": "sedan"} + CAR_GROUND_TRUTH = { "id": "car1", "brand": "Toyota", @@ -33,22 +48,42 @@ PERSON_GROUND_TRUTH = { } -def test_extracted_person(graph_outputs): - (_, person, _, _) = graph_outputs - - run_test_against_ground_truth("person", person, PERSON_GROUND_TRUTH) +def test_extracted_car_type(boris): + nodes, _ = get_graph_from_model(boris) + assert len(nodes) == 3 + car_type = nodes[0] + run_test_against_ground_truth("car_type", car_type, CAR_TYPE_GROUND_TRUTH) -def test_extracted_car(graph_outputs): - (car, _, _, _) = graph_outputs +def test_extracted_car(boris): + nodes, _ = get_graph_from_model(boris) + assert len(nodes) == 3 + car = nodes[1] run_test_against_ground_truth("car", car, CAR_GROUND_TRUTH) -def test_extracted_edge(graph_outputs): - (_, _, edge, _) = graph_outputs +def test_extracted_person(boris): + nodes, _ = get_graph_from_model(boris) + assert len(nodes) == 3 + person = nodes[2] + run_test_against_ground_truth("person", person, PERSON_GROUND_TRUTH) + + +def test_extracted_car_sedan_edge(boris): + _, edges = get_graph_from_model(boris) + edge = edges[0] + + assert CAR_SEDAN_EDGE[:3] == edge[:3], f"{CAR_SEDAN_EDGE[:3] = } != {edge[:3] = }" + for key, ground_truth in CAR_SEDAN_EDGE[3].items(): + assert ground_truth == edge[3][key], f"{ground_truth = } != {edge[3][key] = }" + + +def test_extracted_boris_car_edge(boris): + _, edges = get_graph_from_model(boris) + edge = edges[1] assert ( - EDGE_GROUND_TRUTH[:3] == edge[:3] - ), f"{EDGE_GROUND_TRUTH[:3] = } != {edge[:3] = }" - for key, ground_truth in EDGE_GROUND_TRUTH[3].items(): + BORIS_CAR_EDGE_GROUND_TRUTH[:3] == edge[:3] + ), f"{BORIS_CAR_EDGE_GROUND_TRUTH[:3] = } != {edge[:3] = }" + for key, ground_truth in BORIS_CAR_EDGE_GROUND_TRUTH[3].items(): assert ground_truth == edge[3][key], f"{ground_truth = } != {edge[3][key] = }" diff --git a/cognee/tests/unit/interfaces/graph/get_model_instance_from_graph_test.py b/cognee/tests/unit/interfaces/graph/get_model_instance_from_graph_test.py index 98ba501bd..f1aa7736d 100644 --- a/cognee/tests/unit/interfaces/graph/get_model_instance_from_graph_test.py +++ b/cognee/tests/unit/interfaces/graph/get_model_instance_from_graph_test.py @@ -1,3 +1,7 @@ +from cognee.modules.graph.utils import ( + get_graph_from_model, + get_model_instance_from_graph, +) from cognee.tests.unit.interfaces.graph.util import run_test_against_ground_truth PARSED_PERSON_GROUND_TRUTH = { @@ -21,8 +25,10 @@ CAR_GROUND_TRUTH = { } -def test_parsed_person(graph_outputs): - (_, _, _, parsed_person) = graph_outputs +def test_parsed_person(boris): + nodes, edges = get_graph_from_model(boris) + parsed_person = get_model_instance_from_graph(nodes, edges, "boris") + run_test_against_ground_truth( "parsed_person", parsed_person, PARSED_PERSON_GROUND_TRUTH ) diff --git a/cognee/tests/unit/interfaces/graph/util.py b/cognee/tests/unit/interfaces/graph/util.py index 764eafa11..e5da0201c 100644 --- a/cognee/tests/unit/interfaces/graph/util.py +++ b/cognee/tests/unit/interfaces/graph/util.py @@ -21,6 +21,8 @@ def run_test_against_ground_truth( assert ( ground_truth2 == getattr(test_target_item, key)[key2] ), f"{test_target_item_name}/{key = }/{key2 = }: {ground_truth2 = } != {getattr(test_target_item, key)[key2] = }" + elif isinstance(ground_truth, list): + raise NotImplementedError("Currently not implemented for 'list'") else: assert ground_truth == getattr( test_target_item, key