PFMBench/src/data/esm/utils/misc_test.py

"""Tests for misc.py"""

from src.data.esm.utils.misc import merge_annotations
from src.data.esm.utils.types import FunctionAnnotation


def test_merge_annotations():
    merged = merge_annotations(
        [
            FunctionAnnotation("a", start=1, end=10),
            FunctionAnnotation("b", start=5, end=15),
            FunctionAnnotation("a", start=10, end=20),
            FunctionAnnotation("b", start=2, end=6),
            FunctionAnnotation("c", start=4, end=10),
        ]
    )
    assert len(merged) == 3
    assert FunctionAnnotation("a", start=1, end=20) in merged
    assert FunctionAnnotation("b", start=2, end=15) in merged
    assert FunctionAnnotation("c", start=4, end=10) in merged


def test_merge_annotations_gap():
    merged = merge_annotations(
        [
            FunctionAnnotation("a", start=1, end=10),
            FunctionAnnotation("a", start=13, end=20),  # gap is 2
            FunctionAnnotation("a", start=24, end=30),
        ],
        merge_gap_max=2,
    )

    assert len(merged) == 2
    assert FunctionAnnotation("a", 1, 20) in merged
    assert FunctionAnnotation("a", 24, 30) in merged