40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
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()
|