175 lines
7.4 KiB
Python
175 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from yuxi.channel.constants import CHANNEL_COMPENSATE_LOCK_KEY, CHANNEL_MAX_COMPENSATE_ATTEMPTS
|
|
from yuxi.channel.outbound.retry import _do_compensate, _renew_compensate_lock, compensate_channel_messages
|
|
|
|
|
|
@pytest.fixture
|
|
def redis_mock():
|
|
return MagicMock(
|
|
set=AsyncMock(return_value=True),
|
|
expire=AsyncMock(),
|
|
delete=AsyncMock(),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def message_mock():
|
|
msg = MagicMock()
|
|
msg.id = 1
|
|
msg.conversation_id = 10
|
|
msg.channel_metadata = {}
|
|
msg.delivery_status = "pending"
|
|
return msg
|
|
|
|
|
|
@pytest.fixture
|
|
def conversation_mock():
|
|
conv = MagicMock()
|
|
conv.id = 10
|
|
conv.channel_session_id = "sk-1"
|
|
return conv
|
|
|
|
|
|
@pytest.fixture
|
|
def pg_manager_mock(monkeypatch):
|
|
session_mock = MagicMock(commit=AsyncMock())
|
|
|
|
@asynccontextmanager
|
|
async def ctx():
|
|
yield session_mock
|
|
|
|
mock = MagicMock()
|
|
mock.get_async_session_context = ctx
|
|
monkeypatch.setattr("yuxi.channel.outbound.retry.pg_manager", mock)
|
|
return mock, session_mock
|
|
|
|
|
|
async def test_compensate_acquires_lock_and_releases(redis_mock):
|
|
with patch("yuxi.channel.outbound.retry.get_redis_client", AsyncMock(return_value=redis_mock)):
|
|
with patch("yuxi.channel.outbound.retry._do_compensate", AsyncMock()) as mock_do:
|
|
await compensate_channel_messages()
|
|
|
|
redis_mock.set.assert_awaited_once_with(CHANNEL_COMPENSATE_LOCK_KEY, "1", nx=True, ex=300)
|
|
mock_do.assert_awaited_once()
|
|
redis_mock.delete.assert_awaited_once_with(CHANNEL_COMPENSATE_LOCK_KEY)
|
|
|
|
|
|
async def test_compensate_returns_when_lock_not_acquired(redis_mock):
|
|
redis_mock.set.return_value = None
|
|
with patch("yuxi.channel.outbound.retry.get_redis_client", AsyncMock(return_value=redis_mock)):
|
|
with patch("yuxi.channel.outbound.retry._do_compensate", AsyncMock()) as mock_do:
|
|
await compensate_channel_messages()
|
|
|
|
mock_do.assert_not_awaited()
|
|
redis_mock.delete.assert_not_awaited()
|
|
|
|
|
|
async def test_compensate_releases_lock_even_if_do_compensate_fails(redis_mock):
|
|
with patch("yuxi.channel.outbound.retry.get_redis_client", AsyncMock(return_value=redis_mock)):
|
|
with patch("yuxi.channel.outbound.retry._do_compensate", AsyncMock(side_effect=RuntimeError("boom"))):
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
await compensate_channel_messages()
|
|
|
|
redis_mock.delete.assert_awaited_once_with(CHANNEL_COMPENSATE_LOCK_KEY)
|
|
|
|
|
|
async def test_renew_compensate_lock_expires_until_stopped(redis_mock):
|
|
stop_event = MagicMock()
|
|
stop_event.is_set = MagicMock(side_effect=[False, False, True])
|
|
stop_event.wait = AsyncMock(side_effect=[TimeoutError, TimeoutError])
|
|
|
|
with patch("yuxi.channel.outbound.retry.CHANNEL_PROCESSING_LOCK_RENEW_INTERVAL_SECONDS", 0):
|
|
await _renew_compensate_lock(redis_mock, "lock-key", stop_event)
|
|
|
|
assert redis_mock.expire.await_count == 2
|
|
redis_mock.expire.assert_awaited_with("lock-key", 300)
|
|
|
|
|
|
async def test_do_compensate_publishes_and_increments_attempts(pg_manager_mock, message_mock, conversation_mock):
|
|
_, session_mock = pg_manager_mock
|
|
message_repo = MagicMock()
|
|
conversation_repo = MagicMock()
|
|
message_repo.list_pending_channel_messages = AsyncMock(return_value=[message_mock])
|
|
conversation_repo.get_conversation_by_id = AsyncMock(return_value=conversation_mock)
|
|
|
|
with patch("yuxi.channel.outbound.retry.MessageRepository", return_value=message_repo):
|
|
with patch("yuxi.channel.outbound.retry.ConversationRepository", return_value=conversation_repo):
|
|
with patch(
|
|
"yuxi.channel.outbound.retry.publish_channel_message_event", AsyncMock(return_value="evt-1")
|
|
) as mock_publish:
|
|
await _do_compensate()
|
|
|
|
message_repo.list_pending_channel_messages.assert_awaited_once()
|
|
conversation_repo.get_conversation_by_id.assert_awaited_once_with(10)
|
|
assert message_mock.channel_metadata["compensate_attempts"] == 1
|
|
session_mock.commit.assert_awaited()
|
|
mock_publish.assert_awaited_once_with(message=message_mock, conversation=conversation_mock, session_key="sk-1")
|
|
|
|
|
|
async def test_do_compensate_skips_missing_conversation(pg_manager_mock, message_mock):
|
|
_, session_mock = pg_manager_mock
|
|
message_repo = MagicMock()
|
|
conversation_repo = MagicMock()
|
|
message_repo.list_pending_channel_messages = AsyncMock(return_value=[message_mock])
|
|
conversation_repo.get_conversation_by_id = AsyncMock(return_value=None)
|
|
|
|
with patch("yuxi.channel.outbound.retry.MessageRepository", return_value=message_repo):
|
|
with patch("yuxi.channel.outbound.retry.ConversationRepository", return_value=conversation_repo):
|
|
with patch("yuxi.channel.outbound.retry.publish_channel_message_event", AsyncMock()) as mock_publish:
|
|
await _do_compensate()
|
|
|
|
mock_publish.assert_not_awaited()
|
|
|
|
|
|
async def test_do_compensate_skips_missing_session_key(pg_manager_mock, message_mock, conversation_mock):
|
|
conversation_mock.channel_session_id = None
|
|
message_repo = MagicMock()
|
|
conversation_repo = MagicMock()
|
|
message_repo.list_pending_channel_messages = AsyncMock(return_value=[message_mock])
|
|
conversation_repo.get_conversation_by_id = AsyncMock(return_value=conversation_mock)
|
|
|
|
with patch("yuxi.channel.outbound.retry.MessageRepository", return_value=message_repo):
|
|
with patch("yuxi.channel.outbound.retry.ConversationRepository", return_value=conversation_repo):
|
|
with patch("yuxi.channel.outbound.retry.publish_channel_message_event", AsyncMock()) as mock_publish:
|
|
await _do_compensate()
|
|
|
|
mock_publish.assert_not_awaited()
|
|
|
|
|
|
async def test_do_compensate_marks_dead_letter_after_max_attempts(pg_manager_mock, message_mock, conversation_mock):
|
|
_, session_mock = pg_manager_mock
|
|
message_mock.channel_metadata = {"compensate_attempts": CHANNEL_MAX_COMPENSATE_ATTEMPTS}
|
|
message_repo = MagicMock()
|
|
conversation_repo = MagicMock()
|
|
message_repo.list_pending_channel_messages = AsyncMock(return_value=[message_mock])
|
|
conversation_repo.get_conversation_by_id = AsyncMock(return_value=conversation_mock)
|
|
|
|
with patch("yuxi.channel.outbound.retry.MessageRepository", return_value=message_repo):
|
|
with patch("yuxi.channel.outbound.retry.ConversationRepository", return_value=conversation_repo):
|
|
with patch("yuxi.channel.outbound.retry.publish_channel_message_event", AsyncMock()) as mock_publish:
|
|
await _do_compensate()
|
|
|
|
assert message_mock.delivery_status == "dead_letter"
|
|
mock_publish.assert_not_awaited()
|
|
|
|
|
|
async def test_do_compensate_catches_publish_exception(pg_manager_mock, message_mock, conversation_mock):
|
|
message_repo = MagicMock()
|
|
conversation_repo = MagicMock()
|
|
message_repo.list_pending_channel_messages = AsyncMock(return_value=[message_mock])
|
|
conversation_repo.get_conversation_by_id = AsyncMock(return_value=conversation_mock)
|
|
|
|
with patch("yuxi.channel.outbound.retry.MessageRepository", return_value=message_repo):
|
|
with patch("yuxi.channel.outbound.retry.ConversationRepository", return_value=conversation_repo):
|
|
with patch(
|
|
"yuxi.channel.outbound.retry.publish_channel_message_event",
|
|
AsyncMock(side_effect=RuntimeError("pub failed")),
|
|
):
|
|
# 不应抛出异常
|
|
await _do_compensate()
|