feat: Adds unit tests to CogneeGraph class

cognee/tests/unit/modules/graph/cognee_graph_elements_test.py

@@ -0,0 +1,144 @@
+import pytest
+import numpy as np
+
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
+
+
+def test_node_initialization():
+    """Test that a Node is initialized correctly."""
+    node = Node("node1", {"attr1": "value1"}, dimension=2)
+    assert node.id == "node1"
+    assert node.attributes == {"attr1": "value1"}
+    assert len(node.status) == 2
+    assert np.all(node.status == 1)
+
+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"):
+        Node("node1", dimension=0)
+
+def test_add_skeleton_neighbor():
+    """Test adding a neighbor to a node."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    node1.add_skeleton_neighbor(node2)
+    assert node2 in node1.skeleton_neighbours
+
+def test_remove_skeleton_neighbor():
+    """Test removing a neighbor from a node."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    node1.add_skeleton_neighbor(node2)
+    node1.remove_skeleton_neighbor(node2)
+    assert node2 not in node1.skeleton_neighbours
+
+def test_add_skeleton_edge():
+    """Test adding an edge updates both skeleton_edges and skeleton_neighbours."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2)
+    node1.add_skeleton_edge(edge)
+    assert edge in node1.skeleton_edges
+    assert node2 in node1.skeleton_neighbours
+
+def test_remove_skeleton_edge():
+    """Test removing an edge updates both skeleton_edges and skeleton_neighbours."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2)
+    node1.add_skeleton_edge(edge)
+    node1.remove_skeleton_edge(edge)
+    assert edge not in node1.skeleton_edges
+    assert node2 not in node1.skeleton_neighbours
+
+def test_is_node_alive_in_dimension():
+    """Test checking node's alive status in a specific dimension."""
+    node = Node("node1", dimension=2)
+    assert node.is_node_alive_in_dimension(1)
+    node.status[1] = 0
+    assert not node.is_node_alive_in_dimension(1)
+
+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"):
+        node.is_node_alive_in_dimension(1)
+
+def test_node_equality():
+    """Test equality between nodes."""
+    node1 = Node("node1")
+    node2 = Node("node1")
+    assert node1 == node2
+
+def test_node_hash():
+    """Test hashing for Node."""
+    node = Node("node1")
+    assert hash(node) == hash("node1")
+
+### Tests for Edge ###
+
+def test_edge_initialization():
+    """Test that an Edge is initialized correctly."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, {"weight": 10}, directed=False, dimension=2)
+    assert edge.node1 == node1
+    assert edge.node2 == node2
+    assert edge.attributes == {"weight": 10}
+    assert edge.directed is False
+    assert len(edge.status) == 2
+    assert np.all(edge.status == 1)
+
+def test_edge_invalid_dimension():
+    """Test that initializing an Edge with a non-positive dimension raises an error."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    with pytest.raises(ValueError, match="Dimensions must be a positive integer."):
+        Edge(node1, node2, dimension=0)
+
+def test_is_edge_alive_in_dimension():
+    """Test checking edge's alive status in a specific dimension."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, dimension=2)
+    assert edge.is_edge_alive_in_dimension(1)
+    edge.status[1] = 0
+    assert not edge.is_edge_alive_in_dimension(1)
+
+def test_edge_alive_invalid_dimension():
+    """Test that checking alive status with an invalid dimension raises an error."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, dimension=1)
+    with pytest.raises(ValueError, match="Dimension 1 is out of range"):
+        edge.is_edge_alive_in_dimension(1)
+
+def test_edge_equality_directed():
+    """Test equality between directed edges."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge1 = Edge(node1, node2, directed=True)
+    edge2 = Edge(node1, node2, directed=True)
+    assert edge1 == edge2
+
+def test_edge_equality_undirected():
+    """Test equality between undirected edges."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge1 = Edge(node1, node2, directed=False)
+    edge2 = Edge(node2, node1, directed=False)
+    assert edge1 == edge2
+
+def test_edge_hash_directed():
+    """Test hashing for directed edges."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, directed=True)
+    assert hash(edge) == hash((node1, node2))
+
+def test_edge_hash_undirected():
+    """Test hashing for undirected edges."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, directed=False)
+    assert hash(edge) == hash(frozenset({node1, node2}))

cognee/tests/unit/modules/graph/cognee_graph_test.py

@@ -0,0 +1,79 @@
+import pytest
+
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
+from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
+
+
+@pytest.fixture
+def setup_graph():
+    """Fixture to initialize a CogneeGraph instance."""
+    return CogneeGraph()
+
+def test_add_node_success(setup_graph):
+    """Test successful addition of a node."""
+    graph = setup_graph
+    node = Node("node1")
+    graph.add_node(node)
+    assert graph.get_node("node1") == node
+
+def test_add_duplicate_node(setup_graph):
+    """Test adding a duplicate node raises an exception."""
+    graph = setup_graph
+    node = Node("node1")
+    graph.add_node(node)
+    with pytest.raises(ValueError, match="Node with id node1 already exists."):
+        graph.add_node(node)
+
+def test_add_edge_success(setup_graph):
+    """Test successful addition of an edge."""
+    graph = setup_graph
+    node1 = Node("node1")
+    node2 = Node("node2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    edge = Edge(node1, node2)
+    graph.add_edge(edge)
+    assert edge in graph.edges
+    assert edge in node1.skeleton_edges
+    assert edge in node2.skeleton_edges
+
+def test_add_duplicate_edge(setup_graph):
+    """Test adding a duplicate edge raises an exception."""
+    graph = setup_graph
+    node1 = Node("node1")
+    node2 = Node("node2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    edge = Edge(node1, node2)
+    graph.add_edge(edge)
+    with pytest.raises(ValueError, match="Edge .* already exists in the graph."):
+        graph.add_edge(edge)
+
+def test_get_node_success(setup_graph):
+    """Test retrieving an existing node."""
+    graph = setup_graph
+    node = Node("node1")
+    graph.add_node(node)
+    assert graph.get_node("node1") == node
+
+def test_get_node_nonexistent(setup_graph):
+    """Test retrieving a nonexistent node returns None."""
+    graph = setup_graph
+    assert graph.get_node("nonexistent") is None
+
+def test_get_edges_success(setup_graph):
+    """Test retrieving edges of a node."""
+    graph = setup_graph
+    node1 = Node("node1")
+    node2 = Node("node2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    edge = Edge(node1, node2)
+    graph.add_edge(edge)
+    assert edge in graph.get_edges("node1")
+
+def test_get_edges_nonexistent_node(setup_graph):
+    """Test retrieving edges for a nonexistent node raises an exception."""
+    graph = setup_graph
+    with pytest.raises(ValueError, match="Node with id nonexistent does not exist."):
+        graph.get_edges("nonexistent")

examples/python/dynamic_steps_example.py

@@ -209,7 +209,7 @@ async def main(enable_steps):
     if enable_steps.get("search_insights"):
         search_results = await cognee.search(
             SearchType.INSIGHTS,
-            {'query': 'Which applicant has the most relevant experience?'}
+            {'query': 'Which applicant has the most relevant experience in data science?'}
         )
         print("Search results:")
         for result_text in search_results: