mirror of
https://github.com/langgenius/dify.git
synced 2026-03-13 02:57:41 +08:00
1060 lines
37 KiB
Python
1060 lines
37 KiB
Python
from datetime import datetime
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from libs.infinite_scroll_pagination import InfiniteScrollPagination
|
|
from models.model import App, AppMode, EndUser, Message
|
|
from services.errors.message import (
|
|
FirstMessageNotExistsError,
|
|
LastMessageNotExistsError,
|
|
MessageNotExistsError,
|
|
SuggestedQuestionsAfterAnswerDisabledError,
|
|
)
|
|
from services.message_service import MessageService, attach_message_extra_contents
|
|
|
|
|
|
class TestMessageServiceFactory:
|
|
"""Factory class for creating test data and mock objects for message service tests."""
|
|
|
|
@staticmethod
|
|
def create_app_mock(
|
|
app_id: str = "app-123",
|
|
mode: str = AppMode.ADVANCED_CHAT.value,
|
|
name: str = "Test App",
|
|
) -> MagicMock:
|
|
"""Create a mock App object."""
|
|
app = MagicMock(spec=App)
|
|
app.id = app_id
|
|
app.mode = mode
|
|
app.name = name
|
|
return app
|
|
|
|
@staticmethod
|
|
def create_end_user_mock(
|
|
user_id: str = "user-456",
|
|
session_id: str = "session-789",
|
|
) -> MagicMock:
|
|
"""Create a mock EndUser object."""
|
|
user = MagicMock(spec=EndUser)
|
|
user.id = user_id
|
|
user.session_id = session_id
|
|
return user
|
|
|
|
@staticmethod
|
|
def create_conversation_mock(
|
|
conversation_id: str = "conv-001",
|
|
app_id: str = "app-123",
|
|
) -> MagicMock:
|
|
"""Create a mock Conversation object."""
|
|
conversation = MagicMock()
|
|
conversation.id = conversation_id
|
|
conversation.app_id = app_id
|
|
return conversation
|
|
|
|
@staticmethod
|
|
def create_message_mock(
|
|
message_id: str = "msg-001",
|
|
conversation_id: str = "conv-001",
|
|
query: str = "What is AI?",
|
|
answer: str = "AI stands for Artificial Intelligence.",
|
|
created_at: datetime | None = None,
|
|
) -> MagicMock:
|
|
"""Create a mock Message object."""
|
|
message = MagicMock(spec=Message)
|
|
message.id = message_id
|
|
message.conversation_id = conversation_id
|
|
message.query = query
|
|
message.answer = answer
|
|
message.created_at = created_at or datetime.now()
|
|
return message
|
|
|
|
|
|
class TestMessageServicePaginationByFirstId:
|
|
"""
|
|
Unit tests for MessageService.pagination_by_first_id method.
|
|
|
|
This test suite covers:
|
|
- Basic pagination with and without first_id
|
|
- Order handling (asc/desc)
|
|
- Edge cases (no user, no conversation, invalid first_id)
|
|
- Has_more flag logic
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def factory(self):
|
|
"""Provide test data factory."""
|
|
return TestMessageServiceFactory()
|
|
|
|
# Test 01: No user provided
|
|
def test_pagination_by_first_id_no_user(self, factory):
|
|
"""Test pagination returns empty result when no user is provided."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=None,
|
|
conversation_id="conv-001",
|
|
first_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert isinstance(result, InfiniteScrollPagination)
|
|
assert result.data == []
|
|
assert result.limit == 10
|
|
assert result.has_more is False
|
|
|
|
# Test 02: No conversation_id provided
|
|
def test_pagination_by_first_id_no_conversation(self, factory):
|
|
"""Test pagination returns empty result when no conversation_id is provided."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="",
|
|
first_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert isinstance(result, InfiniteScrollPagination)
|
|
assert result.data == []
|
|
assert result.limit == 10
|
|
assert result.has_more is False
|
|
|
|
# Test 03: Basic pagination without first_id (desc order)
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_pagination_by_first_id_without_first_id_desc(self, mock_conversation_service, mock_db, factory):
|
|
"""Test basic pagination without first_id in descending order."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock()
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
# Create 5 messages
|
|
messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="conv-001",
|
|
first_id=None,
|
|
limit=10,
|
|
order="desc",
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 5
|
|
assert result.has_more is False
|
|
assert result.limit == 10
|
|
# Messages should remain in desc order (not reversed)
|
|
assert result.data[0].id == "msg-000"
|
|
|
|
# Test 04: Basic pagination without first_id (asc order)
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_pagination_by_first_id_without_first_id_asc(self, mock_conversation_service, mock_db, factory):
|
|
"""Test basic pagination without first_id in ascending order."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock()
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
# Create 5 messages (returned in desc order from DB)
|
|
messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, 4 - i), # Descending timestamps
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="conv-001",
|
|
first_id=None,
|
|
limit=10,
|
|
order="asc",
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 5
|
|
assert result.has_more is False
|
|
# Messages should be reversed to asc order
|
|
assert result.data[0].id == "msg-004"
|
|
assert result.data[4].id == "msg-000"
|
|
|
|
# Test 05: Pagination with first_id
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_pagination_by_first_id_with_first_id(self, mock_conversation_service, mock_db, factory):
|
|
"""Test pagination with first_id to get messages before a specific message."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock()
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
first_message = factory.create_message_mock(
|
|
message_id="msg-005",
|
|
created_at=datetime(2024, 1, 1, 12, 5),
|
|
)
|
|
|
|
# Messages before first_message
|
|
history_messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
# Setup query mocks
|
|
mock_query_first = MagicMock()
|
|
mock_query_history = MagicMock()
|
|
|
|
query_calls = []
|
|
|
|
def query_side_effect(*args):
|
|
if args[0] == Message:
|
|
query_calls.append(args)
|
|
if len(query_calls) == 1:
|
|
return mock_query_first
|
|
else:
|
|
return mock_query_history
|
|
|
|
mock_db.session.query.side_effect = [mock_query_first, mock_query_history]
|
|
|
|
# Setup first message query
|
|
mock_query_first.where.return_value = mock_query_first
|
|
mock_query_first.first.return_value = first_message
|
|
|
|
# Setup history messages query
|
|
mock_query_history.where.return_value = mock_query_history
|
|
mock_query_history.order_by.return_value = mock_query_history
|
|
mock_query_history.limit.return_value = mock_query_history
|
|
mock_query_history.all.return_value = history_messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="conv-001",
|
|
first_id="msg-005",
|
|
limit=10,
|
|
order="desc",
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 5
|
|
assert result.has_more is False
|
|
mock_query_first.where.assert_called_once()
|
|
mock_query_history.where.assert_called_once()
|
|
|
|
# Test 06: First message not found
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_pagination_by_first_id_first_message_not_exists(self, mock_conversation_service, mock_db, factory):
|
|
"""Test error handling when first_id doesn't exist."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock()
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.first.return_value = None # Message not found
|
|
|
|
# Act & Assert
|
|
with pytest.raises(FirstMessageNotExistsError):
|
|
MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="conv-001",
|
|
first_id="nonexistent-msg",
|
|
limit=10,
|
|
)
|
|
|
|
# Test 07: Has_more flag when results exceed limit
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_pagination_by_first_id_has_more_true(self, mock_conversation_service, mock_db, factory):
|
|
"""Test has_more flag is True when results exceed limit."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock()
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
# Create limit+1 messages (11 messages for limit=10)
|
|
messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(11)
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="conv-001",
|
|
first_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 10 # Last message trimmed
|
|
assert result.has_more is True
|
|
assert result.limit == 10
|
|
|
|
# Test 08: Empty conversation
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_pagination_by_first_id_empty_conversation(self, mock_conversation_service, mock_db, factory):
|
|
"""Test pagination with conversation that has no messages."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock()
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = []
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_first_id(
|
|
app_model=app,
|
|
user=user,
|
|
conversation_id="conv-001",
|
|
first_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 0
|
|
assert result.has_more is False
|
|
assert result.limit == 10
|
|
|
|
|
|
class TestMessageServicePaginationByLastId:
|
|
"""
|
|
Unit tests for MessageService.pagination_by_last_id method.
|
|
|
|
This test suite covers:
|
|
- Basic pagination with and without last_id
|
|
- Conversation filtering
|
|
- Include_ids filtering
|
|
- Edge cases (no user, invalid last_id)
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def factory(self):
|
|
"""Provide test data factory."""
|
|
return TestMessageServiceFactory()
|
|
|
|
# Test 09: No user provided
|
|
def test_pagination_by_last_id_no_user(self, factory):
|
|
"""Test pagination returns empty result when no user is provided."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=None,
|
|
last_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert isinstance(result, InfiniteScrollPagination)
|
|
assert result.data == []
|
|
assert result.limit == 10
|
|
assert result.has_more is False
|
|
|
|
# Test 10: Basic pagination without last_id
|
|
@patch("services.message_service.db")
|
|
def test_pagination_by_last_id_without_last_id(self, mock_db, factory):
|
|
"""Test basic pagination without last_id."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=user,
|
|
last_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 5
|
|
assert result.has_more is False
|
|
assert result.limit == 10
|
|
|
|
# Test 11: Pagination with last_id
|
|
@patch("services.message_service.db")
|
|
def test_pagination_by_last_id_with_last_id(self, mock_db, factory):
|
|
"""Test pagination with last_id to get messages after a specific message."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
last_message = factory.create_message_mock(
|
|
message_id="msg-005",
|
|
created_at=datetime(2024, 1, 1, 12, 5),
|
|
)
|
|
|
|
# Messages after last_message
|
|
new_messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(6, 10)
|
|
]
|
|
|
|
# Setup base query mock that returns itself for chaining
|
|
mock_base_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_base_query
|
|
|
|
# First where() call for last_id lookup
|
|
mock_query_last = MagicMock()
|
|
mock_query_last.first.return_value = last_message
|
|
|
|
# Second where() call for history messages
|
|
mock_query_history = MagicMock()
|
|
mock_query_history.order_by.return_value = mock_query_history
|
|
mock_query_history.limit.return_value = mock_query_history
|
|
mock_query_history.all.return_value = new_messages
|
|
|
|
# Setup where() to return different mocks on consecutive calls
|
|
mock_base_query.where.side_effect = [mock_query_last, mock_query_history]
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=user,
|
|
last_id="msg-005",
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 4
|
|
assert result.has_more is False
|
|
|
|
# Test 12: Last message not found
|
|
@patch("services.message_service.db")
|
|
def test_pagination_by_last_id_last_message_not_exists(self, mock_db, factory):
|
|
"""Test error handling when last_id doesn't exist."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.first.return_value = None # Message not found
|
|
|
|
# Act & Assert
|
|
with pytest.raises(LastMessageNotExistsError):
|
|
MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=user,
|
|
last_id="nonexistent-msg",
|
|
limit=10,
|
|
)
|
|
|
|
# Test 13: Pagination with conversation_id filter
|
|
@patch("services.message_service.ConversationService")
|
|
@patch("services.message_service.db")
|
|
def test_pagination_by_last_id_with_conversation_filter(self, mock_db, mock_conversation_service, factory):
|
|
"""Test pagination filtered by conversation_id."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
conversation = factory.create_conversation_mock(conversation_id="conv-001")
|
|
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
conversation_id="conv-001",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=user,
|
|
last_id=None,
|
|
limit=10,
|
|
conversation_id="conv-001",
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 5
|
|
assert result.has_more is False
|
|
# Verify conversation_id was used in query
|
|
mock_query.where.assert_called()
|
|
mock_conversation_service.get_conversation.assert_called_once()
|
|
|
|
# Test 14: Pagination with include_ids filter
|
|
@patch("services.message_service.db")
|
|
def test_pagination_by_last_id_with_include_ids(self, mock_db, factory):
|
|
"""Test pagination filtered by include_ids."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
# Only messages with IDs in include_ids should be returned
|
|
messages = [
|
|
factory.create_message_mock(message_id="msg-001"),
|
|
factory.create_message_mock(message_id="msg-003"),
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=user,
|
|
last_id=None,
|
|
limit=10,
|
|
include_ids=["msg-001", "msg-003"],
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 2
|
|
assert result.data[0].id == "msg-001"
|
|
assert result.data[1].id == "msg-003"
|
|
|
|
# Test 15: Has_more flag when results exceed limit
|
|
@patch("services.message_service.db")
|
|
def test_pagination_by_last_id_has_more_true(self, mock_db, factory):
|
|
"""Test has_more flag is True when results exceed limit."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
# Create limit+1 messages (11 messages for limit=10)
|
|
messages = [
|
|
factory.create_message_mock(
|
|
message_id=f"msg-{i:03d}",
|
|
created_at=datetime(2024, 1, 1, 12, i),
|
|
)
|
|
for i in range(11)
|
|
]
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.all.return_value = messages
|
|
|
|
# Act
|
|
result = MessageService.pagination_by_last_id(
|
|
app_model=app,
|
|
user=user,
|
|
last_id=None,
|
|
limit=10,
|
|
)
|
|
|
|
# Assert
|
|
assert len(result.data) == 10 # Last message trimmed
|
|
assert result.has_more is True
|
|
assert result.limit == 10
|
|
|
|
|
|
class TestMessageServiceUtilities:
|
|
"""Unit tests for MessageService module-level utility functions."""
|
|
|
|
@pytest.fixture
|
|
def factory(self):
|
|
"""Provide test data factory."""
|
|
return TestMessageServiceFactory()
|
|
|
|
# Test 16: attach_message_extra_contents with empty list
|
|
def test_attach_message_extra_contents_empty(self):
|
|
"""Test attach_message_extra_contents with empty list does nothing."""
|
|
# Act & Assert (should not raise error)
|
|
attach_message_extra_contents([])
|
|
|
|
# Test 17: attach_message_extra_contents with messages
|
|
@patch("services.message_service._create_execution_extra_content_repository")
|
|
def test_attach_message_extra_contents_with_messages(self, mock_create_repo, factory):
|
|
"""Test attach_message_extra_contents correctly attaches content."""
|
|
# Arrange
|
|
messages = [factory.create_message_mock(message_id="msg-1"), factory.create_message_mock(message_id="msg-2")]
|
|
|
|
mock_repo = MagicMock()
|
|
mock_create_repo.return_value = mock_repo
|
|
|
|
# Mock extra content models
|
|
mock_content1 = MagicMock()
|
|
mock_content1.model_dump.return_value = {"key": "value1"}
|
|
mock_content2 = MagicMock()
|
|
mock_content2.model_dump.return_value = {"key": "value2"}
|
|
|
|
mock_repo.get_by_message_ids.return_value = [[mock_content1], [mock_content2]]
|
|
|
|
# Act
|
|
attach_message_extra_contents(messages)
|
|
|
|
# Assert
|
|
mock_repo.get_by_message_ids.assert_called_once_with(["msg-1", "msg-2"])
|
|
messages[0].set_extra_contents.assert_called_once_with([{"key": "value1"}])
|
|
messages[1].set_extra_contents.assert_called_once_with([{"key": "value2"}])
|
|
|
|
# Test 18: attach_message_extra_contents with index out of bounds
|
|
@patch("services.message_service._create_execution_extra_content_repository")
|
|
def test_attach_message_extra_contents_index_out_of_bounds(self, mock_create_repo, factory):
|
|
"""Test attach_message_extra_contents handles missing content lists."""
|
|
# Arrange
|
|
messages = [factory.create_message_mock(message_id="msg-1")]
|
|
|
|
mock_repo = MagicMock()
|
|
mock_create_repo.return_value = mock_repo
|
|
mock_repo.get_by_message_ids.return_value = [] # Empty returned list
|
|
|
|
# Act
|
|
attach_message_extra_contents(messages)
|
|
|
|
# Assert
|
|
messages[0].set_extra_contents.assert_called_once_with([])
|
|
|
|
# Test 19: _create_execution_extra_content_repository
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.sessionmaker")
|
|
@patch("services.message_service.SQLAlchemyExecutionExtraContentRepository")
|
|
def test_create_execution_extra_content_repository(self, mock_repo_class, mock_sessionmaker, mock_db):
|
|
"""Test _create_execution_extra_content_repository creates expected repository."""
|
|
from services.message_service import _create_execution_extra_content_repository
|
|
|
|
# Act
|
|
_create_execution_extra_content_repository()
|
|
|
|
# Assert
|
|
mock_sessionmaker.assert_called_once()
|
|
mock_repo_class.assert_called_once()
|
|
|
|
|
|
class TestMessageServiceGetMessage:
|
|
"""Unit tests for MessageService.get_message method."""
|
|
|
|
@pytest.fixture
|
|
def factory(self):
|
|
"""Provide test data factory."""
|
|
return TestMessageServiceFactory()
|
|
|
|
# Test 20: get_message success for EndUser
|
|
@patch("services.message_service.db")
|
|
def test_get_message_end_user_success(self, mock_db, factory):
|
|
"""Test get_message returns message for EndUser."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock(user_id="end-user-123")
|
|
message = factory.create_message_mock()
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.first.return_value = message
|
|
|
|
# Act
|
|
result = MessageService.get_message(app_model=app, user=user, message_id="msg-123")
|
|
|
|
# Assert
|
|
assert result == message
|
|
mock_query.where.assert_called_once()
|
|
|
|
# Test 21: get_message success for Account (Admin)
|
|
@patch("services.message_service.db")
|
|
def test_get_message_account_success(self, mock_db, factory):
|
|
"""Test get_message returns message for Account."""
|
|
# Arrange
|
|
from models import Account
|
|
|
|
app = factory.create_app_mock()
|
|
user = MagicMock(spec=Account)
|
|
user.id = "account-123"
|
|
message = factory.create_message_mock()
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.first.return_value = message
|
|
|
|
# Act
|
|
result = MessageService.get_message(app_model=app, user=user, message_id="msg-123")
|
|
|
|
# Assert
|
|
assert result == message
|
|
|
|
# Test 22: get_message not found
|
|
@patch("services.message_service.db")
|
|
def test_get_message_not_found(self, mock_db, factory):
|
|
"""Test get_message raises MessageNotExistsError when not found."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.first.return_value = None
|
|
|
|
# Act & Assert
|
|
with pytest.raises(MessageNotExistsError):
|
|
MessageService.get_message(app_model=app, user=user, message_id="msg-123")
|
|
|
|
|
|
class TestMessageServiceFeedback:
|
|
"""Unit tests for MessageService feedback-related methods."""
|
|
|
|
@pytest.fixture
|
|
def factory(self):
|
|
"""Provide test data factory."""
|
|
return TestMessageServiceFactory()
|
|
|
|
# Test 23: create_feedback - new feedback for EndUser
|
|
@patch("services.message_service.db")
|
|
@patch.object(MessageService, "get_message")
|
|
def test_create_feedback_new_end_user(self, mock_get_message, mock_db, factory):
|
|
"""Test creating new feedback for an end user."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
message = factory.create_message_mock()
|
|
message.user_feedback = None
|
|
mock_get_message.return_value = message
|
|
|
|
# Act
|
|
result = MessageService.create_feedback(
|
|
app_model=app,
|
|
message_id="msg-123",
|
|
user=user,
|
|
rating="like",
|
|
content="Good answer",
|
|
)
|
|
|
|
# Assert
|
|
assert result.rating == "like"
|
|
assert result.content == "Good answer"
|
|
assert result.from_source == "user"
|
|
mock_db.session.add.assert_called_once()
|
|
mock_db.session.commit.assert_called_once()
|
|
|
|
# Test 24: create_feedback - update feedback for Account
|
|
@patch("services.message_service.db")
|
|
@patch.object(MessageService, "get_message")
|
|
def test_create_feedback_update_account(self, mock_get_message, mock_db, factory):
|
|
"""Test updating existing feedback for an account."""
|
|
# Arrange
|
|
from models import Account, MessageFeedback
|
|
|
|
app = factory.create_app_mock()
|
|
user = MagicMock(spec=Account)
|
|
user.id = "account-123"
|
|
message = factory.create_message_mock()
|
|
feedback = MagicMock(spec=MessageFeedback)
|
|
message.admin_feedback = feedback
|
|
mock_get_message.return_value = message
|
|
|
|
# Act
|
|
result = MessageService.create_feedback(
|
|
app_model=app,
|
|
message_id="msg-123",
|
|
user=user,
|
|
rating="dislike",
|
|
content="Bad answer",
|
|
)
|
|
|
|
# Assert
|
|
assert result == feedback
|
|
assert feedback.rating == "dislike"
|
|
assert feedback.content == "Bad answer"
|
|
mock_db.session.commit.assert_called_once()
|
|
|
|
# Test 25: create_feedback - delete feedback (rating is None)
|
|
@patch("services.message_service.db")
|
|
@patch.object(MessageService, "get_message")
|
|
def test_create_feedback_delete(self, mock_get_message, mock_db, factory):
|
|
"""Test deleting feedback by passing rating=None."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
user = factory.create_end_user_mock()
|
|
message = factory.create_message_mock()
|
|
feedback = MagicMock()
|
|
message.user_feedback = feedback
|
|
mock_get_message.return_value = message
|
|
|
|
# Act
|
|
result = MessageService.create_feedback(
|
|
app_model=app,
|
|
message_id="msg-123",
|
|
user=user,
|
|
rating=None,
|
|
content=None,
|
|
)
|
|
|
|
# Assert
|
|
assert result == feedback
|
|
mock_db.session.delete.assert_called_once_with(feedback)
|
|
mock_db.session.commit.assert_called_once()
|
|
|
|
# Test 26: get_all_messages_feedbacks
|
|
@patch("services.message_service.db")
|
|
def test_get_all_messages_feedbacks(self, mock_db, factory):
|
|
"""Test get_all_messages_feedbacks returns list of dicts."""
|
|
# Arrange
|
|
app = factory.create_app_mock()
|
|
feedback = MagicMock()
|
|
feedback.to_dict.return_value = {"id": "fb-1"}
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.limit.return_value = mock_query
|
|
mock_query.offset.return_value = mock_query
|
|
mock_query.all.return_value = [feedback]
|
|
|
|
# Act
|
|
result = MessageService.get_all_messages_feedbacks(app_model=app, page=1, limit=10)
|
|
|
|
# Assert
|
|
assert result == [{"id": "fb-1"}]
|
|
mock_query.limit.assert_called_with(10)
|
|
mock_query.offset.assert_called_with(0)
|
|
|
|
|
|
class TestMessageServiceSuggestedQuestions:
|
|
"""Unit tests for MessageService.get_suggested_questions_after_answer method."""
|
|
|
|
@pytest.fixture
|
|
def factory(self):
|
|
"""Provide test data factory."""
|
|
return TestMessageServiceFactory()
|
|
|
|
# Test 27: get_suggested_questions_after_answer - user is None
|
|
def test_get_suggested_questions_user_none(self, factory):
|
|
app = factory.create_app_mock()
|
|
with pytest.raises(ValueError, match="user cannot be None"):
|
|
MessageService.get_suggested_questions_after_answer(
|
|
app_model=app, user=None, message_id="msg-123", invoke_from=MagicMock()
|
|
)
|
|
|
|
# Test 28: get_suggested_questions_after_answer - Advanced Chat success
|
|
@patch("services.message_service.ModelManager")
|
|
@patch("services.message_service.WorkflowService")
|
|
@patch("services.message_service.AdvancedChatAppConfigManager")
|
|
@patch("services.message_service.TokenBufferMemory")
|
|
@patch("services.message_service.LLMGenerator")
|
|
@patch("services.message_service.TraceQueueManager")
|
|
@patch.object(MessageService, "get_message")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_get_suggested_questions_advanced_chat_success(
|
|
self,
|
|
mock_conversation_service,
|
|
mock_get_message,
|
|
mock_trace_manager,
|
|
mock_llm_gen,
|
|
mock_memory,
|
|
mock_config_manager,
|
|
mock_workflow_service,
|
|
mock_model_manager,
|
|
factory,
|
|
):
|
|
"""Test successful suggested questions generation in Advanced Chat mode."""
|
|
from core.app.entities.app_invoke_entities import InvokeFrom
|
|
|
|
# Arrange
|
|
app = factory.create_app_mock(mode=AppMode.ADVANCED_CHAT.value)
|
|
user = factory.create_end_user_mock()
|
|
message = factory.create_message_mock()
|
|
mock_get_message.return_value = message
|
|
|
|
workflow = MagicMock()
|
|
mock_workflow_service.return_value.get_published_workflow.return_value = workflow
|
|
|
|
app_config = MagicMock()
|
|
app_config.additional_features.suggested_questions_after_answer = True
|
|
mock_config_manager.get_app_config.return_value = app_config
|
|
|
|
mock_llm_gen.generate_suggested_questions_after_answer.return_value = ["Q1?"]
|
|
|
|
# Act
|
|
result = MessageService.get_suggested_questions_after_answer(
|
|
app_model=app, user=user, message_id="msg-123", invoke_from=InvokeFrom.WEB_APP
|
|
)
|
|
|
|
# Assert
|
|
assert result == ["Q1?"]
|
|
mock_workflow_service.return_value.get_published_workflow.assert_called_once()
|
|
mock_llm_gen.generate_suggested_questions_after_answer.assert_called_once()
|
|
|
|
# Test 29: get_suggested_questions_after_answer - Chat app success (no override)
|
|
@patch("services.message_service.db")
|
|
@patch("services.message_service.ModelManager")
|
|
@patch("services.message_service.TokenBufferMemory")
|
|
@patch("services.message_service.LLMGenerator")
|
|
@patch("services.message_service.TraceQueueManager")
|
|
@patch.object(MessageService, "get_message")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_get_suggested_questions_chat_app_success(
|
|
self,
|
|
mock_conversation_service,
|
|
mock_get_message,
|
|
mock_trace_manager,
|
|
mock_llm_gen,
|
|
mock_memory,
|
|
mock_model_manager,
|
|
mock_db,
|
|
factory,
|
|
):
|
|
"""Test successful suggested questions generation in basic Chat mode."""
|
|
# Arrange
|
|
app = factory.create_app_mock(mode=AppMode.CHAT.value)
|
|
user = factory.create_end_user_mock()
|
|
message = factory.create_message_mock()
|
|
mock_get_message.return_value = message
|
|
|
|
conversation = MagicMock()
|
|
conversation.override_model_configs = None
|
|
mock_conversation_service.get_conversation.return_value = conversation
|
|
|
|
app_model_config = MagicMock()
|
|
app_model_config.suggested_questions_after_answer_dict = {"enabled": True}
|
|
app_model_config.model_dict = {"provider": "openai", "name": "gpt-4"}
|
|
|
|
mock_query = MagicMock()
|
|
mock_db.session.query.return_value = mock_query
|
|
mock_query.where.return_value = mock_query
|
|
mock_query.first.return_value = app_model_config
|
|
|
|
mock_llm_gen.generate_suggested_questions_after_answer.return_value = ["Q1?"]
|
|
|
|
# Act
|
|
result = MessageService.get_suggested_questions_after_answer(
|
|
app_model=app, user=user, message_id="msg-123", invoke_from=MagicMock()
|
|
)
|
|
|
|
# Assert
|
|
assert result == ["Q1?"]
|
|
mock_query.first.assert_called_once()
|
|
mock_llm_gen.generate_suggested_questions_after_answer.assert_called_once()
|
|
|
|
# Test 30: get_suggested_questions_after_answer - Disabled Error
|
|
@patch("services.message_service.WorkflowService")
|
|
@patch("services.message_service.AdvancedChatAppConfigManager")
|
|
@patch.object(MessageService, "get_message")
|
|
@patch("services.message_service.ConversationService")
|
|
def test_get_suggested_questions_disabled_error(
|
|
self, mock_conversation_service, mock_get_message, mock_config_manager, mock_workflow_service, factory
|
|
):
|
|
"""Test SuggestedQuestionsAfterAnswerDisabledError is raised when feature is disabled."""
|
|
# Arrange
|
|
app = factory.create_app_mock(mode=AppMode.ADVANCED_CHAT.value)
|
|
user = factory.create_end_user_mock()
|
|
mock_get_message.return_value = factory.create_message_mock()
|
|
|
|
workflow = MagicMock()
|
|
mock_workflow_service.return_value.get_published_workflow.return_value = workflow
|
|
|
|
app_config = MagicMock()
|
|
app_config.additional_features.suggested_questions_after_answer = False
|
|
mock_config_manager.get_app_config.return_value = app_config
|
|
|
|
# Act & Assert
|
|
with pytest.raises(SuggestedQuestionsAfterAnswerDisabledError):
|
|
MessageService.get_suggested_questions_after_answer(
|
|
app_model=app, user=user, message_id="msg-123", invoke_from=MagicMock()
|
|
)
|