Apply _test.py suffix to test files

cognee/tests/integration/run_toy_tasks/conftest.py

@@ -5,7 +5,7 @@ import pytest
 
 @pytest.fixture(autouse=True, scope="session")
 def copy_cognee_db_to_target_location():
-    os.makedirs("cognee/.cognee_system/databases/")
+    os.makedirs("cognee/.cognee_system/databases/", exist_ok=True)
     os.system(
         "cp cognee/tests/integration/run_toy_tasks/data/cognee_db cognee/.cognee_system/databases/cognee_db"
     )

cognee/tests/unit/documents/PdfDocument_test.py

@@ -1,7 +1,8 @@
 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/interfaces/graph/conftest.py

@@ -5,10 +5,8 @@ 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",
@@ -42,18 +40,6 @@ PERSON_GROUND_TRUTH = {
     },
 }
 
-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"
@@ -91,24 +77,6 @@ class Person(DataPoint):
     _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(
@@ -140,32 +108,3 @@ def graph_outputs():
     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)

cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py

@@ -0,0 +1,58 @@
+from cognee.infrastructure.engine import DataPoint
+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
+
+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",
+    },
+}
+
+
+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] = }"

cognee/tests/unit/interfaces/graph/get_model_instance_from_graph_test.py

@@ -0,0 +1,30 @@
+from cognee.tests.unit.interfaces.graph.util import \
+    run_test_against_ground_truth
+
+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",
+    },
+}
+
+CAR_GROUND_TRUTH = {
+    "id": "car1",
+    "brand": "Toyota",
+    "model": "Camry",
+    "year": 2020,
+    "color": "Blue",
+}
+
+
+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)

cognee/tests/unit/interfaces/graph/util.py

@@ -0,0 +1,19 @@
+from datetime import datetime, timezone
+
+
+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() = }"