diff --git a/cognee/tests/integration/run_toy_tasks/conftest.py b/cognee/tests/integration/run_toy_tasks/conftest.py index 94d7e4070..09e98328c 100644 --- a/cognee/tests/integration/run_toy_tasks/conftest.py +++ b/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" ) diff --git a/cognee/tests/integration/run_toy_tasks/test_run_tasks_from_queue.py b/cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py similarity index 100% rename from cognee/tests/integration/run_toy_tasks/test_run_tasks_from_queue.py rename to cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py diff --git a/cognee/tests/integration/run_toy_tasks/test_run_tasks.py b/cognee/tests/integration/run_toy_tasks/run_tasks_test.py similarity index 100% rename from cognee/tests/integration/run_toy_tasks/test_run_tasks.py rename to cognee/tests/integration/run_toy_tasks/run_tasks_test.py diff --git a/cognee/tests/unit/documents/pdf_document.py b/cognee/tests/unit/documents/PdfDocument_test.py similarity index 97% rename from cognee/tests/unit/documents/pdf_document.py rename to cognee/tests/unit/documents/PdfDocument_test.py index 86d59a561..7b51e2838 100644 --- a/cognee/tests/unit/documents/pdf_document.py +++ b/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"}, diff --git a/cognee/tests/unit/interfaces/test_model_to_graph_to_model.py b/cognee/tests/unit/interfaces/graph/conftest.py similarity index 50% rename from cognee/tests/unit/interfaces/test_model_to_graph_to_model.py rename to cognee/tests/unit/interfaces/graph/conftest.py index 308c7ff15..3f70abea2 100644 --- a/cognee/tests/unit/interfaces/test_model_to_graph_to_model.py +++ b/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) 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 new file mode 100644 index 000000000..2311f857c --- /dev/null +++ b/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] = }" 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 new file mode 100644 index 000000000..be24719d5 --- /dev/null +++ b/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) diff --git a/cognee/tests/unit/interfaces/graph/util.py b/cognee/tests/unit/interfaces/graph/util.py new file mode 100644 index 000000000..a9d351087 --- /dev/null +++ b/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() = }" diff --git a/cognee/tests/unit/processing/test_chunking.py b/cognee/tests/unit/processing/chunks/chunk_by_paragraph_test.py similarity index 100% rename from cognee/tests/unit/processing/test_chunking.py rename to cognee/tests/unit/processing/chunks/chunk_by_paragraph_test.py