/test/registered/unit/tokenizer/test_tiktoken_tokenizer.py

"""Unit tests for tiktoken_tokenizer — no server, no model loading."""

import unittest
from unittest.mock import MagicMock, patch

from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=7, suite="base-a-test-cpu")

from sglang.srt.tokenizer.tiktoken_tokenizer import (
    CONTROL_TOKEN_TEXTS,
    DEFAULT_CONTROL_TOKENS,
    DEFAULT_SPECIAL_TOKENS,
    EOS,
    PAD,
    RESERVED_TOKEN_TEXTS,
    SEP,
    TiktokenProcessor,
    TiktokenTokenizer,
)


class TestConstants(CustomTestCase):
    def test_reserved_token_count(self):
        self.assertEqual(len(RESERVED_TOKEN_TEXTS), 125)

    def test_reserved_token_format(self):
        self.assertEqual(RESERVED_TOKEN_TEXTS[0], "<|reserved_3|>")
        self.assertEqual(RESERVED_TOKEN_TEXTS[-1], "<|reserved_127|>")

    def test_control_token_count(self):
        self.assertEqual(len(CONTROL_TOKEN_TEXTS), 704)

    def test_control_token_format(self):
        self.assertEqual(CONTROL_TOKEN_TEXTS[0], "<|control1|>")
        self.assertEqual(CONTROL_TOKEN_TEXTS[-1], "<|control704|>")

    def test_default_special_tokens_contains_all(self):
        self.assertIn(PAD, DEFAULT_SPECIAL_TOKENS)
        self.assertIn(EOS, DEFAULT_SPECIAL_TOKENS)
        self.assertIn(SEP, DEFAULT_SPECIAL_TOKENS)

    def test_default_control_tokens_values(self):
        # Note: "sep" maps to EOS and "eos" maps to SEP in the source code
        self.assertEqual(DEFAULT_CONTROL_TOKENS["pad"], PAD)
        self.assertEqual(DEFAULT_CONTROL_TOKENS["sep"], EOS)
        self.assertEqual(DEFAULT_CONTROL_TOKENS["eos"], SEP)


class TestTiktokenProcessor(CustomTestCase):
    def setUp(self):
        tokenizer_patcher = patch(
            "sglang.srt.tokenizer.tiktoken_tokenizer.TiktokenTokenizer"
        )
        tokenizer_patcher.start()
        self.addCleanup(tokenizer_patcher.stop)
        self.processor = TiktokenProcessor(name="dummy")

    def test_image_processor_wraps_image_in_list(self):
        image = "fake_image_data"
        result = self.processor.image_processor(image)
        self.assertEqual(result["pixel_values"], [image])

    def test_image_processor_with_none(self):
        result = self.processor.image_processor(None)
        self.assertEqual(result["pixel_values"], [None])


class TestTiktokenTokenizer(CustomTestCase):
    def setUp(self):
        from jinja2 import Template

        self.tok = TiktokenTokenizer.__new__(TiktokenTokenizer)
        self.mock_tokenizer = MagicMock()
        self.tok.tokenizer = self.mock_tokenizer
        self.tok.chat_template = "dummy"
        self.tok.chat_template_jinja = Template(
            "{% for message in messages %}"
            "{{ message['role'] }}: {{ message['content'] }}"
            "{% endfor %}"
            "{% if add_generation_prompt %}assistant:{% endif %}"
        )

    def test_batch_decode_flat_list_wraps_each(self):
        self.mock_tokenizer.decode_batch.return_value = ["a", "b"]
        self.tok.batch_decode([1, 2])
        self.mock_tokenizer.decode_batch.assert_called_once_with([[1], [2]])

    def test_call_returns_input_ids(self):
        self.mock_tokenizer.encode.return_value = [1, 2, 3]
        result = self.tok(["hello", "world"])
        self.assertIn("input_ids", result)
        self.assertEqual(len(result["input_ids"]), 2)

    def test_apply_chat_template_no_tokenize(self):
        messages = [{"role": "user", "content": "hello"}]
        result = self.tok.apply_chat_template(
            messages=messages,
            tokenize=False,
            add_generation_prompt=False,
        )
        self.assertIsInstance(result, str)
        self.assertIn("hello", result)

    def test_apply_chat_template_with_tokenize(self):
        self.mock_tokenizer.encode.return_value = [1, 2, 3]
        messages = [{"role": "user", "content": "hello"}]
        result = self.tok.apply_chat_template(
            messages=messages,
            tokenize=True,
            add_generation_prompt=False,
        )
        self.assertIsInstance(result, list)


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