ComfyUI-ShotSplitter/tests/test_ensemble.py

"""Unit tests for ensemble merge logic — pure functions, no video I/O."""

import pytest

from detectors.ensemble import (
    union_boundaries,
    apply_min_shot_filter,
    boundaries_to_intervals,
)


def test_union_within_window_merges_nearby():
    a = [10, 50, 100]
    b = [11, 48, 200]
    # window=3: 10/11 merge -> 10; 50/48 merge -> 48; 100 alone; 200 alone
    result = union_boundaries([a, b], window=3)
    assert result == [10, 48, 100, 200]


def test_union_outside_window_keeps_both():
    a = [10, 50]
    b = [20, 60]
    # window=3: nothing merges (distances are 10)
    result = union_boundaries([a, b], window=3)
    assert result == [10, 20, 50, 60]


def test_union_empty_inputs_returns_empty():
    assert union_boundaries([[], []], window=3) == []
    assert union_boundaries([], window=3) == []


def test_boundaries_to_intervals_adds_sentinels():
    result = boundaries_to_intervals([100, 200], total_frames=300)
    assert result == [(0, 100), (100, 200), (200, 300)]


def test_boundaries_to_intervals_no_cuts_gives_single_shot():
    result = boundaries_to_intervals([], total_frames=500)
    assert result == [(0, 500)]


def test_min_shot_filter_merges_short_with_right_neighbor():
    # fps=30, min=0.5s -> 15 frames
    intervals = [(0, 100), (100, 105), (105, 300)]
    # (100,105) is 5 frames, < 15; merge with right -> (100, 300)
    result = apply_min_shot_filter(intervals, fps=30.0, min_shot_seconds=0.5)
    assert result == [(0, 100), (100, 300)]


def test_min_shot_filter_last_short_merges_with_left():
    intervals = [(0, 100), (100, 300), (300, 305)]
    # (300,305) is at the end; no right neighbor -> merge with left
    result = apply_min_shot_filter(intervals, fps=30.0, min_shot_seconds=0.5)
    assert result == [(0, 100), (100, 305)]


def test_min_shot_filter_keeps_all_when_above_threshold():
    intervals = [(0, 100), (100, 200), (200, 300)]
    result = apply_min_shot_filter(intervals, fps=30.0, min_shot_seconds=0.5)
    assert result == intervals