import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
from uuid import uuid4

from core.storage.message_index import allocate_message_idx


class MessageIndexTests(unittest.TestCase):
    def test_allocator_uses_persisted_counter_when_ahead(self) -> None:
        tid = uuid4()
        task = SimpleNamespace(task_id=tid, next_message_idx=8)
        session = MagicMock()
        session.execute.return_value.scalar_one.return_value = 5

        self.assertEqual(allocate_message_idx(session, tid, locked_task=task), 8)
        self.assertEqual(task.next_message_idx, 9)

    def test_allocator_catches_up_with_legacy_writer(self) -> None:
        tid = uuid4()
        task = SimpleNamespace(task_id=tid, next_message_idx=3)
        session = MagicMock()
        session.execute.return_value.scalar_one.return_value = 6

        self.assertEqual(allocate_message_idx(session, tid, locked_task=task), 7)
        self.assertEqual(task.next_message_idx, 8)

    def test_empty_task_starts_at_zero(self) -> None:
        tid = uuid4()
        task = SimpleNamespace(task_id=tid, next_message_idx=0)
        session = MagicMock()
        session.execute.return_value.scalar_one.return_value = None

        self.assertEqual(allocate_message_idx(session, tid, locked_task=task), 0)
        self.assertEqual(task.next_message_idx, 1)


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