diff --git a/cognee/tests/integration/run_toy_tasks/conftest.py b/cognee/tests/integration/run_toy_tasks/conftest.py index 1c6464ae2..94d7e4070 100644 --- a/cognee/tests/integration/run_toy_tasks/conftest.py +++ b/cognee/tests/integration/run_toy_tasks/conftest.py @@ -1,9 +1,11 @@ - import os + import pytest @pytest.fixture(autouse=True, scope="session") def copy_cognee_db_to_target_location(): os.makedirs("cognee/.cognee_system/databases/") - os.system("cp cognee/tests/integration/run_toy_tasks/data/cognee_db cognee/.cognee_system/databases/cognee_db") + os.system( + "cp cognee/tests/integration/run_toy_tasks/data/cognee_db cognee/.cognee_system/databases/cognee_db" + ) diff --git a/cognee/tests/unit/documents/pdf_document.py b/cognee/tests/unit/documents/pdf_document.py index 7b51e2838..86d59a561 100644 --- a/cognee/tests/unit/documents/pdf_document.py +++ b/cognee/tests/unit/documents/pdf_document.py @@ -1,8 +1,7 @@ import os import uuid -from cognee.modules.data.processing.document_types.PdfDocument import \ - PdfDocument +from cognee.modules.data.processing.document_types.PdfDocument import PdfDocument GROUND_TRUTH = [ {"word_count": 879, "len_text": 5622, "cut_type": "sentence_end"}, diff --git a/cognee/tests/unit/integration/test_model_to_graph_to_model.py b/cognee/tests/unit/integration/test_model_to_graph_to_model.py index e3ecaf0a0..308c7ff15 100644 --- a/cognee/tests/unit/integration/test_model_to_graph_to_model.py +++ b/cognee/tests/unit/integration/test_model_to_graph_to_model.py @@ -5,8 +5,10 @@ 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) +from cognee.modules.graph.utils import ( + get_graph_from_model, + get_model_instance_from_graph, +) EDGE_GROUND_TRUTH = ( "boris", @@ -89,7 +91,7 @@ class Person(DataPoint): _metadata: dict = dict(index_fields=["name"]) -def run_test_agains_ground_truth( +def run_test_against_ground_truth( test_target_item_name, test_target_item, ground_truth_dict ): for key, ground_truth in ground_truth_dict.items(): @@ -132,11 +134,8 @@ def graph_outputs(): ) nodes, edges = get_graph_from_model(boris) - try: - car, person = nodes[0], nodes[1] - edge = edges[0] - except: - print(f"{nodes = }\n{edges = }") + car, person = nodes[0], nodes[1] + edge = edges[0] parsed_person = get_model_instance_from_graph(nodes, edges, "boris") @@ -146,12 +145,12 @@ def graph_outputs(): def test_extracted_person(graph_outputs): (_, person, _, _) = graph_outputs - run_test_agains_ground_truth("person", person, PERSON_GROUND_TRUTH) + run_test_against_ground_truth("person", person, PERSON_GROUND_TRUTH) def test_extracted_car(graph_outputs): (car, _, _, _) = graph_outputs - run_test_agains_ground_truth("car", car, CAR_GROUND_TRUTH) + run_test_against_ground_truth("car", car, CAR_GROUND_TRUTH) def test_extracted_edge(graph_outputs): @@ -166,7 +165,7 @@ def test_extracted_edge(graph_outputs): def test_parsed_person(graph_outputs): (_, _, _, parsed_person) = graph_outputs - run_test_agains_ground_truth( + run_test_against_ground_truth( "parsed_person", parsed_person, PARSED_PERSON_GROUND_TRUTH ) - run_test_agains_ground_truth("car", parsed_person.owns_car[0], CAR_GROUND_TRUTH) + run_test_against_ground_truth("car", parsed_person.owns_car[0], CAR_GROUND_TRUTH) diff --git a/cognee/tests/unit/interfaces/test_model_to_graph_to_model.py b/cognee/tests/unit/interfaces/test_model_to_graph_to_model.py new file mode 100644 index 000000000..308c7ff15 --- /dev/null +++ b/cognee/tests/unit/interfaces/test_model_to_graph_to_model.py @@ -0,0 +1,171 @@ +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, +) + +EDGE_GROUND_TRUTH = ( + "boris", + "car1", + "owns_car", + { + "source_node_id": "boris", + "target_node_id": "car1", + "relationship_name": "owns_car", + "metadata": {"type": "list"}, + }, +) + +CAR_GROUND_TRUTH = { + "id": "car1", + "brand": "Toyota", + "model": "Camry", + "year": 2020, + "color": "Blue", +} + +PERSON_GROUND_TRUTH = { + "id": "boris", + "name": "Boris", + "age": 30, + "driving_license": { + "issued_by": "PU Vrsac", + "issued_on": "2025-11-06", + "number": "1234567890", + "expires_on": "2025-11-06", + }, +} + +PARSED_PERSON_GROUND_TRUTH = { + "id": "boris", + "name": "Boris", + "age": 30, + "driving_license": { + "issued_by": "PU Vrsac", + "issued_on": "2025-11-06", + "number": "1234567890", + "expires_on": "2025-11-06", + }, +} + + +class CarTypeName(Enum): + Pickup = "Pickup" + Sedan = "Sedan" + SUV = "SUV" + Coupe = "Coupe" + Convertible = "Convertible" + Hatchback = "Hatchback" + Wagon = "Wagon" + Minivan = "Minivan" + Van = "Van" + + +class CarType(DataPoint): + id: str + name: CarTypeName + _metadata: dict = dict(index_fields=["name"]) + + +class Car(DataPoint): + id: str + brand: str + model: str + year: int + color: str + is_type: CarType + + +class Person(DataPoint): + id: str + name: str + age: int + owns_car: list[Car] + driving_license: Optional[dict] + _metadata: dict = dict(index_fields=["name"]) + + +def run_test_against_ground_truth( + test_target_item_name, test_target_item, ground_truth_dict +): + for key, ground_truth in ground_truth_dict.items(): + if isinstance(ground_truth, dict): + for key2, ground_truth2 in ground_truth.items(): + assert ( + ground_truth2 == getattr(test_target_item, key)[key2] + ), f"{test_target_item_name}/{key = }/{key2 = }: {ground_truth2 = } != {getattr(test_target_item, key)[key2] = }" + else: + assert ground_truth == getattr( + test_target_item, key + ), f"{test_target_item_name}/{key = }: {ground_truth = } != {getattr(test_target_item, key) = }" + time_delta = datetime.now(timezone.utc) - getattr(test_target_item, "updated_at") + + assert time_delta.total_seconds() < 20, f"{ time_delta.total_seconds() = }" + + +@pytest.fixture(scope="session") +def graph_outputs(): + boris = Person( + id="boris", + name="Boris", + age=30, + owns_car=[ + Car( + id="car1", + brand="Toyota", + model="Camry", + year=2020, + color="Blue", + is_type=CarType(id="sedan", name=CarTypeName.Sedan), + ) + ], + driving_license={ + "issued_by": "PU Vrsac", + "issued_on": "2025-11-06", + "number": "1234567890", + "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) + + +def test_extracted_person(graph_outputs): + (_, person, _, _) = graph_outputs + + run_test_against_ground_truth("person", person, PERSON_GROUND_TRUTH) + + +def test_extracted_car(graph_outputs): + (car, _, _, _) = graph_outputs + run_test_against_ground_truth("car", car, CAR_GROUND_TRUTH) + + +def test_extracted_edge(graph_outputs): + (_, _, edge, _) = graph_outputs + + assert ( + EDGE_GROUND_TRUTH[:3] == edge[:3] + ), f"{EDGE_GROUND_TRUTH[:3] = } != {edge[:3] = }" + for key, ground_truth in EDGE_GROUND_TRUTH[3].items(): + assert ground_truth == edge[3][key], f"{ground_truth = } != {edge[3][key] = }" + + +def test_parsed_person(graph_outputs): + (_, _, _, parsed_person) = graph_outputs + run_test_against_ground_truth( + "parsed_person", parsed_person, PARSED_PERSON_GROUND_TRUTH + ) + run_test_against_ground_truth("car", parsed_person.owns_car[0], CAR_GROUND_TRUTH)