import unittest

from sglang.srt.layers.attention.nsa.utils import (
    _get_in_seq_last_token_owner_and_offset,
)
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=1, suite="stage-a-test-cpu")


class TestNSAInSeqCPUtils(unittest.TestCase):
    def test_last_token_owner_uses_actual_token_count_when_batch_is_padded(self):
        # Padded prefill can have 64 model tokens while the real prompt has only
        # 11 tokens.  In in-seq split with cp_size=8, the real last token is in
        # segment 2, not in rank 0's trailing padded segment.
        split_list = [4] * 16

        owner, local_offset = _get_in_seq_last_token_owner_and_offset(
            split_list=split_list,
            cp_size=8,
            actual_token_count=11,
        )

        self.assertEqual(owner, 2)
        self.assertEqual(local_offset, 2)

    def test_last_token_owner_keeps_existing_unpadded_fast_path_location(self):
        split_list = [4] * 16

        owner, local_offset = _get_in_seq_last_token_owner_and_offset(
            split_list=split_list,
            cp_size=8,
            actual_token_count=64,
        )

        self.assertEqual(owner, 0)
        self.assertEqual(local_offset, 7)


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