castle-batch-v1/tests/test_runner.py

import itertools
import threading
import time
import pytest
from castle_pipeline.runner import bounded_map, Checkpoints, clip_windows, memory_estimate, source_lease


def test_lazy_scheduler_does_not_exhaust_video_iterator():
    produced = []
    release = threading.Event()
    def inputs():
        for i in range(1000):
            produced.append(i)
            yield i
    def worker(i):
        assert release.wait(3)
        return i
    result = []
    thread = threading.Thread(target=lambda: result.append(next(bounded_map(inputs(), worker, 3))))
    thread.start()
    time.sleep(.1)
    assert len(produced) == 3
    release.set()
    thread.join(4)
    assert not thread.is_alive()
    assert len(produced) == 3


def test_tail_clip_is_not_dropped_and_no_zero_length_window():
    assert list(clip_windows(65, 30)) == [(0, 0., 30.), (1, 30., 30.), (2, 60., 5.)]
    assert list(clip_windows(60, 30)) == [(0, 0., 30.), (1, 30., 30.)]
    with pytest.raises(ValueError):
        list(clip_windows(60, 0))


def test_checkpoint_fingerprint_prevents_stale_resume(tmp_path):
    checkpoint = Checkpoints(tmp_path, 'fingerprint-a')
    checkpoint.save('audio', {'data': {'utterances': []}})
    assert checkpoint.load('audio')['data'] == {'utterances': []}
    assert Checkpoints(tmp_path, 'fingerprint-b').load('audio') is None
    (tmp_path / 'audio.json').write_text('{partial', encoding='utf-8')
    assert checkpoint.load('audio') is None


def test_memory_bound_depends_on_workers_not_source_length():
    one = memory_estimate(source_seconds=3600, workers=3)
    ten = memory_estimate(source_seconds=36000, workers=3)
    assert one['bounded_working_set_mib_estimate'] == ten['bounded_working_set_mib_estimate']
    assert ten['all_raw_1fps_gib'] == pytest.approx(one['all_raw_1fps_gib'] * 10)


def test_two_jobs_cannot_write_the_same_source_run(tmp_path):
    with source_lease(tmp_path):
        with pytest.raises(RuntimeError):
            with source_lease(tmp_path):
                pytest.fail('Overlapping run acquired the same output lease')
    assert not (tmp_path/'run.lock').exists()