test_generation.py

"""Regression tests for generation configuration."""

from __future__ import annotations

import unittest

from generation import (
    ensure_bos_token,
    gpu_duration_seconds,
    head_tail_token_counts,
    merge_eos_token_ids,
)


class GenerationConfigTests(unittest.TestCase):
    def test_adds_missing_bos_token(self) -> None:
        self.assertEqual(
            ensure_bos_token("<|im_start|>user\nHello", "<|begin_of_text|>"),
            "<|begin_of_text|><|im_start|>user\nHello",
        )

    def test_does_not_duplicate_bos_token(self) -> None:
        prompt = "<|begin_of_text|><|im_start|>user\nHello"
        self.assertEqual(ensure_bos_token(prompt, "<|begin_of_text|>"), prompt)

    def test_preserves_all_model_terminators(self) -> None:
        self.assertEqual(
            merge_eos_token_ids([128001, 128008, 128009, 128256], 128256),
            [128001, 128008, 128009, 128256],
        )

    def test_adds_tokenizer_fallback(self) -> None:
        self.assertEqual(merge_eos_token_ids(128001, 128256), [128001, 128256])

    def test_returns_none_without_valid_ids(self) -> None:
        self.assertIsNone(merge_eos_token_ids(None, None))

    def test_short_prompt_is_not_split(self) -> None:
        self.assertEqual(head_tail_token_counts(100, 200, 50), (100, 0))

    def test_large_prompt_preserves_prefix_and_maximizes_tail(self) -> None:
        self.assertEqual(
            head_tail_token_counts(30_000, 16_000, 4_096),
            (4_096, 11_904),
        )

    def test_tiny_budget_still_preserves_latest_token(self) -> None:
        self.assertEqual(head_tail_token_counts(100, 1, 4_096), (0, 1))

    def test_gpu_duration_accounts_for_large_context_prefill(self) -> None:
        self.assertEqual(
            gpu_duration_seconds(90_000, 32, 16_384),
            110,
        )

    def test_gpu_duration_keeps_small_requests_efficient(self) -> None:
        self.assertEqual(gpu_duration_seconds(300, 32, 16_384), 30)

    def test_gpu_duration_is_capped(self) -> None:
        self.assertEqual(
            gpu_duration_seconds(1_000_000, 1_024, 16_384),
            120,
        )


if __name__ == "__main__":
    unittest.main()