lionagi 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

lionagi/tests/test_core/collections/test_pile.py

@@ -0,0 +1,171 @@
+import unittest
+import asyncio
+from lionagi.core.collections.pile import Pile
+from lionagi import Node
+
+
+class TestPile(unittest.TestCase):
+
+    def setUp(self):
+        """Create initial nodes and piles for testing."""
+        self.nodes1 = [Node(content=i) for i in ["A", "B", "C"]]
+        self.nodes2 = [Node(content=i) for i in ["D", "E", "F"]]
+        self.p1 = Pile(self.nodes1)
+        self.p2 = Pile(self.nodes2)
+
+    def test_add_piles(self):
+        """Test adding two piles."""
+        combined_pile = self.p1 + self.p2
+        self.assertEqual(len(combined_pile), 6)
+        contents = [node.content for node in combined_pile.values()]
+        self.assertListEqual(contents, ["A", "B", "C", "D", "E", "F"])
+
+    def test_subtract_piles(self):
+        """Test subtracting nodes from a pile."""
+        with self.assertRaises(Exception):
+            result_pile = self.p1 - self.nodes2[0]  # D not in p1
+
+    def test_containment(self):
+        """Test containment check."""
+        self.assertIn(self.nodes1[1], self.p1)  # B in p1
+        self.assertNotIn(self.nodes2[0], self.p1)  # D not in p1
+
+    def test_include_nodes(self):
+        """Test including nodes in a pile."""
+        self.p1.include(self.nodes2)
+        self.assertEqual(len(self.p1), 6)
+        contents = [node.content for node in self.p1.values()]
+        self.assertListEqual(contents, ["A", "B", "C", "D", "E", "F"])
+
+    def test_exclude_nodes(self):
+        """Test excluding nodes from a pile."""
+        self.p1.include(self.nodes2)  # Include first to have something to exclude
+        self.p1.exclude(self.nodes2)
+        self.assertEqual(len(self.p1), 5)
+        contents = [node.content for node in self.p1.values()]
+        self.assertListEqual(contents, ["A", "B", "C", "E", "F"])
+
+    def test_get_item_by_index(self):
+        """Test getting an item by index."""
+        node = self.p1[0]
+        self.assertEqual(node.content, "A")
+
+    def test_set_item_by_index(self):
+        """Test setting an item at a specific index."""
+        new_node = Node(content="Updated Node")
+        self.p1[0] = new_node
+        self.assertEqual(self.p1[0].content, "Updated Node")
+
+    def test_insert_item(self):
+        """Test inserting an item at a specific position."""
+        new_node = Node(content="Inserted Node")
+        self.p1.insert(0, new_node)
+        self.assertEqual(self.p1[0].content, "Inserted Node")
+        self.assertEqual(len(self.p1), 4)
+
+    def test_homogenous(self):
+        """Test if all items in the pile are of the same type."""
+        self.assertTrue(self.p1.is_homogenous())
+        with self.assertRaises(AttributeError):
+            mixed_pile = Pile(self.nodes1 + [123])  # Adding an int to mix types
+
+    def test_is_empty(self):
+        """Test if the pile is empty."""
+        empty_pile = Pile()
+        self.assertTrue(empty_pile.is_empty())
+        self.assertFalse(self.p1.is_empty())
+
+    def test_pop_item(self):
+        """Test popping an item from the pile."""
+        popped_node = self.p1.pop(self.nodes1[0].ln_id)
+        self.assertEqual(popped_node.content, "A")
+        self.assertEqual(len(self.p1), 2)
+
+    def test_clear_pile(self):
+        """Test clearing the pile."""
+        self.p1.clear()
+        self.assertTrue(self.p1.is_empty())
+
+    def test_update_pile(self):
+        """Test updating the pile with another pile."""
+        self.p1.update(self.p2)
+        self.assertEqual(len(self.p1), 6)
+        contents = [node.content for node in self.p1.values()]
+        self.assertListEqual(contents, ["A", "B", "C", "D", "E", "F"])
+
+    def test_to_dataframe(self):
+        """Test converting the pile to a DataFrame."""
+        df = self.p1.to_df()
+        self.assertEqual(len(df), 3)
+        self.assertListEqual(df["content"].tolist(), ["A", "B", "C"])
+
+    def test_create_index(self):
+        """Test creating an index for the pile."""
+        with self.assertRaises(ValueError):
+            self.p1.create_index(index_type="llama_index")
+
+    def test_create_query_engine(self):
+        """Test creating a query engine for the pile."""
+        with self.assertRaises(ValueError):
+            self.p1.create_query_engine(index_type="llama_index")
+
+    def test_create_chat_engine(self):
+        """Test creating a chat engine for the pile."""
+        with self.assertRaises(ValueError):
+            self.p1.create_chat_engine(index_type="llama_index")
+
+    # async def test_query_pile(self):
+    #     """Test querying the pile."""
+    #     self.p1.create_query_engine(index_type="llama_index")
+    #     response = await self.p1.query_pile(query="test query")
+    #     self.assertIsInstance(response, str)
+
+    # async def test_chat_pile(self):
+    #     """Test chatting with the pile."""
+    #     self.p1.create_chat_engine(index_type="llama_index")
+    #     response = await self.p1.chat_pile(query="test chat")
+    #     self.assertIsInstance(response, str)
+
+    # async def test_embed_pile(self):
+    #     """Test embedding the items in the pile."""
+    #     await self.p1.embed_pile()
+    #     for node in self.p1.values():
+    #         self.assertIn("embedding", node._all_fields)
+
+    def test_to_csv(self):
+        """Test saving the pile to a CSV file."""
+        self.p1.to_csv("test_pile.csv")
+        with open("test_pile.csv", "r") as f:
+            content = f.read()
+        self.assertIn("content", content)
+
+    def test_from_csv(self):
+        """Test loading a pile from a CSV file."""
+        self.p1.to_csv("test_pile.csv")
+        loaded_pile = Pile.from_csv("test_pile.csv")
+        self.assertEqual(len(loaded_pile), 3)
+        self.assertEqual(loaded_pile[0].content, "A")
+
+    def test_from_dataframe(self):
+        """Test loading a pile from a DataFrame."""
+        df = self.p1.to_df()
+        loaded_pile = Pile.from_df(df)
+        self.assertEqual(len(loaded_pile), 3)
+        self.assertEqual(loaded_pile[0].content, "A")
+
+    def test_as_query_tool(self):
+        """Test creating a query tool for the pile."""
+        with self.assertRaises(ValueError):
+            tool = self.p1.as_query_tool(index_type="llama_index")
+
+    def test_str_representation(self):
+        """Test the string representation of the pile."""
+        self.assertIsInstance(str(self.p1), str)
+
+    def test_repr_representation(self):
+        """Test the representation of the pile."""
+        self.assertIsInstance(repr(self.p1), str)
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/collections/test_progression.py

@@ -0,0 +1,129 @@
+import unittest
+from lionagi.core.collections.abc import ItemNotFoundError
+from lionagi.core.collections import Progression
+from lionagi import Node
+
+
+class TestProgression(unittest.TestCase):
+
+    def setUp(self):
+        self.nodes = [Node() for _ in range(5)]
+        self.p = Progression(order=[node.ln_id for node in self.nodes])
+
+    def test_initial_order(self):
+        self.assertEqual(self.p.order, [node.ln_id for node in self.nodes])
+
+    def test_append_node(self):
+        new_node = Node()
+        self.p.append(new_node)
+        self.assertIn(new_node.ln_id, self.p.order)
+        self.assertEqual(len(self.p), 6)
+
+    def test_extend_single_node(self):
+        new_node = Node()
+        self.p.extend(new_node)
+        self.assertIn(new_node.ln_id, self.p.order)
+        self.assertEqual(len(self.p), 6)
+
+    def test_extend_multiple_nodes(self):
+        new_nodes = [Node() for _ in range(3)]
+        self.p.extend(new_nodes)
+        for node in new_nodes:
+            self.assertIn(node.ln_id, self.p.order)
+        self.assertEqual(len(self.p), 8)
+
+    def test_include_node(self):
+        new_node = Node()
+        self.p.include(new_node)
+        self.assertIn(new_node.ln_id, self.p.order)
+
+    def test_getitem_slice(self):
+        slice_result = self.p[1:5]
+        self.assertEqual(slice_result.order, self.p.order[1:5])
+
+    def test_setitem(self):
+        new_node = Node()
+        self.p[0] = new_node.ln_id
+        self.assertEqual(self.p[0], new_node.ln_id)
+
+    def test_add_node(self):
+        new_node = Node()
+        new_prog = self.p + new_node
+        self.assertIn(new_node.ln_id, new_prog.order)
+
+    def test_iadd_node(self):
+        new_node = Node()
+        self.p += new_node
+        self.assertIn(new_node.ln_id, self.p.order)
+
+    def test_subtract_node(self):
+        node_to_remove = self.nodes[0]
+        self.p -= node_to_remove.ln_id
+        self.assertNotIn(node_to_remove.ln_id, self.p.order)
+
+    def test_isubtract_node(self):
+        node_to_remove = self.nodes[0]
+        self.p -= node_to_remove.ln_id
+        self.assertNotIn(node_to_remove.ln_id, self.p.order)
+
+    def test_popleft(self):
+        leftmost = self.p.order[0]
+        self.assertEqual(self.p.popleft(), leftmost)
+        self.assertNotIn(leftmost, self.p.order)
+
+    def test_exclude_node(self):
+        node_to_exclude = self.nodes[0]
+        self.p.exclude(node_to_exclude.ln_id)
+        self.assertNotIn(node_to_exclude.ln_id, self.p.order)
+
+    def test_clear(self):
+        self.p.clear()
+        self.assertEqual(len(self.p), 0)
+
+    def test_to_dict(self):
+        p_dict = self.p.to_dict()
+        self.assertEqual(p_dict["order"], self.p.order)
+        self.assertEqual(p_dict["name"], self.p.name)
+
+    def test_bool(self):
+        self.assertTrue(bool(self.p))
+
+    def test_remove_nonexistent_node(self):
+        non_existent_node = Node()
+        with self.assertRaises(ItemNotFoundError):
+            self.p.remove(non_existent_node.ln_id)
+
+    def test_pop_index_error(self):
+        with self.assertRaises(ItemNotFoundError):
+            self.p.pop(100)
+
+    def test_popleft_index_error(self):
+        empty_prog = Progression()
+        with self.assertRaises(ItemNotFoundError):
+            empty_prog.popleft()
+
+    def test_iteration(self):
+        nodes_list = [node.ln_id for node in self.nodes]
+        for i, node in enumerate(self.p):
+            self.assertEqual(node, nodes_list[i])
+
+    def test_contains(self):
+        node_to_check = self.nodes[0]
+        self.assertIn(node_to_check.ln_id, self.p)
+
+    def test_not_contains(self):
+        non_existent_node = Node()
+        self.assertNotIn(non_existent_node.ln_id, self.p)
+
+    def test_equality(self):
+        new_prog = Progression(order=[node.ln_id for node in self.nodes])
+        self.assertEqual(self.p.order, new_prog.order)
+
+    def test_inequality(self):
+        new_node = Node()
+        new_prog = self.p + new_node
+        self.assertNotEqual(self.p.order, new_prog)
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/generic/test_edge.py

@@ -0,0 +1,67 @@
+import unittest
+from unittest.mock import AsyncMock, patch
+from pydantic import ValidationError
+from lionagi.core.collections.abc import Component, get_lion_id, LionIDable, Condition
+from lionagi.core.generic.edge_condition import EdgeCondition
+from lionagi.core.generic.edge import Edge
+
+
+class TestEdge(unittest.TestCase):
+
+    def setUp(self):
+        self.node1 = Component()
+        self.node2 = Component()
+        self.edge = Edge(head=self.node1.ln_id, tail=self.node2.ln_id)
+
+    def test_initialization(self):
+        """Test initialization of the Edge class."""
+        self.assertEqual(self.edge.head, self.node1.ln_id)
+        self.assertEqual(self.edge.tail, self.node2.ln_id)
+        self.assertIsNone(self.edge.condition)
+        self.assertIsNone(self.edge.label)
+        self.assertFalse(self.edge.bundle)
+
+    async def test_check_condition_without_condition(self):
+        """Test check_condition raises ValueError when condition is not set."""
+        with self.assertRaises(ValueError):
+            await self.edge.check_condition({})
+
+    @patch("lionagi.core.models.edge.EdgeCondition.applies", new_callable=AsyncMock)
+    async def test_check_condition_with_condition(self, mock_applies):
+        """Test check_condition with a set condition."""
+        mock_applies.return_value = True
+        condition = EdgeCondition()
+        self.edge.condition = condition
+        result = await self.edge.check_condition({})
+        self.assertTrue(result)
+        mock_applies.assert_awaited_once()
+
+    def test_validate_head_tail(self):
+        """Test the head and tail validation."""
+        try:
+            valid_id = get_lion_id("valid_head")
+        except Exception as e:
+            if e.__class__.__name__ == "LionTypeError":
+                return
+        self.assertEqual(Edge._validate_head_tail(valid_id), valid_id)
+        with self.assertRaises(ValidationError):
+            Edge._validate_head_tail("invalid_id")
+
+    def test_string_condition_none(self):
+        """Test string_condition method when condition is None."""
+        self.assertIsNone(self.edge.string_condition())
+
+    def test_len(self):
+        """Test the __len__ method."""
+        self.assertEqual(len(self.edge), 1)
+
+    def test_contains(self):
+        """Test the __contains__ method."""
+        self.assertIn(self.node1.ln_id, self.edge)
+        self.assertIn(self.node2.ln_id, self.edge)
+        fake_node = Component()
+        self.assertNotIn(fake_node.ln_id, self.edge)
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/generic/test_graph.py

@@ -0,0 +1,96 @@
+import unittest
+from lionagi.core.collections.abc import LionTypeError, ItemNotFoundError
+from lionagi.core.generic import Graph, Node, Edge
+from lionagi.libs.ln_convert import to_list
+
+
+class TestGraph(unittest.TestCase):
+
+    def setUp(self):
+        """Set up a basic graph for testing."""
+        self.node1 = Node()
+        self.node2 = Node()
+        self.node3 = Node()
+        self.node4 = Node()
+        self.g = Graph()
+
+    def test_initial_size(self):
+        """Test the initial size of the graph."""
+        self.assertEqual(self.g.size(), 0)
+
+    def test_initial_internal_edges(self):
+        """Test the initial number of internal edges."""
+        self.assertEqual(len(self.g.internal_edges), 0)
+
+    def test_add_node_no_edge(self):
+        """Test that adding nodes does not create any edges."""
+        self.g.add_node([self.node1, self.node2, self.node3, self.node4])
+        self.assertEqual(len(self.g.internal_edges), 0)
+
+    def test_add_edges(self):
+        """Test adding edges between nodes."""
+        self.g.add_edge(self.node1, self.node2)
+        self.g.add_edge(self.node2, self.node3)
+        self.g.add_edge(self.node3, self.node4)
+        self.assertEqual(len(self.g.internal_edges), 3)
+
+    def test_get_node(self):
+        """Test retrieving a node by its identifier."""
+        self.g.add_node(self.node2)
+        node5 = self.g.get_node(self.node2)
+        self.assertEqual(node5, self.node2)
+
+    def test_get_node_edges(self):
+        """Test getting the edges of a node."""
+        self.g.add_edge(self.node1, self.node2)
+        self.assertEqual(len(self.g.get_node_edges(self.node1)), 1)
+
+    def test_get_heads(self):
+        """Test getting head nodes."""
+        self.g.add_edge(self.node1, self.node2)
+        self.g.add_edge(self.node2, self.node3)
+        self.g.add_edge(self.node3, self.node4)
+        self.assertEqual(len(self.g.get_heads()), 1)
+
+    def test_is_acyclic(self):
+        """Test if the graph is acyclic."""
+        self.g.add_edge(self.node1, self.node2)
+        self.g.add_edge(self.node2, self.node3)
+        self.g.add_edge(self.node3, self.node4)
+        self.assertTrue(self.g.is_acyclic())
+
+    def test_remove_edge(self):
+        """Test removing an edge from the graph."""
+        self.g.add_edge(self.node1, self.node2)
+        self.g.add_edge(self.node2, self.node3)
+        self.g.add_edge(self.node3, self.node4)
+        edge = self.g.internal_edges[0]
+        self.g.remove_edge(edge)
+        self.assertEqual(len(self.g.internal_edges), 2)
+
+    def test_remove_nonexistent_edge(self):
+        """Test removing a nonexistent edge."""
+        self.g.add_edge(self.node1, self.node2)
+        with self.assertRaises(ItemNotFoundError):
+            self.g.remove_edge(edge=Edge(head=self.node2, tail=self.node3))
+
+    def test_clear(self):
+        """Test clearing all nodes and edges from the graph."""
+        self.g.add_node([self.node1, self.node2, self.node3, self.node4])
+        self.g.add_edge(self.node1, self.node2)
+        self.g.add_edge(self.node2, self.node3)
+        self.g.add_edge(self.node3, self.node4)
+        self.g.clear()
+        self.assertTrue(self.g.is_empty())
+        self.assertEqual(len(self.g.internal_edges), 0)
+
+    def test_display(self):
+        """Test displaying the graph (visual check)."""
+        self.g.add_edge(self.node1, self.node2)
+        self.g.add_edge(self.node2, self.node3)
+        self.g.add_edge(self.node3, self.node4)
+        self.g.display()  # This should display the graph visually
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/generic/test_node.py

@@ -0,0 +1,106 @@
+import unittest
+from lionagi.core.generic import Node
+from lionagi.core.collections.abc import RelationError, Condition
+from lionagi.core.generic.edge import Edge
+
+
+class TestNode(unittest.TestCase):
+
+    def setUp(self):
+        """Setup nodes for testing."""
+        self.node_a = Node()
+        self.node_b = Node()
+        self.node_c = Node()
+
+    def test_initialization(self):
+        """Test that a Node initializes with empty relations."""
+        self.assertEqual(len(self.node_a.relations["in"]), 0)
+        self.assertEqual(len(self.node_a.relations["out"]), 0)
+
+    def test_relate_outgoing(self):
+        """Test relating nodes with an outgoing edge."""
+        self.node_a.relate(self.node_b, "out")
+        self.assertIn(self.node_b.ln_id, self.node_a.related_nodes)
+        self.assertIn(self.node_a.ln_id, self.node_b.related_nodes)
+        self.assertIn(self.node_b.ln_id, self.node_a.successors)
+        self.assertIn(self.node_a.ln_id, self.node_b.predecessors)
+
+    def test_relate_incoming(self):
+        """Test relating nodes with an incoming edge."""
+        self.node_a.relate(self.node_c, "in")
+        self.assertIn(self.node_c.ln_id, self.node_a.related_nodes)
+        self.assertIn(self.node_a.ln_id, self.node_c.related_nodes)
+        self.assertIn(self.node_a.ln_id, self.node_c.successors)
+        self.assertIn(self.node_c.ln_id, self.node_a.predecessors)
+
+    def test_invalid_relate(self):
+        """Test that relating nodes with an invalid direction raises ValueError."""
+        with self.assertRaises(ValueError):
+            self.node_a.relate(self.node_b, "invalid_direction")
+
+    def test_remove_edge(self):
+        """Test removing a specific edge between nodes."""
+        self.node_a.relate(self.node_b, "out")
+        edge = next(iter(self.node_a.relations["out"]))
+        self.assertTrue(self.node_a.remove_edge(self.node_b, edge))
+
+    def test_unrelate_all_edges(self):
+        """Test unrelating all edges between nodes."""
+        self.node_a.relate(self.node_b, "out")
+        self.assertTrue(self.node_a.unrelate(self.node_b))
+        self.assertNotIn(self.node_b.ln_id, self.node_a.related_nodes)
+        self.assertNotIn(self.node_a.ln_id, self.node_b.related_nodes)
+
+    def test_unrelate_specific_edge(self):
+        """Test unrelating a specific edge between nodes."""
+        self.node_a.relate(self.node_b, "out")
+        edge = next(iter(self.node_a.relations["out"]))
+        self.assertTrue(self.node_a.unrelate(self.node_b, edge))
+        self.assertNotIn(self.node_b.ln_id, self.node_a.related_nodes)
+        self.assertNotIn(self.node_a.ln_id, self.node_b.related_nodes)
+
+    def test_node_relations(self):
+        """Test that node relations are correctly categorized."""
+        self.node_a.relate(self.node_b, "out")
+        self.node_a.relate(self.node_c, "in")
+        relations = self.node_a.node_relations
+        self.assertIn(self.node_b.ln_id, relations["out"])
+        self.assertIn(self.node_c.ln_id, relations["in"])
+
+    def test_related_nodes(self):
+        """Test that related nodes list is correct."""
+        self.node_a.relate(self.node_b, "out")
+        self.node_a.relate(self.node_c, "in")
+        related_nodes = self.node_a.related_nodes
+        self.assertIn(self.node_b.ln_id, related_nodes)
+        self.assertIn(self.node_c.ln_id, related_nodes)
+
+    def test_predecessors(self):
+        """Test that predecessors list is correct."""
+        self.node_a.relate(self.node_c, "in")
+        self.assertIn(self.node_c.ln_id, self.node_a.predecessors)
+
+    def test_successors(self):
+        """Test that successors list is correct."""
+        self.node_a.relate(self.node_b, "out")
+        self.assertIn(self.node_b.ln_id, self.node_a.successors)
+
+    def test_str(self):
+        """Test the string representation of a node."""
+        self.node_a.relate(self.node_b, "out")
+        node_str = str(self.node_a)
+        self.assertIn(self.node_a.ln_id, node_str)
+        self.assertIn("relations", node_str)
+        self.assertIn("1", node_str)  # 1 relation
+
+    def test_repr(self):
+        """Test the repr representation of a node."""
+        self.node_a.relate(self.node_b, "out")
+        node_repr = repr(self.node_a)
+        self.assertIn(self.node_a.ln_id, node_repr)
+        self.assertIn("relations", node_repr)
+        self.assertIn("1", node_repr)  # 1 relation
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/generic/test_tree_node.py

@@ -0,0 +1,73 @@
+import unittest
+from lionagi.core.generic.tree_node import TreeNode
+from lionagi.core.generic.tree import Tree
+
+
+class TestTreeNode(unittest.TestCase):
+
+    def setUp(self):
+        self.tree_node1 = TreeNode()
+        self.tree_node2 = TreeNode()
+        self.tree_node3 = TreeNode()
+        self.tree_node4 = TreeNode()
+        self.tree = Tree()
+
+    def test_initialization(self):
+        self.assertIsNone(self.tree_node1.parent)
+        self.assertEqual(self.tree_node1.children, [])
+        self.assertIsInstance(self.tree_node1, TreeNode)
+        self.assertIsInstance(self.tree_node1.ln_id, str)
+
+    def test_relate_child(self):
+        self.tree_node1.relate_child([self.tree_node2, self.tree_node3])
+        self.assertIn(self.tree_node2.ln_id, self.tree_node1.children)
+        self.assertIn(self.tree_node3.ln_id, self.tree_node1.children)
+        self.assertEqual(self.tree_node2.parent, self.tree_node1)
+        self.assertEqual(self.tree_node3.parent, self.tree_node1)
+
+    def test_relate_parent(self):
+        self.tree_node1.relate_parent(self.tree_node4)
+        self.assertEqual(self.tree_node1.parent, self.tree_node4)
+        self.assertIn(self.tree_node1.ln_id, self.tree_node4.children)
+
+    def test_unrelate_parent(self):
+        self.tree_node1.relate_parent(self.tree_node4)
+        self.tree_node1.unrelate_parent()
+        self.assertIsNone(self.tree_node1.parent)
+
+    def test_unrelate_child(self):
+        self.tree_node1.relate_child([self.tree_node2, self.tree_node3])
+        self.tree_node1.unrelate_child([self.tree_node2, self.tree_node3])
+        self.assertNotIn(self.tree_node2.ln_id, self.tree_node1.children)
+        self.assertNotIn(self.tree_node3.ln_id, self.tree_node1.children)
+        self.assertIsNone(self.tree_node2.parent)
+        self.assertIsNone(self.tree_node3.parent)
+
+    def test_tree_integration(self):
+        self.tree.relate_parent_child(
+            self.tree_node1, [self.tree_node2, self.tree_node3, self.tree_node4]
+        )
+        self.assertEqual(len(self.tree.internal_nodes), 4)
+        self.assertEqual(len(self.tree_node1.children), 3)
+        self.assertEqual(self.tree.root, self.tree_node1)
+
+    def test_tree_node_edges(self):
+        self.tree_node1.relate_child([self.tree_node2, self.tree_node3])
+        df = self.tree_node1.edges.to_df()
+        self.assertEqual(len(df), 2)
+        self.assertIn(self.tree_node2.ln_id, df["tail"].values)
+        self.assertIn(self.tree_node3.ln_id, df["tail"].values)
+
+    def test_tree_display(self):
+        # This test is mainly to check if the display function runs without errors
+        self.tree.relate_parent_child(
+            self.tree_node1, [self.tree_node2, self.tree_node3, self.tree_node4]
+        )
+        try:
+            self.tree.display()
+        except Exception as e:
+            self.fail(f"Tree display method raised an exception: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()