tests/test_collators.py

"""
tests/test_collators.py
Unit tests for data collators.
"""

import sys
sys.path.insert(0, "src")

import pytest


def test_find_last_subseq():
    from ankahi.data.collators import _find_last_subseq
    seq = [1, 2, 3, 1, 2, 3, 4]
    assert _find_last_subseq(seq, [1, 2]) == 3
    assert _find_last_subseq(seq, [3, 4]) == 5
    assert _find_last_subseq(seq, [9]) is None
    assert _find_last_subseq([], [1]) is None


def test_find_last_subseq_single():
    from ankahi.data.collators import _find_last_subseq
    assert _find_last_subseq([5, 5, 5], [5]) == 2
    assert _find_last_subseq([1, 2, 3], [1, 2, 3]) == 0


def test_pack_texts_empty():
    from ankahi.data.collators import pack_texts

    class MockTokenizer:
        def encode(self, text, add_special_tokens=True):
            return list(range(len(text.split())))
        def decode(self, tokens):
            return " ".join(str(t) for t in tokens)

    packed = pack_texts([], MockTokenizer(), max_length=512)
    assert packed == []


def test_pack_texts_basic():
    from ankahi.data.collators import pack_texts

    class MockTokenizer:
        def encode(self, text, add_special_tokens=True):
            return [0] * 100  # 100 tokens per text
        def decode(self, tokens):
            return " ".join(str(t) for t in tokens)

    texts = ["text " * 20] * 5
    # Each text = 100 tokens, max = 512, should pack ~5 into each chunk
    packed = pack_texts(texts, MockTokenizer(), max_length=512)
    assert len(packed) >= 1