From aa1480ca2cedb0049511f5b95dbafa9f71b67d94 Mon Sep 17 00:00:00 2001
From: Leon Luithlen <leon@topoteretes.com>
Date: Wed, 13 Nov 2024 14:26:05 +0100
Subject: [PATCH] Remove unused imports and make PdfDocument_test a pytest
 function

---
 cognee/tests/unit/documents/PdfDocument_test.py            | 2 +-
 .../unit/interfaces/graph/get_graph_from_model_test.py     | 5 -----
 cognee/tests/unit/interfaces/graph/util.py                 | 7 ++-----
 3 files changed, 3 insertions(+), 11 deletions(-)

diff --git a/cognee/tests/unit/documents/PdfDocument_test.py b/cognee/tests/unit/documents/PdfDocument_test.py
index 86d59a561..917e9c3e0 100644
--- a/cognee/tests/unit/documents/PdfDocument_test.py
+++ b/cognee/tests/unit/documents/PdfDocument_test.py
@@ -9,7 +9,7 @@ GROUND_TRUTH = [
 ]
 
 
-if __name__ == "__main__":
+def test_PdfDocument():
     test_file_path = os.path.join(
         os.sep,
         *(os.path.dirname(__file__).split(os.sep)[:-2]),
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 90423cc33..17dd69a0e 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,8 +1,3 @@
-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 = (
diff --git a/cognee/tests/unit/interfaces/graph/util.py b/cognee/tests/unit/interfaces/graph/util.py
index b7598cc94..764eafa11 100644
--- a/cognee/tests/unit/interfaces/graph/util.py
+++ b/cognee/tests/unit/interfaces/graph/util.py
@@ -1,12 +1,9 @@
 from datetime import datetime, timezone
+from typing import Any, Dict
 
 
-from typing import Dict, Any
-
 def run_test_against_ground_truth(
-    test_target_item_name: str,
-    test_target_item: Any,
-    ground_truth_dict: Dict[str, Any]
+    test_target_item_name: str, test_target_item: Any, ground_truth_dict: Dict[str, Any]
 ):
     """Validates test target item attributes against ground truth values.