mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-20 09:20:59 +00:00
feat(test): Phase 1.5 coverage expansion - COV-001 to COV-013
Coverage baseline raised from 13.65% to 26% (+12.35%) Gate raised from 12% to 18% Tasks completed: - COV-001: Command system unit tests (100% coverage) - COV-002: API service unit tests batch 1 (user/apikey/model/provider) - COV-003: Provider model manager unit tests - COV-004: Pipeline remaining stage tests (aggregator/cntfilter/longtext/msgtrun) - COV-005: Storage and utils coverage pass - COV-006: Gate ratchet 12%→15% - COV-007: Gate ratchet 15%→18% - COV-008: API service batch 2 (bot/pipeline/webhook/space/maintenance/mcp) - COV-009: Blocked - API controller circular import issue documented - COV-010: Plugin runtime unit tests (+0.08%) - COV-011: RAG and vector unit tests (+0.68%) - COV-012: Core boot and migration unit tests - COV-013: Provider requester logic unit tests (+0.62%) Key additions: - tests/utils/import_isolation.py: sys.modules isolation for circular imports - Provider requester mock tests: proved HTTP-dependent code can be tested locally - Vector filter utilities: 100% coverage on pure functions - API services: fake persistence pattern for unit testing Blocked issue COV-009 documented in langbot-test-plan/1.5/issues/ Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,596 @@
|
||||
"""
|
||||
Unit tests for MessageAggregator (aggregator) module.
|
||||
|
||||
Tests cover:
|
||||
- Message buffering and merging
|
||||
- Timer-based flush behavior
|
||||
- MAX_BUFFER_MESSAGES limit
|
||||
- Aggregation enabled/disabled
|
||||
- Config delay clamping
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from importlib import import_module
|
||||
|
||||
from tests.factories import (
|
||||
FakeApp,
|
||||
text_chain,
|
||||
friend_message_event,
|
||||
mock_adapter,
|
||||
)
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
|
||||
def get_aggregator_module():
|
||||
"""Lazy import to avoid circular import issues."""
|
||||
return import_module('langbot.pkg.pipeline.aggregator')
|
||||
|
||||
|
||||
def make_aggregator_app():
|
||||
"""Create a FakeApp with necessary mocks for aggregator tests."""
|
||||
app = FakeApp()
|
||||
# Ensure query_pool has add_query method
|
||||
app.query_pool.add_query = AsyncMock()
|
||||
# Add pipeline_mgr mock
|
||||
app.pipeline_mgr = AsyncMock()
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=None)
|
||||
return app
|
||||
|
||||
|
||||
class TestPendingMessage:
|
||||
"""Tests for PendingMessage dataclass."""
|
||||
|
||||
def test_pending_message_creation(self):
|
||||
"""PendingMessage should be created with correct fields."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
pending = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid='test-pipeline',
|
||||
)
|
||||
|
||||
assert pending.bot_uuid == 'test-bot'
|
||||
assert pending.launcher_type == provider_session.LauncherTypes.PERSON
|
||||
assert pending.message_chain == chain
|
||||
assert pending.timestamp is not None
|
||||
|
||||
|
||||
class TestSessionBuffer:
|
||||
"""Tests for SessionBuffer dataclass."""
|
||||
|
||||
def test_session_buffer_creation(self):
|
||||
"""SessionBuffer should be created with correct fields."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
buffer = aggregator.SessionBuffer(session_id='test-session')
|
||||
|
||||
assert buffer.session_id == 'test-session'
|
||||
assert buffer.messages == []
|
||||
assert buffer.timer_task is None
|
||||
assert buffer.last_message_time is not None
|
||||
|
||||
def test_session_buffer_with_messages(self):
|
||||
"""SessionBuffer should accept initial messages."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
pending = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
buffer = aggregator.SessionBuffer(
|
||||
session_id='test-session',
|
||||
messages=[pending],
|
||||
)
|
||||
|
||||
assert len(buffer.messages) == 1
|
||||
|
||||
|
||||
class TestMessageAggregatorInit:
|
||||
"""Tests for MessageAggregator initialization."""
|
||||
|
||||
def test_aggregator_init(self):
|
||||
"""MessageAggregator should initialize with correct fields."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
assert agg.ap == app
|
||||
assert agg.buffers == {}
|
||||
assert isinstance(agg.lock, asyncio.Lock)
|
||||
|
||||
|
||||
class TestMessageAggregatorSessionId:
|
||||
"""Tests for session ID generation."""
|
||||
|
||||
def test_session_id_format(self):
|
||||
"""Session ID should be correctly formatted."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
session_id = agg._get_session_id(
|
||||
bot_uuid='bot-123',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=45678,
|
||||
)
|
||||
|
||||
assert session_id == 'bot-123:person:45678'
|
||||
|
||||
def test_session_id_different_launchers(self):
|
||||
"""Different launcher types should produce different IDs."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
person_id = agg._get_session_id(
|
||||
bot_uuid='bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=123,
|
||||
)
|
||||
|
||||
group_id = agg._get_session_id(
|
||||
bot_uuid='bot',
|
||||
launcher_type=provider_session.LauncherTypes.GROUP,
|
||||
launcher_id=123,
|
||||
)
|
||||
|
||||
assert person_id != group_id
|
||||
|
||||
|
||||
class TestMessageAggregatorConfig:
|
||||
"""Tests for aggregation config retrieval."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_none_pipeline(self):
|
||||
"""None pipeline_uuid should return default config."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
enabled, delay = await agg._get_aggregation_config(None)
|
||||
|
||||
assert enabled == False
|
||||
assert delay == 1.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_pipeline_not_found(self):
|
||||
"""Non-existent pipeline should return default config."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=None)
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
enabled, delay = await agg._get_aggregation_config('unknown-pipeline')
|
||||
|
||||
assert enabled == False
|
||||
assert delay == 1.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_enabled(self):
|
||||
"""Pipeline with enabled aggregation should return True."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.pipeline_entity = Mock()
|
||||
mock_pipeline.pipeline_entity.config = {
|
||||
'trigger': {
|
||||
'message-aggregation': {
|
||||
'enabled': True,
|
||||
'delay': 2.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=mock_pipeline)
|
||||
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
enabled, delay = await agg._get_aggregation_config('test-pipeline')
|
||||
|
||||
assert enabled == True
|
||||
assert delay == 2.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_delay_clamped_low(self):
|
||||
"""Delay below 1.0 should be clamped to 1.0."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.pipeline_entity = Mock()
|
||||
mock_pipeline.pipeline_entity.config = {
|
||||
'trigger': {
|
||||
'message-aggregation': {
|
||||
'enabled': True,
|
||||
'delay': 0.5, # Below minimum
|
||||
}
|
||||
}
|
||||
}
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=mock_pipeline)
|
||||
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
enabled, delay = await agg._get_aggregation_config('test-pipeline')
|
||||
|
||||
assert delay == 1.0 # Clamped to minimum
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_delay_clamped_high(self):
|
||||
"""Delay above 10.0 should be clamped to 10.0."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.pipeline_entity = Mock()
|
||||
mock_pipeline.pipeline_entity.config = {
|
||||
'trigger': {
|
||||
'message-aggregation': {
|
||||
'enabled': True,
|
||||
'delay': 15.0, # Above maximum
|
||||
}
|
||||
}
|
||||
}
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=mock_pipeline)
|
||||
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
enabled, delay = await agg._get_aggregation_config('test-pipeline')
|
||||
|
||||
assert delay == 10.0 # Clamped to maximum
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_delay_invalid_type(self):
|
||||
"""Invalid delay type should use default."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.pipeline_entity = Mock()
|
||||
mock_pipeline.pipeline_entity.config = {
|
||||
'trigger': {
|
||||
'message-aggregation': {
|
||||
'enabled': True,
|
||||
'delay': 'invalid', # Not a number
|
||||
}
|
||||
}
|
||||
}
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=mock_pipeline)
|
||||
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
enabled, delay = await agg._get_aggregation_config('test-pipeline')
|
||||
|
||||
assert delay == 1.5 # Default
|
||||
|
||||
|
||||
class TestMessageAggregatorAddMessage:
|
||||
"""Tests for add_message behavior."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_adds_to_query_pool(self):
|
||||
"""Disabled aggregation should directly add to query_pool."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
await agg.add_message(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None, # None -> disabled
|
||||
)
|
||||
|
||||
# Should have called query_pool.add_query
|
||||
assert app.query_pool.add_query.called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enabled_buffers_message(self):
|
||||
"""Enabled aggregation should buffer message."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.pipeline_entity = Mock()
|
||||
mock_pipeline.pipeline_entity.config = {
|
||||
'trigger': {
|
||||
'message-aggregation': {
|
||||
'enabled': True,
|
||||
'delay': 2.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=mock_pipeline)
|
||||
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
await agg.add_message(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid='test-pipeline',
|
||||
)
|
||||
|
||||
# Should have buffered the message
|
||||
assert len(agg.buffers) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_buffer_flushes_immediately(self):
|
||||
"""Reaching MAX_BUFFER_MESSAGES should flush immediately."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.pipeline_entity = Mock()
|
||||
mock_pipeline.pipeline_entity.config = {
|
||||
'trigger': {
|
||||
'message-aggregation': {
|
||||
'enabled': True,
|
||||
'delay': 10.0, # Long delay
|
||||
}
|
||||
}
|
||||
}
|
||||
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=mock_pipeline)
|
||||
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
# Add messages up to MAX_BUFFER_MESSAGES
|
||||
for i in range(aggregator.MAX_BUFFER_MESSAGES):
|
||||
await agg.add_message(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid='test-pipeline',
|
||||
)
|
||||
|
||||
# Buffer should be flushed (empty or no buffer)
|
||||
session_id = agg._get_session_id('test-bot', provider_session.LauncherTypes.PERSON, 12345)
|
||||
assert session_id not in agg.buffers or len(agg.buffers[session_id].messages) == 0
|
||||
|
||||
|
||||
class TestMessageAggregatorMerge:
|
||||
"""Tests for message merging."""
|
||||
|
||||
def test_merge_single_message(self):
|
||||
"""Single message should return unchanged."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
pending = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
merged = agg._merge_messages([pending])
|
||||
|
||||
assert merged.message_chain == chain
|
||||
|
||||
def test_merge_multiple_messages(self):
|
||||
"""Multiple messages should be merged with newline separator."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain1 = text_chain("hello")
|
||||
chain2 = text_chain("world")
|
||||
event = friend_message_event(chain1)
|
||||
adapter = mock_adapter()
|
||||
|
||||
pending1 = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain1,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
pending2 = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain2,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
merged = agg._merge_messages([pending1, pending2])
|
||||
|
||||
# Should contain both messages with separator
|
||||
merged_str = str(merged.message_chain)
|
||||
assert "hello" in merged_str
|
||||
assert "world" in merged_str
|
||||
|
||||
|
||||
class TestMessageAggregatorFlush:
|
||||
"""Tests for buffer flush behavior."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_empty_buffer(self):
|
||||
"""Flushing empty buffer should do nothing."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
await agg._flush_buffer('nonexistent-session')
|
||||
|
||||
# Should not call query_pool
|
||||
assert not app.query_pool.add_query.called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_single_message(self):
|
||||
"""Flushing single message should add directly to query_pool."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
pending = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
buffer = aggregator.SessionBuffer(
|
||||
session_id='test-session',
|
||||
messages=[pending],
|
||||
)
|
||||
|
||||
agg.buffers['test-session'] = buffer
|
||||
|
||||
await agg._flush_buffer('test-session')
|
||||
|
||||
assert app.query_pool.add_query.called
|
||||
assert 'test-session' not in agg.buffers
|
||||
|
||||
|
||||
class TestMessageAggregatorFlushAll:
|
||||
"""Tests for flush_all behavior."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_all_empty(self):
|
||||
"""flush_all with no buffers should do nothing."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
await agg.flush_all()
|
||||
|
||||
# Should not call query_pool
|
||||
assert not app.query_pool.add_query.called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_all_with_buffers(self):
|
||||
"""flush_all should flush all pending buffers."""
|
||||
aggregator = get_aggregator_module()
|
||||
|
||||
app = make_aggregator_app()
|
||||
agg = aggregator.MessageAggregator(app)
|
||||
|
||||
chain = text_chain("hello")
|
||||
event = friend_message_event(chain)
|
||||
adapter = mock_adapter()
|
||||
|
||||
# Create two buffers
|
||||
pending1 = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
pending2 = aggregator.PendingMessage(
|
||||
bot_uuid='test-bot',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=67890,
|
||||
sender_id=67890,
|
||||
message_event=event,
|
||||
message_chain=chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=None,
|
||||
)
|
||||
|
||||
buffer1 = aggregator.SessionBuffer(session_id='session-1', messages=[pending1])
|
||||
buffer2 = aggregator.SessionBuffer(session_id='session-2', messages=[pending2])
|
||||
|
||||
agg.buffers['session-1'] = buffer1
|
||||
agg.buffers['session-2'] = buffer2
|
||||
|
||||
await agg.flush_all()
|
||||
|
||||
# Both buffers should be flushed
|
||||
assert len(agg.buffers) == 0
|
||||
assert app.query_pool.add_query.call_count == 2
|
||||
@@ -0,0 +1,514 @@
|
||||
"""
|
||||
Unit tests for ContentFilterStage (cntfilter) pipeline stage.
|
||||
|
||||
Tests cover:
|
||||
- Pre-filter behavior (income message filtering)
|
||||
- Post-filter behavior (output message filtering)
|
||||
- Content ignore rules (prefix/regexp)
|
||||
- Pass/Block/Masked result handling
|
||||
- CONTINUE/INTERRUPT flow control
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock
|
||||
from importlib import import_module
|
||||
|
||||
from tests.factories import (
|
||||
FakeApp,
|
||||
text_query,
|
||||
image_query,
|
||||
)
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
|
||||
def get_cntfilter_module():
|
||||
"""Lazy import to avoid circular import issues."""
|
||||
# Import pipelinemgr first to trigger stage registration
|
||||
import_module('langbot.pkg.pipeline.pipelinemgr')
|
||||
return import_module('langbot.pkg.pipeline.cntfilter.cntfilter')
|
||||
|
||||
|
||||
def get_filter_module():
|
||||
"""Lazy import for filter base."""
|
||||
return import_module('langbot.pkg.pipeline.cntfilter.filter')
|
||||
|
||||
|
||||
def get_entities_module():
|
||||
"""Lazy import for pipeline entities."""
|
||||
return import_module('langbot.pkg.pipeline.entities')
|
||||
|
||||
|
||||
def get_filter_entities_module():
|
||||
"""Lazy import for filter entities."""
|
||||
return import_module('langbot.pkg.pipeline.cntfilter.entities')
|
||||
|
||||
|
||||
def make_pipeline_config(**overrides):
|
||||
"""Create a pipeline config with defaults for content filter tests."""
|
||||
base_config = {
|
||||
'safety': {
|
||||
'content-filter': {
|
||||
'check-sensitive-words': False,
|
||||
'scope': 'both',
|
||||
}
|
||||
},
|
||||
'trigger': {
|
||||
'ignore-rules': {
|
||||
'prefix': [],
|
||||
'regexp': [],
|
||||
}
|
||||
},
|
||||
}
|
||||
# Deep merge for nested dicts
|
||||
for key, value in overrides.items():
|
||||
if key in base_config and isinstance(base_config[key], dict) and isinstance(value, dict):
|
||||
for sub_key, sub_value in value.items():
|
||||
if sub_key in base_config[key] and isinstance(base_config[key][sub_key], dict) and isinstance(sub_value, dict):
|
||||
base_config[key][sub_key].update(sub_value)
|
||||
else:
|
||||
base_config[key][sub_key] = sub_value
|
||||
else:
|
||||
base_config[key] = value
|
||||
return base_config
|
||||
|
||||
|
||||
class TestContentFilterStageInit:
|
||||
"""Tests for ContentFilterStage initialization."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_basic_filters(self):
|
||||
"""Initialize should load required filters."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
assert stage.filter_chain is not None
|
||||
# Should have at least 'content-ignore' filter
|
||||
assert len(stage.filter_chain) >= 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_with_sensitive_words(self):
|
||||
"""Initialize with sensitive words should load ban-word-filter."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
|
||||
app = FakeApp()
|
||||
# Mock sensitive_meta for ban-word-filter
|
||||
app.sensitive_meta = Mock()
|
||||
app.sensitive_meta.data = {
|
||||
'words': [],
|
||||
'mask': '*',
|
||||
'mask_word': '',
|
||||
}
|
||||
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
safety={
|
||||
'content-filter': {
|
||||
'check-sensitive-words': True,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
# Should have content-ignore and ban-word-filter
|
||||
assert len(stage.filter_chain) >= 2
|
||||
|
||||
|
||||
class TestPreContentFilter:
|
||||
"""Tests for PreContentFilterStage (income message filtering)."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normal_text_continues(self):
|
||||
"""Normal text message should continue pipeline."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello world")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
assert result.new_query is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_text_continues(self):
|
||||
"""Empty text message should continue pipeline."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
# Empty message chain
|
||||
query = text_query("")
|
||||
query.message_chain = platform_message.MessageChain([])
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Empty messages should continue
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_whitespace_only_continues(self):
|
||||
"""Whitespace-only message should continue pipeline."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query(" ") # Only whitespace
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Whitespace-only should continue (stripped becomes empty)
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_text_component_continues(self):
|
||||
"""Message with non-text components should continue (skip filter)."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
# Image message (non-text)
|
||||
query = image_query()
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Non-text messages should continue (skip filter)
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_scope_skip_pre_filter(self):
|
||||
"""scope=output-msg should skip pre-filter."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
safety={
|
||||
'content-filter': {
|
||||
'scope': 'output-msg', # Only check output
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello world")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Should continue without filtering
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
|
||||
class TestContentIgnoreFilter:
|
||||
"""Tests for content-ignore filter rules."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prefix_rule_blocks(self):
|
||||
"""Message matching prefix ignore rule should be blocked."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
trigger={
|
||||
'ignore-rules': {
|
||||
'prefix': ['/help', '/ping'],
|
||||
'regexp': [],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("/help me")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Should be interrupted due to prefix rule
|
||||
assert result.result_type == entities.ResultType.INTERRUPT
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regexp_rule_blocks(self):
|
||||
"""Message matching regexp ignore rule should be blocked."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
trigger={
|
||||
'ignore-rules': {
|
||||
'prefix': [],
|
||||
'regexp': ['^http://.*', r'\d{10}'],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("http://example.com")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Should be interrupted due to regexp rule
|
||||
assert result.result_type == entities.ResultType.INTERRUPT
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_rule_match_continues(self):
|
||||
"""Message not matching any rule should continue."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
trigger={
|
||||
'ignore-rules': {
|
||||
'prefix': ['/help', '/ping'],
|
||||
'regexp': ['^http://.*'],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("normal message")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Should continue (no rule match)
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_rules_continues(self):
|
||||
"""Empty ignore rules should not block any message."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("/help me")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
# Should continue (empty rules)
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
|
||||
class TestPostContentFilter:
|
||||
"""Tests for PostContentFilterStage (output message filtering)."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normal_response_continues(self):
|
||||
"""Normal response message should continue pipeline."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
# Add a response message
|
||||
query.resp_messages = [
|
||||
provider_message.Message(role='assistant', content='Hello back!')
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'PostContentFilterStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_income_scope_skip_post_filter(self):
|
||||
"""scope=income-msg should skip post-filter."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
safety={
|
||||
'content-filter': {
|
||||
'scope': 'income-msg', # Only check income
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_messages = [
|
||||
provider_message.Message(role='assistant', content='Response')
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'PostContentFilterStage')
|
||||
|
||||
# Should continue without filtering
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_string_content_continues(self):
|
||||
"""Non-string content should continue (skip filter)."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
# Non-string content - use model_construct to bypass validation
|
||||
# The actual content type could be a list of ContentElement objects
|
||||
non_string_msg = provider_message.Message.model_construct(
|
||||
role='assistant',
|
||||
content=[Mock()], # Mock content element
|
||||
)
|
||||
query.resp_messages = [non_string_msg]
|
||||
|
||||
result = await stage.process(query, 'PostContentFilterStage')
|
||||
|
||||
# Should continue (skip filter for non-string)
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_response_continues(self):
|
||||
"""Empty response should continue pipeline."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_messages = [
|
||||
provider_message.Message(role='assistant', content='')
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'PostContentFilterStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
|
||||
class TestContentFilterStageInvalidName:
|
||||
"""Tests for invalid stage_inst_name handling."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_stage_name_raises(self):
|
||||
"""Unknown stage_inst_name should raise ValueError."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
with pytest.raises(ValueError, match='未知的 stage_inst_name'):
|
||||
await stage.process(query, 'UnknownStage')
|
||||
|
||||
|
||||
class TestContentIgnoreFilterDirect:
|
||||
"""Direct tests for ContentIgnore filter."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_content_ignore_pass(self):
|
||||
"""ContentIgnore should PASS for non-matching messages."""
|
||||
cntfilter = get_cntfilter_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
stage = cntfilter.ContentFilterStage(app)
|
||||
|
||||
pipeline_config = make_pipeline_config(
|
||||
trigger={
|
||||
'ignore-rules': {
|
||||
'prefix': ['/test'],
|
||||
'regexp': [],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("normal message without prefix")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
result = await stage.process(query, 'PreContentFilterStage')
|
||||
|
||||
assert result.result_type == cntfilter.entities.ResultType.CONTINUE
|
||||
@@ -0,0 +1,370 @@
|
||||
"""
|
||||
Unit tests for LongTextProcessStage (longtext) pipeline stage.
|
||||
|
||||
Tests cover:
|
||||
- Strategy selection (none/image/forward)
|
||||
- Threshold boundary handling
|
||||
- Plain/non-Plain component handling
|
||||
- Strategy initialization and process
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock
|
||||
from importlib import import_module
|
||||
|
||||
from tests.factories import (
|
||||
FakeApp,
|
||||
text_query,
|
||||
)
|
||||
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
|
||||
def get_longtext_module():
|
||||
"""Lazy import to avoid circular import issues."""
|
||||
# Import pipelinemgr first to trigger stage registration
|
||||
import_module('langbot.pkg.pipeline.pipelinemgr')
|
||||
return import_module('langbot.pkg.pipeline.longtext.longtext')
|
||||
|
||||
|
||||
def get_strategy_module():
|
||||
"""Lazy import for strategy base."""
|
||||
return import_module('langbot.pkg.pipeline.longtext.strategy')
|
||||
|
||||
|
||||
def get_entities_module():
|
||||
"""Lazy import for pipeline entities."""
|
||||
return import_module('langbot.pkg.pipeline.entities')
|
||||
|
||||
|
||||
def make_longtext_config(strategy: str = 'none', threshold: int = 1000):
|
||||
"""Create a pipeline config for long text processing."""
|
||||
return {
|
||||
'output': {
|
||||
'long-text-processing': {
|
||||
'strategy': strategy,
|
||||
'threshold': threshold,
|
||||
'font-path': '/nonexistent/font.ttf', # For image strategy
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class TestLongTextProcessStageInit:
|
||||
"""Tests for LongTextProcessStage initialization."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_none_strategy(self):
|
||||
"""Initialize with strategy='none' should set strategy_impl to None."""
|
||||
longtext = get_longtext_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
pipeline_config = make_longtext_config(strategy='none')
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
assert stage.strategy_impl is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_forward_strategy(self):
|
||||
"""Initialize with strategy='forward' should use ForwardComponentStrategy."""
|
||||
longtext = get_longtext_module()
|
||||
strategy = get_strategy_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
pipeline_config = make_longtext_config(strategy='forward')
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
assert stage.strategy_impl is not None
|
||||
assert isinstance(stage.strategy_impl, strategy.LongTextStrategy)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_unknown_strategy_raises(self):
|
||||
"""Initialize with unknown strategy should raise ValueError."""
|
||||
longtext = get_longtext_module()
|
||||
strategy = get_strategy_module()
|
||||
|
||||
# Save original preregistered_strategies
|
||||
original_strategies = strategy.preregistered_strategies.copy()
|
||||
|
||||
try:
|
||||
# Clear registered strategies to simulate unknown
|
||||
strategy.preregistered_strategies = []
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
pipeline_config = make_longtext_config(strategy='unknown')
|
||||
|
||||
with pytest.raises(ValueError, match='Long message processing strategy not found'):
|
||||
await stage.initialize(pipeline_config)
|
||||
finally:
|
||||
# Restore original strategies
|
||||
strategy.preregistered_strategies = original_strategies
|
||||
|
||||
|
||||
class TestLongTextProcessStageProcess:
|
||||
"""Tests for LongTextProcessStage process behavior."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_strategy_continues(self):
|
||||
"""strategy='none' should always continue."""
|
||||
longtext = get_longtext_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
pipeline_config = make_longtext_config(strategy='none')
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = [
|
||||
platform_message.MessageChain([platform_message.Plain(text="very long response")])
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
assert result.new_query is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_short_text_continues_without_transform(self):
|
||||
"""Text shorter than threshold should not be transformed."""
|
||||
longtext = get_longtext_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
# High threshold so text won't trigger transform
|
||||
pipeline_config = make_longtext_config(strategy='forward', threshold=10000)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = [
|
||||
platform_message.MessageChain([platform_message.Plain(text="short response")])
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# Should not transform short text
|
||||
assert result.new_query.resp_message_chain is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_plain_component_skips(self):
|
||||
"""resp_message_chain with non-Plain components should skip processing."""
|
||||
longtext = get_longtext_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
pipeline_config = make_longtext_config(strategy='forward', threshold=10) # Low threshold
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
# Non-Plain component (Image)
|
||||
query.resp_message_chain = [
|
||||
platform_message.MessageChain([
|
||||
platform_message.Plain(text="short"),
|
||||
platform_message.Image(url="https://example.com/img.png")
|
||||
])
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# Should skip due to non-Plain component
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_resp_message_chain(self):
|
||||
"""Empty resp_message_chain should be handled gracefully."""
|
||||
longtext = get_longtext_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
pipeline_config = make_longtext_config(strategy='forward')
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Should handle gracefully (may raise or return CONTINUE)
|
||||
# This tests the defensive behavior
|
||||
try:
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
# If it returns, should be CONTINUE
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
except (IndexError, AttributeError):
|
||||
# Expected if resp_message_chain is empty
|
||||
pass
|
||||
|
||||
|
||||
class TestForwardStrategy:
|
||||
"""Tests for ForwardComponentStrategy."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forward_strategy_processes(self):
|
||||
"""ForwardComponentStrategy should create Forward component."""
|
||||
longtext = get_longtext_module()
|
||||
get_strategy_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
# Low threshold to trigger
|
||||
pipeline_config = make_longtext_config(strategy='forward', threshold=10)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
# Create a mock adapter with bot_account_id
|
||||
mock_adapter = Mock()
|
||||
mock_adapter.bot_account_id = '12345'
|
||||
query.adapter = mock_adapter
|
||||
|
||||
# Long text exceeding threshold
|
||||
long_text = "This is a very long response that exceeds the threshold"
|
||||
query.resp_message_chain = [
|
||||
platform_message.MessageChain([platform_message.Plain(text=long_text)])
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# Check that message chain was transformed
|
||||
assert result.new_query.resp_message_chain is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forward_strategy_direct_process(self):
|
||||
"""Test ForwardComponentStrategy process method directly."""
|
||||
strategy = get_strategy_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Get ForwardComponentStrategy from preregistered
|
||||
for strat_cls in strategy.preregistered_strategies:
|
||||
if strat_cls.name == 'forward':
|
||||
strat = strat_cls(app)
|
||||
break
|
||||
else:
|
||||
pytest.skip('ForwardComponentStrategy not registered')
|
||||
|
||||
await strat.initialize()
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = make_longtext_config()
|
||||
mock_adapter = Mock()
|
||||
mock_adapter.bot_account_id = '12345'
|
||||
query.adapter = mock_adapter
|
||||
|
||||
components = await strat.process("test message", query)
|
||||
|
||||
assert len(components) == 1
|
||||
assert isinstance(components[0], platform_message.Forward)
|
||||
|
||||
|
||||
class TestLongTextThreshold:
|
||||
"""Tests for threshold boundary handling."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exact_threshold_continues(self):
|
||||
"""Text exactly at threshold should trigger processing."""
|
||||
longtext = get_longtext_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
threshold = 50
|
||||
pipeline_config = make_longtext_config(strategy='forward', threshold=threshold)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
mock_adapter = Mock()
|
||||
mock_adapter.bot_account_id = '12345'
|
||||
query.adapter = mock_adapter
|
||||
|
||||
# Text exactly at threshold
|
||||
exact_text = "x" * threshold
|
||||
query.resp_message_chain = [
|
||||
platform_message.MessageChain([platform_message.Plain(text=exact_text)])
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_below_threshold_not_processed(self):
|
||||
"""Text below threshold should not be transformed."""
|
||||
longtext = get_longtext_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
threshold = 100
|
||||
pipeline_config = make_longtext_config(strategy='forward', threshold=threshold)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
|
||||
# Text below threshold
|
||||
short_text = "x" * (threshold - 1)
|
||||
query.resp_message_chain = [
|
||||
platform_message.MessageChain([platform_message.Plain(text=short_text)])
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'LongTextProcessStage')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# Original chain should remain unchanged
|
||||
|
||||
|
||||
class TestLongTextProcessStageImageStrategy:
|
||||
"""Tests for image strategy handling (requires PIL/font)."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_strategy_missing_font_fallback(self):
|
||||
"""Missing font should fallback to forward strategy."""
|
||||
longtext = get_longtext_module()
|
||||
strategy = get_strategy_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = longtext.LongTextProcessStage(app)
|
||||
|
||||
# Use non-existent font path
|
||||
pipeline_config = make_longtext_config(strategy='image')
|
||||
|
||||
# On non-Windows without font, should fallback to forward
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
# Should have initialized (possibly with fallback strategy)
|
||||
if stage.strategy_impl is not None:
|
||||
assert isinstance(stage.strategy_impl, strategy.LongTextStrategy)
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
Unit tests for ConversationMessageTruncator (msgtrun) pipeline stage.
|
||||
|
||||
Tests cover:
|
||||
- Normal truncation behavior based on max-round
|
||||
- Boundary length handling
|
||||
- Empty message handling
|
||||
- Multi-message chain truncation
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from importlib import import_module
|
||||
|
||||
from tests.factories import (
|
||||
FakeApp,
|
||||
text_query,
|
||||
)
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
|
||||
|
||||
def get_msgtrun_module():
|
||||
"""Lazy import to avoid circular import issues."""
|
||||
# Import pipelinemgr first to trigger stage registration
|
||||
import_module('langbot.pkg.pipeline.pipelinemgr')
|
||||
return import_module('langbot.pkg.pipeline.msgtrun.msgtrun')
|
||||
|
||||
|
||||
def get_truncator_module():
|
||||
"""Lazy import for truncator base."""
|
||||
return import_module('langbot.pkg.pipeline.msgtrun.truncator')
|
||||
|
||||
|
||||
def get_entities_module():
|
||||
"""Lazy import for pipeline entities."""
|
||||
return import_module('langbot.pkg.pipeline.entities')
|
||||
|
||||
|
||||
def get_round_truncator_module():
|
||||
"""Lazy import for round truncator."""
|
||||
return import_module('langbot.pkg.pipeline.msgtrun.truncators.round')
|
||||
|
||||
|
||||
def make_truncate_config(max_round: int = 5):
|
||||
"""Create a pipeline config with max-round setting."""
|
||||
return {
|
||||
'ai': {
|
||||
'local-agent': {
|
||||
'max-round': max_round,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class TestConversationMessageTruncatorInit:
|
||||
"""Tests for ConversationMessageTruncator initialization."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_round_truncator(self):
|
||||
"""Initialize should select 'round' truncator by default."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
truncator = get_truncator_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
assert stage.trun is not None
|
||||
assert isinstance(stage.trun, truncator.Truncator)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_unknown_truncator_raises(self):
|
||||
"""Initialize with unknown truncator method should raise ValueError."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
truncator = get_truncator_module()
|
||||
|
||||
# Save original preregistered_truncators
|
||||
original_truncators = truncator.preregistered_truncators.copy()
|
||||
|
||||
try:
|
||||
# Clear registered truncators to simulate unknown method
|
||||
truncator.preregistered_truncators = []
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config()
|
||||
|
||||
with pytest.raises(ValueError, match='Unknown truncator'):
|
||||
await stage.initialize(pipeline_config)
|
||||
finally:
|
||||
# Restore original truncators
|
||||
truncator.preregistered_truncators = original_truncators
|
||||
|
||||
|
||||
class TestRoundTruncatorProcess:
|
||||
"""Tests for RoundTruncator truncation behavior."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_within_limit(self):
|
||||
"""Messages within max-round limit should not be truncated."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config(max_round=5)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
# Create query with 3 messages (within limit)
|
||||
query = text_query("current message")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.messages = [
|
||||
provider_message.Message(role='user', content='message 1'),
|
||||
provider_message.Message(role='assistant', content='response 1'),
|
||||
provider_message.Message(role='user', content='message 2'),
|
||||
provider_message.Message(role='assistant', content='response 2'),
|
||||
provider_message.Message(role='user', content='current message'),
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'ConversationMessageTruncator')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# All messages should be preserved
|
||||
assert len(result.new_query.messages) == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_exceeds_limit(self):
|
||||
"""Messages exceeding max-round should be truncated."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config(max_round=2) # Only keep 2 rounds
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
# Create query with many messages exceeding limit
|
||||
query = text_query("current message")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.messages = [
|
||||
provider_message.Message(role='user', content='message 1'),
|
||||
provider_message.Message(role='assistant', content='response 1'),
|
||||
provider_message.Message(role='user', content='message 2'),
|
||||
provider_message.Message(role='assistant', content='response 2'),
|
||||
provider_message.Message(role='user', content='message 3'),
|
||||
provider_message.Message(role='assistant', content='response 3'),
|
||||
provider_message.Message(role='user', content='current message'),
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'ConversationMessageTruncator')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# Should only keep last 2 rounds (2 user messages)
|
||||
# Each round = user + assistant, so 2 rounds = 4 messages + current = 5
|
||||
assert len(result.new_query.messages) <= 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_empty_messages(self):
|
||||
"""Empty messages list should return empty list."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.messages = []
|
||||
|
||||
result = await stage.process(query, 'ConversationMessageTruncator')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
assert len(result.new_query.messages) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_single_message(self):
|
||||
"""Single message should be preserved."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.messages = [
|
||||
provider_message.Message(role='user', content='hello'),
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'ConversationMessageTruncator')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
assert len(result.new_query.messages) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_preserves_order(self):
|
||||
"""Truncation should preserve message order."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config(max_round=2)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("current")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.messages = [
|
||||
provider_message.Message(role='user', content='user1'),
|
||||
provider_message.Message(role='assistant', content='asst1'),
|
||||
provider_message.Message(role='user', content='user2'),
|
||||
provider_message.Message(role='assistant', content='asst2'),
|
||||
provider_message.Message(role='user', content='user3'),
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'ConversationMessageTruncator')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
# Check order is preserved (user2 -> asst2 -> user3)
|
||||
messages = result.new_query.messages
|
||||
if len(messages) >= 3:
|
||||
assert messages[0].role == 'user'
|
||||
assert messages[0].content == 'user2'
|
||||
assert messages[1].role == 'assistant'
|
||||
assert messages[1].content == 'asst2'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_max_round_one(self):
|
||||
"""max-round=1 should only keep last user message."""
|
||||
msgtrun = get_msgtrun_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = msgtrun.ConversationMessageTruncator(app)
|
||||
|
||||
pipeline_config = make_truncate_config(max_round=1)
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("current")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.messages = [
|
||||
provider_message.Message(role='user', content='old1'),
|
||||
provider_message.Message(role='assistant', content='old1_resp'),
|
||||
provider_message.Message(role='user', content='current'),
|
||||
]
|
||||
|
||||
result = await stage.process(query, 'ConversationMessageTruncator')
|
||||
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
# Only last round (user + assistant pair) should remain
|
||||
messages = result.new_query.messages
|
||||
# At most 2 messages (user + assistant before current)
|
||||
assert len(messages) <= 2
|
||||
|
||||
|
||||
class TestRoundTruncatorDirect:
|
||||
"""Direct tests for RoundTruncator class."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_round_truncator_direct_process(self):
|
||||
"""Test RoundTruncator truncate method directly."""
|
||||
truncator_mod = get_truncator_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Get the RoundTruncator class from preregistered
|
||||
for trun_cls in truncator_mod.preregistered_truncators:
|
||||
if trun_cls.name == 'round':
|
||||
trun = trun_cls(app)
|
||||
break
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = make_truncate_config(max_round=3)
|
||||
query.messages = [
|
||||
provider_message.Message(role='user', content='m1'),
|
||||
provider_message.Message(role='assistant', content='r1'),
|
||||
provider_message.Message(role='user', content='m2'),
|
||||
provider_message.Message(role='assistant', content='r2'),
|
||||
provider_message.Message(role='user', content='hello'),
|
||||
]
|
||||
|
||||
result = await trun.truncate(query)
|
||||
|
||||
assert result is not None
|
||||
assert hasattr(result, 'messages')
|
||||
@@ -0,0 +1,475 @@
|
||||
"""
|
||||
Unit tests for ResponseWrapper (wrapper) pipeline stage.
|
||||
|
||||
Tests cover:
|
||||
- MessageChain wrapping
|
||||
- Command response wrapping
|
||||
- Plugin response wrapping
|
||||
- Assistant response wrapping with content/tool_calls
|
||||
- Plugin event emission and INTERRUPT handling
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from importlib import import_module
|
||||
|
||||
from tests.factories import (
|
||||
FakeApp,
|
||||
text_query,
|
||||
)
|
||||
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
|
||||
def get_wrapper_module():
|
||||
"""Lazy import to avoid circular import issues."""
|
||||
# Import pipelinemgr first to trigger stage registration
|
||||
import_module('langbot.pkg.pipeline.pipelinemgr')
|
||||
return import_module('langbot.pkg.pipeline.wrapper.wrapper')
|
||||
|
||||
|
||||
def get_entities_module():
|
||||
"""Lazy import for pipeline entities."""
|
||||
return import_module('langbot.pkg.pipeline.entities')
|
||||
|
||||
|
||||
def make_wrapper_config():
|
||||
"""Create a pipeline config for wrapper tests."""
|
||||
return {
|
||||
'output': {
|
||||
'misc': {
|
||||
'at-sender': False,
|
||||
'quote-origin': False,
|
||||
'track-function-calls': False,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def make_session():
|
||||
"""Create a valid Session object for tests."""
|
||||
return provider_session.Session(
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
use_prompt_name="default",
|
||||
using_conversation=None,
|
||||
conversations=[],
|
||||
)
|
||||
|
||||
|
||||
class TestResponseWrapperInit:
|
||||
"""Tests for ResponseWrapper initialization."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_passes(self):
|
||||
"""Initialize should complete without error."""
|
||||
wrapper = get_wrapper_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = {}
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
|
||||
class TestResponseWrapperMessageChain:
|
||||
"""Tests for MessageChain wrapping."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_chain_direct_append(self):
|
||||
"""MessageChain in resp_messages should be directly appended."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_messages = [
|
||||
platform_message.MessageChain([platform_message.Plain(text="response")])
|
||||
]
|
||||
query.resp_message_chain = []
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].result_type == entities.ResultType.CONTINUE
|
||||
assert len(results[0].new_query.resp_message_chain) == 1
|
||||
|
||||
|
||||
class TestResponseWrapperCommand:
|
||||
"""Tests for command response wrapping."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_command_response_prefix(self):
|
||||
"""Command response should have [bot] prefix."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create a command response message
|
||||
command_resp = Mock()
|
||||
command_resp.role = 'command'
|
||||
command_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Help info")])
|
||||
)
|
||||
query.resp_messages = [command_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].result_type == entities.ResultType.CONTINUE
|
||||
# Check that prefix was added (via get_content_platform_message_chain)
|
||||
command_resp.get_content_platform_message_chain.assert_called_once()
|
||||
|
||||
|
||||
class TestResponseWrapperPlugin:
|
||||
"""Tests for plugin response wrapping."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_response_direct(self):
|
||||
"""Plugin response should be wrapped without prefix."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create a plugin response message
|
||||
plugin_resp = Mock()
|
||||
plugin_resp.role = 'plugin'
|
||||
plugin_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Plugin response")])
|
||||
)
|
||||
query.resp_messages = [plugin_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].result_type == entities.ResultType.CONTINUE
|
||||
|
||||
|
||||
class TestResponseWrapperAssistant:
|
||||
"""Tests for assistant response wrapping."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assistant_content_response(self):
|
||||
"""Assistant with content should emit event and wrap."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Mock session manager to return a valid Session
|
||||
session = make_session()
|
||||
app.sess_mgr.get_session = AsyncMock(return_value=session)
|
||||
|
||||
# Mock plugin connector - normal event (not prevented)
|
||||
mock_event_ctx = Mock()
|
||||
mock_event_ctx.is_prevented_default = Mock(return_value=False)
|
||||
mock_event_ctx.event = Mock()
|
||||
mock_event_ctx.event.reply_message_chain = None
|
||||
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
|
||||
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create assistant response with content
|
||||
assistant_resp = Mock()
|
||||
assistant_resp.role = 'assistant'
|
||||
assistant_resp.content = "Hello back!"
|
||||
assistant_resp.tool_calls = None
|
||||
assistant_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Hello back!")])
|
||||
)
|
||||
query.resp_messages = [assistant_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].result_type == entities.ResultType.CONTINUE
|
||||
# Event should have been emitted
|
||||
app.plugin_connector.emit_event.assert_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assistant_empty_content(self):
|
||||
"""Assistant with empty content should not emit event."""
|
||||
wrapper = get_wrapper_module()
|
||||
get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create assistant response with empty content
|
||||
assistant_resp = Mock()
|
||||
assistant_resp.role = 'assistant'
|
||||
assistant_resp.content = None
|
||||
assistant_resp.tool_calls = None
|
||||
query.resp_messages = [assistant_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
# Should have at least one result (for empty content case)
|
||||
assert len(results) >= 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assistant_tool_calls(self):
|
||||
"""Assistant with tool_calls should show function call message."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Mock session manager to return a valid Session
|
||||
session = make_session()
|
||||
app.sess_mgr.get_session = AsyncMock(return_value=session)
|
||||
|
||||
# Mock plugin connector
|
||||
mock_event_ctx = Mock()
|
||||
mock_event_ctx.is_prevented_default = Mock(return_value=False)
|
||||
mock_event_ctx.event = Mock()
|
||||
mock_event_ctx.event.reply_message_chain = None
|
||||
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
|
||||
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
pipeline_config['output']['misc']['track-function-calls'] = True
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create assistant response with tool_calls
|
||||
mock_tool_call = Mock()
|
||||
mock_tool_call.function = Mock()
|
||||
mock_tool_call.function.name = 'test_function'
|
||||
|
||||
assistant_resp = Mock()
|
||||
assistant_resp.role = 'assistant'
|
||||
assistant_resp.content = "Processing..."
|
||||
assistant_resp.tool_calls = [mock_tool_call]
|
||||
assistant_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Processing...")])
|
||||
)
|
||||
query.resp_messages = [assistant_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
# Should have results for content and tool_calls
|
||||
assert len(results) >= 1
|
||||
for result in results:
|
||||
assert result.result_type == entities.ResultType.CONTINUE
|
||||
|
||||
|
||||
class TestResponseWrapperInterrupt:
|
||||
"""Tests for INTERRUPT behavior when plugin prevents default."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_prevented_interrupts(self):
|
||||
"""Plugin event prevented should return INTERRUPT."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Mock session manager to return a valid Session
|
||||
session = make_session()
|
||||
app.sess_mgr.get_session = AsyncMock(return_value=session)
|
||||
|
||||
# Mock plugin connector - event is prevented
|
||||
mock_event_ctx = Mock()
|
||||
mock_event_ctx.is_prevented_default = Mock(return_value=True)
|
||||
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
|
||||
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create assistant response with content
|
||||
assistant_resp = Mock()
|
||||
assistant_resp.role = 'assistant'
|
||||
assistant_resp.content = "Hello!"
|
||||
assistant_resp.tool_calls = None
|
||||
assistant_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Hello!")])
|
||||
)
|
||||
query.resp_messages = [assistant_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].result_type == entities.ResultType.INTERRUPT
|
||||
|
||||
|
||||
class TestResponseWrapperCustomReply:
|
||||
"""Tests for custom reply from plugin event."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_reply_chain_used(self):
|
||||
"""Plugin reply_message_chain should replace default."""
|
||||
wrapper = get_wrapper_module()
|
||||
entities = get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Mock session manager to return a valid Session
|
||||
session = make_session()
|
||||
app.sess_mgr.get_session = AsyncMock(return_value=session)
|
||||
|
||||
# Mock plugin connector with custom reply
|
||||
custom_chain = platform_message.MessageChain([platform_message.Plain(text="Custom reply")])
|
||||
mock_event_ctx = Mock()
|
||||
mock_event_ctx.is_prevented_default = Mock(return_value=False)
|
||||
mock_event_ctx.event = Mock()
|
||||
mock_event_ctx.event.reply_message_chain = custom_chain
|
||||
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
|
||||
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
|
||||
# Create assistant response
|
||||
assistant_resp = Mock()
|
||||
assistant_resp.role = 'assistant'
|
||||
assistant_resp.content = "Default reply"
|
||||
assistant_resp.tool_calls = None
|
||||
assistant_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Default reply")])
|
||||
)
|
||||
query.resp_messages = [assistant_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].result_type == entities.ResultType.CONTINUE
|
||||
# Custom chain should be in resp_message_chain
|
||||
assert len(results[0].new_query.resp_message_chain) == 1
|
||||
# Should be the custom chain
|
||||
chain = results[0].new_query.resp_message_chain[0]
|
||||
assert "Custom reply" in str(chain)
|
||||
|
||||
|
||||
class TestResponseWrapperVariables:
|
||||
"""Tests for bound plugins variable."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bound_plugins_passed_to_event(self):
|
||||
"""_pipeline_bound_plugins should be passed to emit_event."""
|
||||
wrapper = get_wrapper_module()
|
||||
get_entities_module()
|
||||
|
||||
app = FakeApp()
|
||||
|
||||
# Mock session manager to return a valid Session
|
||||
session = make_session()
|
||||
app.sess_mgr.get_session = AsyncMock(return_value=session)
|
||||
|
||||
# Mock plugin connector
|
||||
mock_event_ctx = Mock()
|
||||
mock_event_ctx.is_prevented_default = Mock(return_value=False)
|
||||
mock_event_ctx.event = Mock()
|
||||
mock_event_ctx.event.reply_message_chain = None
|
||||
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
|
||||
|
||||
stage = wrapper.ResponseWrapper(app)
|
||||
|
||||
pipeline_config = make_wrapper_config()
|
||||
|
||||
await stage.initialize(pipeline_config)
|
||||
|
||||
query = text_query("hello")
|
||||
query.pipeline_config = pipeline_config
|
||||
query.resp_message_chain = []
|
||||
query.variables['_pipeline_bound_plugins'] = ['plugin1', 'plugin2']
|
||||
|
||||
# Create assistant response
|
||||
assistant_resp = Mock()
|
||||
assistant_resp.role = 'assistant'
|
||||
assistant_resp.content = "Hello"
|
||||
assistant_resp.tool_calls = None
|
||||
assistant_resp.get_content_platform_message_chain = Mock(
|
||||
return_value=platform_message.MessageChain([platform_message.Plain(text="Hello")])
|
||||
)
|
||||
query.resp_messages = [assistant_resp]
|
||||
|
||||
results = []
|
||||
async for result in stage.process(query, 'ResponseWrapper'):
|
||||
results.append(result)
|
||||
|
||||
# Check that bound_plugins was passed
|
||||
emit_call = app.plugin_connector.emit_event.call_args
|
||||
assert emit_call[0][1] == ['plugin1', 'plugin2'] # Second argument is bound_plugins
|
||||
Reference in New Issue
Block a user