privateGPT/tests/components/readers/nodes/test_node_structure.py

from typing import TYPE_CHECKING, Any

import pandas as pd
import pytest
from llama_index.core.schema import MetadataMode
from llama_index.core.vector_stores.utils import (
    node_to_metadata_dict,
)

from private_gpt.components.readers.nodes.document_node import DocumentRootNode
from private_gpt.components.readers.nodes.list_node import ListItemNode, ListNode
from private_gpt.components.readers.nodes.partial_node import PartialNode
from private_gpt.components.readers.nodes.section_node import SectionNode
from private_gpt.components.readers.nodes.table_node import TableNode, TableRowNode
from private_gpt.components.readers.nodes.text_node import TextNode
from private_gpt.components.readers.nodes.tree_node import (
    CURRENT_VERSION,
    TreeMetadataMode,
)
from private_gpt.components.readers.nodes.utils import combine_trees
from private_gpt.components.readers.nodes.v2.document_node_v2 import DocumentRootNodeV2

if TYPE_CHECKING:
    from private_gpt.components.readers.nodes.tree_node import TreeNode


def test_document_tree_structure() -> None:
    # Create a document tree with multiple node types
    root = DocumentRootNode()

    # Add a section
    section1 = SectionNode(text="Introduction")
    root.add_child(section1)

    # Add text to section
    text_node = TextNode(text="This is an introduction paragraph.")
    section1.add_child(text_node)

    # Create a table
    df = pd.DataFrame(
        {
            "Name": ["Alice", "Bob"],
            "Age": [30, 25],
            "City": ["New York", "San Francisco"],
        }
    )
    table = TableNode(df=df, description="Sample table")
    root.add_child(table)

    # Add table rows
    row1 = TableRowNode(header=df.columns.tolist(), content=df.iloc[0].tolist())
    table.add_child(row1)
    row2 = TableRowNode(header=df.columns.tolist(), content=df.iloc[1].tolist())
    table.add_child(row2)

    # Test tree structure
    assert len(root.children) == 2
    assert root.children[0] == section1
    assert root.children[1] == table

    assert len(section1.children) == 1
    assert section1.children[0] == text_node

    assert len(table.children) == 2
    assert table.children[0] == row1
    assert table.children[1] == row2


def test_node_serialization() -> None:
    # Create a sample document tree
    root = DocumentRootNode()
    section = SectionNode(text="Test Section")
    text = TextNode(text="Sample text")

    root.add_child(section)
    section.add_child(text)

    # Test serialization
    serialized = root.dict()
    assert serialized["class_name"] == f"DocumentRootNode-{CURRENT_VERSION}"
    assert "parent" not in serialized
    assert "children" not in serialized

    serialized = section.dict()
    assert serialized["class_name"] == f"SectionNode-{CURRENT_VERSION}"
    assert serialized["text"] == "Test Section"
    assert serialized["root_id"] == root.id_

    serialized = text.dict()
    assert serialized["class_name"] == f"TextNode-{CURRENT_VERSION}"
    assert serialized["text"] == "Sample text"
    assert serialized["root_id"] == root.id_


def test_node_serialization_utils() -> None:
    # Create a sample document tree
    root = DocumentRootNode()

    # Create a section
    section = SectionNode(text="Test Section")
    text = TextNode(text="Sample text")

    root.add_child(section)
    section.add_child(text)

    # Create a table
    df = pd.DataFrame(
        {
            "Name": ["Alice", "Bob"],
            "Age": [30, 25],
            "City": ["New York", "San Francisco"],
        }
    )
    table = TableNode(df=df, description="Sample table")
    section.add_child(table)

    # Add table rows
    row1 = TableRowNode(header=df.columns.tolist(), content=df.iloc[0].tolist())
    table.add_child(row1)
    row2 = TableRowNode(header=df.columns.tolist(), content=df.iloc[1].tolist())
    table.add_child(row2)

    # Test document root serialization
    serialized = node_to_metadata_dict(root)
    assert serialized["_node_type"] == f"DocumentRootNode-{CURRENT_VERSION}"

    # Test section serialization
    serialized = node_to_metadata_dict(section)
    assert serialized["_node_type"] == f"SectionNode-{CURRENT_VERSION}"

    # Test table serialization
    serialized = node_to_metadata_dict(table)
    assert serialized["_node_type"] == f"TableNode-{CURRENT_VERSION}"


def test_node_deserialization() -> None:
    # Create a sample document tree
    root = DocumentRootNode()
    section = SectionNode(text="Test Section")
    text = TextNode(text="Sample text")

    root.add_child(section)
    section.add_child(text)

    # Serialize and deserialize
    serialized = root.dict()
    doc_deserialized = DocumentRootNode.parse_obj(serialized)
    assert doc_deserialized.id_ == root.id_
    assert doc_deserialized.hash == root.hash
    assert doc_deserialized.idx == root.idx
    assert doc_deserialized.abs_idx == root.abs_idx
    assert doc_deserialized.depth == root.depth
    assert doc_deserialized.height == root.height

    serialized = section.dict()
    section_deserialized = SectionNode.parse_obj(serialized)
    assert section_deserialized.hash == section.hash
    assert section_deserialized.text == section.text
    assert section_deserialized.parent_id == section.parent_id
    assert section_deserialized.root_id == section.root_id
    assert section_deserialized.idx == section.idx
    assert section_deserialized.abs_idx == section.abs_idx
    assert section_deserialized.depth == section.depth
    assert section_deserialized.height == section.height

    serialized = text.dict()
    text_deserialized = TextNode.parse_obj(serialized)
    assert text_deserialized.hash == text.hash
    assert text_deserialized.text == text.text
    assert text_deserialized.parent_id == text.parent_id
    assert text_deserialized.root_id == text.root_id
    assert text_deserialized.idx == text.idx
    assert text_deserialized.abs_idx == text.abs_idx
    assert text_deserialized.depth == text.depth
    assert text_deserialized.height == text.height


def test_node_deserialization_utils() -> None:
    # Create a sample document tree
    root = DocumentRootNode()

    # Create a section
    section = SectionNode(text="Test Section")
    text = TextNode(text="Sample text")

    root.add_child(section)
    section.add_child(text)

    # Create a table
    df = pd.DataFrame(
        {
            "Name": ["Alice", "Bob"],
            "Age": [30, 25],
            "City": ["New York", "San Francisco"],
        }
    )
    table = TableNode(df=df, description="Sample table")
    section.add_child(table)

    # Add table rows
    row1 = TableRowNode(header=df.columns.tolist(), content=df.iloc[0].tolist())
    table.add_child(row1)
    row2 = TableRowNode(header=df.columns.tolist(), content=df.iloc[1].tolist())
    table.add_child(row2)

    # Test document root serialization
    serialized = node_to_metadata_dict(root)
    json = root.model_dump_json()
    node_content: str = serialized.get("_node_content") or ""
    assert node_content
    node_type: str = serialized.get("_node_type") or ""
    assert node_type == f"DocumentRootNode-{CURRENT_VERSION}"
    doc_obj = DocumentRootNode.from_json(node_content)
    doc_json = doc_obj.model_dump_json()
    assert doc_obj.hash == root.hash
    assert doc_obj.idx == root.idx
    assert doc_obj.abs_idx == root.abs_idx
    assert doc_obj.depth == root.depth
    assert doc_obj.height == root.height
    assert doc_json == json

    # Test section serialization
    serialized = node_to_metadata_dict(section)
    json = root.model_dump_json()
    node_content = serialized.get("_node_content") or ""
    assert node_content
    node_type = serialized.get("_node_type") or ""
    assert node_type == f"SectionNode-{CURRENT_VERSION}"
    section_obj = SectionNode.from_json(node_content)
    section_json = doc_obj.model_dump_json()
    assert section_obj.hash == section.hash
    assert section_obj.text == section.text
    assert section_obj.parent_id == section.parent_id
    assert section_obj.root_id == section.root_id
    assert section_obj.idx == section.idx
    assert section_obj.abs_idx == section.abs_idx
    assert section_obj.depth == section.depth
    assert section_obj.height == section.height
    assert json == section_json

    # Test table serialization
    serialized = node_to_metadata_dict(table)
    json = table.model_dump_json()
    node_content = serialized.get("_node_content") or ""
    assert node_content
    node_type = serialized.get("_node_type") or ""
    assert node_type == f"TableNode-{CURRENT_VERSION}"
    table_obj = TableNode.from_json(node_content)
    table_json = table_obj.model_dump_json()
    assert table_obj.hash == table.hash
    assert table_obj.description == table.description
    assert table_obj.df.equals(table.df)
    assert table_obj.parent_id == table.parent_id
    assert table_obj.root_id == table.root_id
    assert table_obj.idx == table.idx
    assert table_obj.abs_idx == table.abs_idx
    assert table_obj.depth == table.depth
    assert table_obj.height == table.height
    assert json == table_json


def test_node_flattening() -> None:
    # Create a nested document tree
    root = DocumentRootNode()
    section1 = SectionNode(text="Section 1")
    section2 = SectionNode(text="Section 2")
    text1 = TextNode(text="Text in Section 1")
    text2 = TextNode(text="Text in Section 2")

    root.add_child(section1)
    root.add_child(section2)
    section1.add_child(text1)
    section2.add_child(text2)

    # Test flattening
    flattened_nodes = list(root.flatten())
    assert len(flattened_nodes) == 5  # root + 2 sections + 2 texts
    assert flattened_nodes[0] == root
    assert flattened_nodes[1] == section1
    assert flattened_nodes[2] == text1
    assert flattened_nodes[3] == section2
    assert flattened_nodes[4] == text2


def test_node_metadata() -> None:
    # Test metadata functionality
    text_node = TextNode(text="Hello World")
    text_node.metadata["source"] = "test"

    # Test metadata string generation
    default_metadata_str = text_node.get_metadata_str()
    assert "source: test" in default_metadata_str

    # Test LLM metadata exclusion
    text_node.excluded_llm_metadata_keys = ["source"]
    llm_metadata_str = text_node.get_metadata_str(mode=MetadataMode.LLM)
    assert "source: test" not in llm_metadata_str


def test_node_hash() -> None:
    # Ensure hash is consistent for same node
    node1 = TextNode(text="Test")
    node2 = TextNode(text="Test")

    assert node1.hash == node2.hash

    # Ensure different nodes have different hashes
    node3 = TextNode(text="Different")
    assert node1.hash != node3.hash


def test_parent_child_relationship() -> None:
    root = DocumentRootNode()
    section = SectionNode(text="Test")
    text = TextNode(text="Content")

    root.add_child(section)
    section.add_child(text)

    # Test parent-child relationships
    assert text.parent == section
    assert section.parent == root
    assert root.parent is None

    assert text in section.children
    assert section in root.children


def test_document_node_tree_serialization() -> None:
    root = DocumentRootNode()
    section = SectionNode(text="Test")
    text = TextNode(text="Content")

    root.add_child(section)
    section.add_child(text)

    json = root.to_tree_serialization()
    assert json

    deserialized = DocumentRootNode()
    deserialized.from_tree_serialization(json)
    assert len(deserialized.children) == 1
    assert len(deserialized.children[0].children) == 1


@pytest.mark.parametrize("num_nodes", [1000, 10000, 100000])
def test_document_node_v1_tree_serialization(num_nodes: int) -> None:
    root = DocumentRootNode()
    num_nodes_by_type = num_nodes // 3
    text_num_nodes, table_num_nodes, list_num_nodes = (
        num_nodes_by_type,
        num_nodes_by_type,
        num_nodes_by_type,
    )

    texts: list[TreeNode] = []
    for i in range(text_num_nodes):
        section = SectionNode(text=f"Section {i}")
        text = TextNode(text=f"Text in Section {i}")
        section.add_child(text)
        root.add_child(section)
        texts.append(text)

    tables: list[TreeNode] = []
    for i in range(table_num_nodes):
        df = pd.DataFrame(
            {
                "Name": [f"Name {i}", f"Name {i + 1}"],
                "Age": [i, i + 1],
                "City": [f"City {i}", f"City {i + 1}"],
            }
        )
        table = TableNode(df=df, description=f"Table {i}")
        row1 = TableRowNode(header=df.columns.tolist(), content=df.iloc[0].tolist())
        row2 = TableRowNode(header=df.columns.tolist(), content=df.iloc[1].tolist())
        table.add_children(row1, row2)
        tables.append(table)

    lists: list[TreeNode] = []
    for i in range(list_num_nodes):
        list_node = ListNode()
        list_item = ListItemNode(text=f"List item {i}-{0}")
        list_node.add_child(list_item)
        list_item = ListItemNode(text=f"List item {i}-{1}")
        list_node.add_child(list_item)
        lists.append(list_node)

    nodes: list[TreeNode] = texts + tables + lists
    root.add_children(*nodes)
    root.update_references()

    original_tree_len = len(list(root.flatten()))

    serialization = root.to_tree_serialization()
    assert serialization

    deserialized = DocumentRootNode(id_=root.id_)
    deserialized.from_tree_serialization(serialization)
    deserialized_tree_len = len(list(deserialized.flatten()))

    assert deserialized_tree_len == original_tree_len


@pytest.mark.parametrize("num_nodes", [1000, 10000, 100000])
def test_document_node_v2_tree_serialization(num_nodes: int) -> None:
    root = DocumentRootNodeV2()
    num_nodes_by_type = num_nodes // 3
    text_num_nodes, table_num_nodes, list_num_nodes = (
        num_nodes_by_type,
        num_nodes_by_type,
        num_nodes_by_type,
    )

    texts: list[TreeNode] = []
    for i in range(text_num_nodes):
        section = SectionNode(text=f"Section {i}")
        text = TextNode(text=f"Text in Section {i}")
        section.add_child(text)
        root.add_child(section)
        texts.append(text)

    tables: list[TreeNode] = []
    for i in range(table_num_nodes):
        df = pd.DataFrame(
            {
                "Name": [f"Name {i}", f"Name {i + 1}"],
                "Age": [i, i + 1],
                "City": [f"City {i}", f"City {i + 1}"],
            }
        )
        table = TableNode(df=df, description=f"Table {i}")
        row1 = TableRowNode(header=df.columns.tolist(), content=df.iloc[0].tolist())
        row2 = TableRowNode(header=df.columns.tolist(), content=df.iloc[1].tolist())
        table.add_children(row1, row2)
        tables.append(table)

    lists: list[TreeNode] = []
    for i in range(list_num_nodes):
        list_node = ListNode()
        list_item = ListItemNode(text=f"List item {i}-{0}")
        list_node.add_child(list_item)
        list_item = ListItemNode(text=f"List item {i}-{1}")
        list_node.add_child(list_item)
        lists.append(list_node)

    nodes: list[TreeNode] = texts + tables + lists
    root.add_children(*nodes)
    root.update_references()

    original_tree_len = len(list(root.flatten()))

    serialization = root.to_tree_serialization()
    assert serialization

    deserialized = DocumentRootNodeV2(id_=root.id_)
    deserialized.from_tree_serialization(serialization)
    deserialized_tree_len = len(list(deserialized.flatten()))

    assert deserialized_tree_len == original_tree_len


def test_combine_trees() -> None:
    root1 = DocumentRootNode()
    section1 = SectionNode(text="Test")
    text1 = TextNode(text="Content")

    root1.add_child(section1)
    section1.add_child(text1)

    root2 = DocumentRootNode()
    section2 = SectionNode(text="Test2")
    text2 = TextNode(text="Content2")
    text3 = TextNode(text="Content3")

    root2.add_child(section2)
    section2.add_child(text2)
    section2.add_child(text3)

    root3 = DocumentRootNode()
    section3 = SectionNode(text="Test3")
    section4 = SectionNode(text="Test4")
    text4 = TextNode(text="Content4")

    root3.add_child(section3)
    root3.add_child(section4)
    section3.add_child(text4)

    combined = combine_trees(root1, root2, root3)
    assert len(combined.children) == 4

    # Check first level
    assert combined.children[0].text == "Test"
    assert combined.children[1].text == "Test2"
    assert combined.children[2].text == "Test3"
    assert combined.children[3].text == "Test4"

    # Check that original nodes are not modified
    assert root1.children[0].text == "Test"

    # Check rest of content
    tree_as_list = list(combined.flatten())
    assert len(tree_as_list) == 9
    assert combined.id_ == root1.id_
    assert all(node.root_id == combined.id_ for node in tree_as_list[1:])


def test_prune_trees() -> None:
    # Case 1. Empty section
    root: TreeNode = DocumentRootNode()
    section = SectionNode(text="Test Section")
    root.add_child(section)
    assert root.children[0] == section
    root = root.prune()
    assert root is None

    # Case 2. Non-empty section
    root: TreeNode = DocumentRootNode()
    text = TextNode(text="Sample text")
    root.add_child(section)
    section.add_child(text)
    assert root.children[0] == section
    assert section.children[0] == text
    root = root.prune()
    assert root is not None
    assert len(root.children) == 1
    assert len(section.children) == 1

    # Case 3. Empty document
    root: TreeNode = DocumentRootNode()
    root = root.prune()
    assert root is None


@pytest.mark.parametrize(
    ("metadata_mode", "metadata", "expected"),
    [
        (
            TreeMetadataMode.ALL,
            {"key1": "value1", "key2": "value2"},
            "key1: value1\nkey2: value2",
        ),
        (
            MetadataMode.LLM,
            {"key1": "value1", "key2": "value2"},
            "key1: value1\nkey2: value2",
        ),
        ("embed", {"key1": "value1", "key2": "value2"}, "key1: value1\nkey2: value2"),
        (TreeMetadataMode.NONE, {"key1": "value1", "key2": "value2"}, ""),
    ],
)
def test_get_metadata_str(
    metadata_mode: TreeMetadataMode | MetadataMode | str,
    metadata: dict[str, str],
    expected: str,
) -> None:
    node = TextNode(text="", extra_info=metadata)
    result = node.get_metadata_str(metadata_mode)
    assert result == expected


@pytest.mark.parametrize(
    ("metadata_mode", "expected"),
    [
        (TreeMetadataMode.ALL, "Content for mode: all"),
        (MetadataMode.LLM, "Content for mode: llm"),
        ("rag", "Content for mode: rag"),
        (TreeMetadataMode.NONE, "Content for mode: none"),
    ],
)
def test_get_content(
    metadata_mode: TreeMetadataMode | MetadataMode | str, expected: str
) -> None:
    node = TextNode(text=expected)
    result = node.get_content(metadata_mode)
    assert result == expected


@pytest.mark.parametrize(
    ("instance", "real_class", "expected"),
    [
        (PartialNode(type=TextNode.get_type()), TextNode, True),
        (PartialNode(type=TextNode.get_type()), SectionNode, False),
        (PartialNode(type=""), None, False),
        (DocumentRootNode(), DocumentRootNode, True),
        (DocumentRootNode(), SectionNode, False),
        (DocumentRootNode(), None, False),
    ],
)
def test_isinstance(instance: Any, real_class: Any, expected: bool) -> None:
    result = instance.isinstance(real_class)
    assert result == expected


def test_add_children() -> None:
    root = DocumentRootNode()

    # Add first section with text
    section1 = SectionNode(text="Introduction")
    root.add_child(section1)
    text1 = TextNode(text="This is the introduction.")
    section1.add_child(text1)

    # Add second section with text
    section2 = SectionNode(text="Body")
    root.add_child(section2)
    text2 = TextNode(text="This is the body text.")
    section2.add_child(text2)

    # Add a subsection with text
    subsection = SectionNode(text="Subsection")
    section2.add_child(subsection)
    subtext = TextNode(text="This is the subsection text.")
    subsection.add_child(subtext)

    # Add a table with rows
    df = pd.DataFrame(
        {
            "Name": ["Alice", "Bob"],
            "Age": [30, 25],
            "City": ["New York", "San Francisco"],
        }
    )
    table = TableNode(df=df, description="Sample table")
    section2.add_child(table)
    row1 = TableRowNode(header=df.columns.tolist(), content=df.iloc[0].tolist())
    row2 = TableRowNode(header=df.columns.tolist(), content=df.iloc[1].tolist())
    table.add_child(row1)
    table.add_child(row2)

    # Add a list with items
    list_node = ListNode()
    root.add_child(list_node)
    list_item1 = ListItemNode(text="First item")
    list_item2 = ListItemNode(text="Second item")
    list_node.add_child(list_item1)
    list_node.add_child(list_item2)

    # Add nested list items
    nested_list_item = ListItemNode(text="Nested item")
    list_item2.add_child(nested_list_item)

    # Root update references
    root.update_references()

    # Assertions to check the tree structure
    assert len(root.children) == 3
    assert root.children[0] == section1
    assert root.children[1] == section2
    assert root.children[2] == list_node

    assert len(section1.children) == 1
    assert section1.children[0] == text1

    assert len(section2.children) == 3
    assert section2.children[0] == text2
    assert section2.children[1] == subsection
    assert section2.children[2] == table

    assert len(subsection.children) == 1
    assert subsection.children[0] == subtext

    assert len(table.children) == 2
    assert table.children[0] == row1
    assert table.children[1] == row2

    assert len(list_node.children) == 2
    assert list_node.children[0] == list_item1
    assert list_node.children[1] == list_item2

    assert len(list_item2.children) == 1
    assert list_item2.children[0] == nested_list_item

    # Check parents
    assert section1.parent == root
    assert section2.parent == root
    assert text1.parent == section1
    assert text2.parent == section2
    assert subsection.parent == section2
    assert subtext.parent == subsection
    assert table.parent == section2
    assert row1.parent == table
    assert row2.parent == table
    assert list_node.parent == root
    assert list_item1.parent == list_node
    assert list_item2.parent == list_node
    assert nested_list_item.parent == list_item2

    # Check root id
    assert section1.root_id == root.id_
    assert section2.root_id == root.id_
    assert text1.root_id == root.id_
    assert text2.root_id == root.id_
    assert subsection.root_id == root.id_
    assert subtext.root_id == root.id_
    assert table.root_id == root.id_
    assert row1.root_id == root.id_
    assert row2.root_id == root.id_
    assert list_node.root_id == root.id_
    assert list_item1.root_id == root.id_
    assert list_item2.root_id == root.id_
    assert nested_list_item.root_id == root.id_

    # Check depth
    assert root.depth == 0
    assert section1.depth == 1
    assert section2.depth == 1
    assert text1.depth == 2
    assert text2.depth == 2
    assert subsection.depth == 2
    assert subtext.depth == 3
    assert table.depth == 2
    assert row1.depth == 3
    assert row2.depth == 3
    assert list_node.depth == 1
    assert list_item1.depth == 2
    assert list_item2.depth == 2
    assert nested_list_item.depth == 3

    # Check order index
    assert root.idx == 0
    assert section1.idx == 0
    assert section2.idx == 1
    assert text1.idx == 0
    assert text2.idx == 0
    assert subsection.idx == 1
    assert subtext.idx == 0
    assert table.idx == 2
    assert row1.idx == 0
    assert row2.idx == 1
    assert list_node.idx == 2
    assert list_item1.idx == 0
    assert list_item2.idx == 1
    assert nested_list_item.idx == 0

    # Check height
    assert root.height == 3
    assert section1.height == 1
    assert section2.height == 2
    assert text1.height == 0
    assert text2.height == 0
    assert subsection.height == 1
    assert subtext.height == 0
    assert table.height == 1
    assert row1.height == 0
    assert row2.height == 0
    assert list_node.height == 2
    assert list_item1.height == 0
    assert list_item2.height == 1
    assert nested_list_item.height == 0

    # Check absolute index
    for i, node in enumerate(root.flatten()):
        assert node.abs_idx == i


def test_insert_node() -> None:
    root = DocumentRootNode()
    section1 = SectionNode(text="Introduction")
    root.add_child(section1)
    text1 = TextNode(text="This is the introduction.")
    section1.add_child(text1)
    section2 = SectionNode(text="Body")
    root.add_child(section2)
    text2 = TextNode(text="This is the body text.")
    section2.add_child(text2)
    subsection = SectionNode(text="Subsection")
    section2.add_child(subsection)
    subtext = TextNode(text="This is the subsection text.")
    subsection.add_child(subtext)
    root.update_references()

    # Case 1. Insert element at the end
    new_text = TextNode(text="This is the new text.")
    root.insert_child(2, new_text, update_references=True)
    sorted_nodes = list(root.flatten())

    assert len(root.children) == 3
    assert root.children[2] == new_text

    assert root.height == 3
    assert new_text.height == 0
    assert new_text.abs_idx == sorted_nodes.index(new_text)

    # Case 2. Insert element at the beginning
    new_section = SectionNode(text="New Section")
    root.insert_child(0, new_section)
    root.update_references()
    sorted_nodes = list(root.flatten())

    assert len(root.children) == 4
    assert root.children[0] == new_section

    assert root.height == 3
    assert new_section.height == 0
    assert new_section.abs_idx == sorted_nodes.index(new_section)
    for i, child in enumerate(root.children):
        assert child.idx == i

    # Case 3. Insert element in the middle
    new_subsection = SectionNode(text="New Subsection")
    root.insert_child(2, new_subsection, update_references=True)
    sorted_nodes = list(root.flatten())

    assert len(root.children) == 5
    assert root.children[2] == new_subsection

    assert root.height == 3
    assert new_subsection.height == 0
    assert new_subsection.abs_idx == sorted_nodes.index(new_subsection)
    for i, child in enumerate(root.children):
        assert child.idx == i

    # Case 4. Negative index
    new_text = TextNode(text="This is the new text.")
    root.insert_child(-1, new_text)
    root.update_references()
    sorted_nodes = list(root.flatten())

    assert len(root.children) == 6
    assert root.children[-2] == new_text

    assert root.height == 3
    assert new_text.height == 0
    assert new_text.abs_idx == sorted_nodes.index(new_text)
    for i, child in enumerate(root.children):
        assert child.idx == i

    # Case 5.1 Index out of bounds (negative)
    new_text = TextNode(text="This is the new text.")
    with pytest.raises(IndexError):
        root.insert_child(-100, new_text)

    # Case 5.2 Index out of bounds (positive)
    new_text = TextNode(text="This is the new text.")
    with pytest.raises(IndexError):
        root.insert_child(100, new_text)


def test_partial_node() -> None:
    def compare_nodes(node1: Any, node2: Any) -> None:
        assert node1.id_ == node2.id_
        assert node1.metadata == node2.metadata
        assert node1.type == node2.get_type()
        assert node1.idx == node2.idx
        assert node1.depth == node2.depth
        assert node1.parent_id == node2.parent_id
        assert node1.root_id == node2.root_id
        assert node1.hash == node2.hash
        assert node1.height == node2.height
        assert node1.abs_idx == node2.abs_idx
        assert len(node1.children) == len(node2.children)
        for child1, child2 in zip(node1.children, node2.children, strict=False):
            compare_nodes(child1, child2)

    # Case 1. Create a partial node from a tree node
    root = DocumentRootNode()
    section1 = SectionNode(text="Introduction")
    root.add_child(section1)
    text1 = TextNode(text="This is the introduction.")
    section1.add_child(text1)
    section2 = SectionNode(text="Body")
    root.add_child(section2)

    partial_node = PartialNode.from_node(root)
    compare_nodes(partial_node, root)

    # Case 2. Create a partial node from a dictionary
    dict = partial_node.dict()

    partial_node2 = PartialNode.from_partial_dict(dict)
    compare_nodes(partial_node2, root)