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.