Clean up unit test pull request

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"
+    )

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"},

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)

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)