diff --git a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py new file mode 100644 index 00000000..137b9f7e --- /dev/null +++ b/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})) \ No newline at end of file diff --git a/cognee/tests/unit/modules/graph/cognee_graph_test.py b/cognee/tests/unit/modules/graph/cognee_graph_test.py new file mode 100644 index 00000000..235ccf11 --- /dev/null +++ b/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") diff --git a/examples/python/dynamic_steps_example.py b/examples/python/dynamic_steps_example.py index a10c26c7..309aea82 100644 --- a/examples/python/dynamic_steps_example.py +++ b/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: