diff --git a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py
index d7dfcec15..7f7493ed1 100644
--- a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py
+++ b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py
@@ -16,7 +16,7 @@ def test_node_initialization():
 
 def test_node_invalid_dimension():
     """Test that initializing a Node with a non-positive dimension raises an error."""
-    with pytest.raises(ValueError, match="Dimension must be a positive integer"):
+    with pytest.raises(InvalidValueError, match="Dimension must be a positive integer"):
         Node("node1", dimension=0)
 
 
@@ -69,7 +69,7 @@ def test_is_node_alive_in_dimension():
 def test_node_alive_invalid_dimension():
     """Test that checking alive status with an invalid dimension raises an error."""
     node = Node("node1", dimension=1)
-    with pytest.raises(ValueError, match="Dimension 1 is out of range"):
+    with pytest.raises(InvalidValueError, match="Dimension 1 is out of range"):
         node.is_node_alive_in_dimension(1)
 
 
diff --git a/cognee/tests/unit/modules/graph/cognee_graph_test.py b/cognee/tests/unit/modules/graph/cognee_graph_test.py
index d77bab8ad..bdcd9725c 100644
--- a/cognee/tests/unit/modules/graph/cognee_graph_test.py
+++ b/cognee/tests/unit/modules/graph/cognee_graph_test.py
@@ -1,6 +1,6 @@
 import pytest
 
-from cognee.exceptions import InvalidValueError
+from cognee.exceptions import EntityNotFoundError, EntityAlreadyExistsError
 from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node
 
@@ -24,7 +24,7 @@ def test_add_duplicate_node(setup_graph):
     graph = setup_graph
     node = Node("node1")
     graph.add_node(node)
-    with pytest.raises(InvalidValueError, match="Node with id node1 already exists."):
+    with pytest.raises(EntityAlreadyExistsError, match="Node with id node1 already exists."):
         graph.add_node(node)
 
 
@@ -51,7 +51,7 @@ def test_add_duplicate_edge(setup_graph):
     graph.add_node(node2)
     edge = Edge(node1, node2)
     graph.add_edge(edge)
-    with pytest.raises(InvalidValueError, match="Edge .* already exists in the graph."):
+    with pytest.raises(EntityAlreadyExistsError, match="Edge .* already exists in the graph."):
         graph.add_edge(edge)
 
 
@@ -84,5 +84,5 @@ def test_get_edges_success(setup_graph):
 def test_get_edges_nonexistent_node(setup_graph):
     """Test retrieving edges for a nonexistent node raises an exception."""
     graph = setup_graph
-    with pytest.raises(InvalidValueError, match="Node with id nonexistent does not exist."):
+    with pytest.raises(EntityNotFoundError, match="Node with id nonexistent does not exist."):
         graph.get_edges("nonexistent")