chore: import upstream snapshot with attribution
Test Migrations / Migrations (SQLite) (push) Has been cancelled
Build Dev Image / build-dev-image (push) Has been cancelled
Check i18n Keys / Check i18n Key Consistency (push) Has been cancelled
Lint / Ruff Lint & Format (push) Has been cancelled
Lint / Frontend Lint (push) Has been cancelled
Test Migrations / Migrations (PostgreSQL) (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:44:22 +08:00
commit 534bb94eea
1288 changed files with 266913 additions and 0 deletions
+179
View File
@@ -0,0 +1,179 @@
# 单元测试覆盖率排除说明
## 排除范围
以下外部适配器模块不纳入测试覆盖目标,因为它们需要实际外部环境才能测试:
### 1. 消息平台适配器 (`platform/sources/`)
- **路径**: `src/langbot/pkg/platform/sources/`
- **模块**: aiocqhttp, dingtalk, discord, feishu, gestep, kook, lark, slack, telegram, wecom, wechatpv, wechatmp, qqbot
- **排除原因**: 需要真实消息平台账号和 webhook 连接,无法纯单元测试
- **测试方式**: 需要 mock 平台 API 或集成测试环境
- **状态**: 后续可补充 mock 测试
### 2. LLM Requester (`provider/modelmgr/requesters/`)
- **路径**: `src/langbot/pkg/provider/modelmgr/requesters/`
- **模块**: deepseek, openai, anthropic, gemini, moonshot, ollama, zhipuai 等 20+ 个 requester
- **排除原因**: 需要真实 LLM API 密钥和网络请求,涉及付费 API 调用
- **测试方式**: 需要 mock HTTP 响应或使用 fake LLM server
- **状态**: 后续可补充 mock HTTP 测试
### 3. Agent Runner (`provider/runners/`)
- **路径**: `src/langbot/pkg/provider/runners/`
- **模块**: cozeapi, difysvapi, n8nsvapi, langflowapi, dashscopeapi, localagent, tboxapi
- **排除原因**: 需要真实 Agent 平台(Coze、Dify、n8n 等)的 API 连接
- **测试方式**: 需要 mock Agent 平台响应
- **状态**: 后续可补充 mock 测试
### 4. 向量数据库 (`vector/vdbs/`)
- **路径**: `src/langbot/pkg/vector/vdbs/`
- **模块**: chroma, milvus, pgvector, qdrant, seekdb, valkey_search
- **排除原因**: 需要真实向量数据库实例运行
- **测试方式**: 需要 Docker 启动测试数据库或 mock
- **状态**: 后续可补充 mock 测试
---
## 覆盖率计算(排除外部适配器)
### 统计方法
```bash
# 排除外部适配器后计算覆盖率
pytest tests/unit_tests/ --cov=langbot.pkg \
--cov-fail-under=0 \
-o "cov_exclude_patterns=platform/sources/*,provider/modelmgr/requesters/*,provider/runners/*,vector/vdbs/*"
```
### 当前覆盖率(排除后)
| 模块 | 覆盖率 | 状态 |
|------|--------|------|
| `command` | **99%** | ✅ 完成 |
| `entity` | **99%** | ✅ 完成 |
| `vector` | **76%** | ✅ 完成 |
| `survey` | **84%** | ✅ 完成 |
| `pipeline` | **72%** | ✅ 核心流程 |
| `rag` | **66%** | ✅ 完成 |
| `telemetry` | **87%** | ✅ 完成 |
| `storage` | **80%** | ✅ 完成 |
| `provider` | **83%** | ✅ 完成 |
| `discover` | **61%** | ✅ 完成 |
| `config` | **70%** | ✅ 完成 |
| `utils` | **48%** | 🔄 部分完成 |
| `api` | **34%** | 🔄 需补充 controller |
| `platform` | **35%** | 🔄 需补充 adapter base |
| `plugin` | **27%** | 🔄 需补充 handler |
| `core` | **28%** | 🔄 需补充 app 启动 |
| `persistence` | **24%** | 🔄 需补充 mgr |
---
## 后续计划
### 可补充的 Mock 测试(优先级排序)
1. **`provider/modelmgr/requesters/`** (优先级:中)
- 使用 `httpx` mock 测试 API 响应解析
- 测试重试逻辑、错误处理
2. **`provider/runners/`** (优先级:中)
- Mock Agent 平台响应
- 测试 session 管理、错误处理
3. **`platform/sources/`** (优先级:低)
- Mock 平台 webhook 事件
- 测试消息解析、事件处理
4. **`vector/vdbs/`** (优先级:低)
- Mock 向量数据库操作
- 测试 CRUD、查询逻辑
---
## 测试文件结构
```
tests/unit_tests/
├── api/
│ └── service/
│ ├── test_knowledge_service.py # 22 tests ✅
│ └── ...
├── core/
│ ├── test_taskmgr.py # 21 tests ✅
│ ├── test_load_config.py # 21 tests ✅ (含env override)
│ └── ...
├── plugin/
│ ├── test_connector_static.py # 8 tests ✅
│ ├── test_connector_pure.py # 7 tests ✅
│ ├── test_connector_methods.py # 24 tests ✅
│ ├── test_extract_deps.py # 7 tests ✅
│ ├── test_handler_actions.py # 15 tests ✅ (新增)
│ └── ...
├── provider/
│ ├── test_session_manager.py # 11 tests ✅ (新增)
│ ├── test_tool_manager.py # 14 tests ✅ (新增)
│ └── ...
├── rag/
│ ├── test_i18n_conversion.py # 8 tests ✅
│ ├── test_kbmgr.py # 39 tests ✅
│ ├── test_file_storage.py # 21 tests ✅ (新增)
│ └── ...
├── storage/
│ ├── test_s3storage.py # 16 tests ✅ (新增)
│ ├── test_localstorage_path_traversal.py # 11 tests ✅
│ └── ...
├── survey/
│ └── test_survey_manager.py # 22 tests ✅
├── telemetry/
│ └── test_telemetry.py # 25 tests ✅ (重写)
├── vector/
│ ├── test_filter_utils.py # 21 tests ✅
│ ├── test_vdb_filter_conversion.py # 30 tests ✅ (新增)
│ └── ...
├── utils/
│ ├── test_platform.py # 7 tests ✅
│ ├── test_funcschema.py # 9 tests ✅
│ └── ...
├── pipeline/
│ ├── test_ratelimit.py # 12 tests ✅ (新增真实算法)
│ ├── test_msgtrun.py # 9 tests ✅ (强化断言)
│ └── ...
└── persistence/
├── test_serialize_model.py # 6 tests ✅
├── test_database_decorator.py # 7 tests ✅
└── ...
```
---
## 总结
- **总测试数**: 1193 passed
- **总体覆盖率**: 30%
- **核心模块覆盖率**: **51.2%** (6549/12825 语句) - 排除外部适配器
- **外部适配器覆盖率**: 5.6% (535/9483 语句) - 不纳入目标
### 核心模块覆盖率详情
| 模块 | 覆盖率 | 语句数 | 说明 |
|------|--------|--------|------|
| `command` | **99%** | 93 | ✅ 完成 |
| `entity` | **99%** | 335 | ✅ 完成 |
| `vector` | **76%** | 139 | ✅ 完成 (新增filter转换测试) |
| `survey` | **84%** | 95 | ✅ 完成 |
| `pipeline` | **72%** | 1761 | ✅ 核心流程 (新增算法测试) |
| `rag` | **66%** | 347 | ✅ 完成 (新增ZIP处理测试) |
| `telemetry` | **87%** | 70 | ✅ 完成 (重写假测试) |
| `storage` | **80%** | 170 | ✅ 完成 (新增S3测试) |
| `provider` | **83%** | 854 | ✅ 完成 (新增Session/Tool测试) |
| `discover` | **61%** | 188 | ✅ 完成 |
| `config` | **70%** | 198 | ✅ 完成 |
| `utils` | **48%** | 478 | 🔄 部分完成 |
| `api` | **34%** | 4061 | 🔄 需补充 controller |
| `platform` | **35%** | 433 | 🔄 需补充 adapter base |
| `plugin` | **27%** | 815 | 🔄 需补充 handler (新增action测试) |
| `core` | **28%** | 1289 | 🔄 需补充 app 启动 |
| `persistence` | **24%** | 1099 | 🔄 需补充 mgr |
外部适配器测试需要 mock 环境或集成测试,不属于纯单元测试范畴。
View File
+1
View File
@@ -0,0 +1 @@
"""Unit tests for LangBot API HTTP service layer."""
@@ -0,0 +1,62 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock
from sqlalchemy.sql.dml import Update
from langbot.pkg.api.http.service.bot import BotService
class _FakeResult:
def __init__(self, value):
self.value = value
def first(self):
return self.value
class _PersistenceManager:
def __init__(self):
self.update_values = None
async def execute_async(self, statement):
if isinstance(statement, Update):
self.update_values = {
key: value for key, value in statement.compile().params.items() if not key.startswith('uuid_')
}
return None
return _FakeResult(SimpleNamespace(name='Updated Pipeline'))
async def test_update_bot_copies_input_before_filtering_and_setting_pipeline_name():
persistence_mgr = _PersistenceManager()
runtime_bot = SimpleNamespace(enable=False)
platform_mgr = SimpleNamespace(
remove_bot=AsyncMock(),
load_bot=AsyncMock(return_value=runtime_bot),
)
ap = SimpleNamespace(
persistence_mgr=persistence_mgr,
platform_mgr=platform_mgr,
sess_mgr=SimpleNamespace(session_list=[]),
)
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'bot-1'})
payload = {
'uuid': 'caller-owned-uuid',
'name': 'Test Bot',
'use_pipeline_uuid': 'pipeline-1',
}
await service.update_bot('bot-1', payload)
assert payload == {
'uuid': 'caller-owned-uuid',
'name': 'Test Bot',
'use_pipeline_uuid': 'pipeline-1',
}
assert persistence_mgr.update_values == {
'name': 'Test Bot',
'use_pipeline_uuid': 'pipeline-1',
'use_pipeline_name': 'Updated Pipeline',
}
+16
View File
@@ -0,0 +1,16 @@
"""Unit tests for API HTTP service layer.
Tests real service business logic with mocked dependencies:
- persistence_mgr (database operations)
- model_mgr (runtime model management)
- platform_mgr (platform management)
- plugin_connector (plugin runtime)
- adjacent services (cross-service calls)
Does NOT:
- Start real Quart server
- Access real database
- Call real provider/platform/network
Uses tests.factories.FakeApp as base mock application.
"""
@@ -0,0 +1,482 @@
"""
Unit tests for ApiKeyService.
Tests API key CRUD operations with mocked persistence layer.
Source: src/langbot/pkg/api/http/service/apikey.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, patch
from types import SimpleNamespace
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.entity.persistence.apikey import ApiKey
pytestmark = pytest.mark.asyncio
class TestApiKeyServiceGetApiKeys:
"""Tests for get_api_keys method."""
async def test_get_api_keys_empty_list(self):
"""Returns empty list when no API keys exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.all = Mock(return_value=[])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'key': entity.key,
'description': entity.description,
}
if entity
else {}
)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_keys()
# Verify
assert result == []
ap.persistence_mgr.execute_async.assert_called_once()
async def test_get_api_keys_returns_serialized_list(self):
"""Returns serialized list of API keys."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Create mock API key entities
key1 = Mock(spec=ApiKey)
key1.id = 1
key1.name = 'Test Key 1'
key1.key = 'lbk_test_key_1'
key1.description = 'First test key'
key2 = Mock(spec=ApiKey)
key2.id = 2
key2.name = 'Test Key 2'
key2.key = 'lbk_test_key_2'
key2.description = 'Second test key'
mock_result = Mock()
mock_result.all = Mock(return_value=[key1, key2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'key': entity.key,
'description': entity.description,
}
)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_keys()
# Verify
assert len(result) == 2
assert result[0]['name'] == 'Test Key 1'
assert result[1]['name'] == 'Test Key 2'
class TestApiKeyServiceCreateApiKey:
"""Tests for create_api_key method."""
async def test_create_api_key_generates_key_with_prefix(self):
"""Creates API key with 'lbk_' prefix."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
created_key = Mock(spec=ApiKey)
created_key.id = 1
created_key.name = 'New Key'
created_key.key = 'lbk_fixed-token'
created_key.description = 'Test description'
select_result = Mock()
select_result.first = Mock(return_value=created_key)
insert_params = []
async def mock_execute(query):
params = query.compile().params
if {'name', 'key', 'description'}.issubset(params):
insert_params.append(params)
return Mock()
return select_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': 1,
'name': entity.name,
'key': entity.key,
'description': entity.description,
}
)
service = ApiKeyService(ap)
with patch('langbot.pkg.api.http.service.apikey.secrets.token_urlsafe', return_value='fixed-token'):
result = await service.create_api_key('New Key', 'Test description')
assert insert_params == [{'name': 'New Key', 'key': 'lbk_fixed-token', 'description': 'Test description'}]
assert result['key'].startswith('lbk_')
assert result['key'] == 'lbk_fixed-token'
assert result['name'] == 'New Key'
assert result['description'] == 'Test description'
async def test_create_api_key_without_description(self):
"""Creates API key with empty description when not provided."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
created_key = Mock(spec=ApiKey)
created_key.id = 1
created_key.name = 'No Desc Key'
created_key.key = 'lbk_no_desc_key'
created_key.description = ''
select_result = Mock()
select_result.first = Mock(return_value=created_key)
insert_result = Mock()
async def mock_execute(query):
if hasattr(query, 'values'):
return insert_result
return select_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'No Desc Key',
'key': 'lbk_no_desc_key',
'description': '',
}
)
service = ApiKeyService(ap)
# Execute
result = await service.create_api_key('No Desc Key')
# Verify
assert result['description'] == ''
class TestApiKeyServiceGetApiKey:
"""Tests for get_api_key method."""
async def test_get_api_key_by_id_found(self):
"""Returns API key when found by ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
key = Mock(spec=ApiKey)
key.id = 1
key.name = 'Found Key'
key.key = 'lbk_found_key'
key.description = 'Found'
mock_result = Mock()
mock_result.first = Mock(return_value=key)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'Found Key',
'key': 'lbk_found_key',
'description': 'Found',
}
)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_key(1)
# Verify
assert result is not None
assert result['id'] == 1
assert result['name'] == 'Found Key'
async def test_get_api_key_by_id_not_found(self):
"""Returns None when API key not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_key(999)
# Verify
assert result is None
async def test_get_api_key_by_id_zero(self):
"""Handles ID=0 (edge case) correctly."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_key(0)
# Verify - should return None (no key with ID 0)
assert result is None
class TestApiKeyServiceVerifyApiKey:
"""Tests for verify_api_key method."""
@staticmethod
def _make_ap(db_key=None, global_api_key=''):
"""Build a mock Application with persistence + instance_config."""
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=db_key)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.instance_config = SimpleNamespace(data={'api': {'global_api_key': global_api_key}})
return ap
async def test_verify_api_key_valid(self):
"""Returns True for valid API key."""
# Setup
key = Mock(spec=ApiKey)
ap = self._make_ap(db_key=key)
service = ApiKeyService(ap)
# Execute
result = await service.verify_api_key('lbk_valid_key')
# Verify
assert result is True
async def test_verify_api_key_invalid(self):
"""Returns False for invalid API key."""
# Setup
ap = self._make_ap(db_key=None)
service = ApiKeyService(ap)
# Execute
result = await service.verify_api_key('lbk_invalid_key')
# Verify
assert result is False
async def test_verify_api_key_empty_string(self):
"""Returns False for empty key string."""
# Setup
ap = self._make_ap(db_key=None)
service = ApiKeyService(ap)
# Execute
result = await service.verify_api_key('')
# Verify
assert result is False
async def test_verify_api_key_unknown_key(self):
"""Returns False when the key is not present in persistence."""
# Setup
ap = self._make_ap(db_key=None)
service = ApiKeyService(ap)
# Execute
result = await service.verify_api_key('unknown_key')
# Verify
assert result is False
async def test_verify_global_api_key_match(self):
"""Returns True when key matches the config.yaml global API key (no DB lookup)."""
# Setup: no DB record, but a global key is configured
ap = self._make_ap(db_key=None, global_api_key='my-global-secret')
service = ApiKeyService(ap)
# Execute
result = await service.verify_api_key('my-global-secret')
# Verify: accepted purely on config match
assert result is True
# DB should not have been consulted for the global-key path
ap.persistence_mgr.execute_async.assert_not_called()
async def test_verify_global_api_key_no_prefix_required(self):
"""Global API key is accepted even without the lbk_ prefix."""
ap = self._make_ap(db_key=None, global_api_key='plainsecret123')
service = ApiKeyService(ap)
result = await service.verify_api_key('plainsecret123')
assert result is True
async def test_verify_global_api_key_mismatch_falls_back_to_db(self):
"""A non-matching key still falls through to the DB lookup."""
# Global key set, but request uses a different lbk_ key that IS in DB
key = Mock(spec=ApiKey)
ap = self._make_ap(db_key=key, global_api_key='my-global-secret')
service = ApiKeyService(ap)
result = await service.verify_api_key('lbk_db_key')
assert result is True
ap.persistence_mgr.execute_async.assert_called_once()
async def test_verify_empty_global_api_key_disabled(self):
"""An empty global_api_key must never authenticate an empty/blank request."""
ap = self._make_ap(db_key=None, global_api_key='')
service = ApiKeyService(ap)
# Empty request key is rejected, and a blank global key never matches
assert await service.verify_api_key('') is False
assert await service.verify_api_key(' ') is False
async def test_verify_api_key_missing_global_config_key(self):
"""Works even when api.global_api_key is absent (existing installs)."""
# instance_config without the global_api_key field at all
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.instance_config = SimpleNamespace(data={'api': {}})
service = ApiKeyService(ap)
result = await service.verify_api_key('lbk_some_key')
assert result is False
class TestApiKeyServiceDeleteApiKey:
"""Tests for delete_api_key method."""
async def test_delete_api_key_by_id(self):
"""Deletes API key by ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute
await service.delete_api_key(1)
# Verify - execute_async was called (delete operation)
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_api_key_nonexistent_id(self):
"""Delete operation completes even for nonexistent ID (no error raised)."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute - should not raise error
await service.delete_api_key(999)
# Verify - execute_async was called regardless
ap.persistence_mgr.execute_async.assert_called_once()
class TestApiKeyServiceUpdateApiKey:
"""Tests for update_api_key method."""
async def test_update_api_key_name_only(self):
"""Updates only the name field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute
await service.update_api_key(1, name='Updated Name')
# Verify - execute_async was called with update
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_api_key_description_only(self):
"""Updates only the description field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute
await service.update_api_key(1, description='Updated description')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_api_key_both_fields(self):
"""Updates both name and description."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute
await service.update_api_key(1, name='New Name', description='New description')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_api_key_no_fields(self):
"""Does nothing when no fields provided."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute
await service.update_api_key(1)
# Verify - no execute call since no update_data
ap.persistence_mgr.execute_async.assert_not_called()
@@ -0,0 +1,650 @@
"""
Unit tests for BotService.
Tests bot CRUD operations with mocked persistence and runtime managers.
Source: src/langbot/pkg/api/http/service/bot.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, patch
from types import SimpleNamespace
import uuid
from langbot.pkg.api.http.service.bot import BotService
from langbot.pkg.entity.persistence.bot import Bot
pytestmark = pytest.mark.asyncio
def _create_mock_bot(
bot_uuid: str = None,
name: str = 'Test Bot',
description: str = 'Test Description',
adapter: str = 'telegram',
adapter_config: dict = None,
enable: bool = True,
use_pipeline_uuid: str = None,
use_pipeline_name: str = None,
) -> Mock:
"""Helper to create mock Bot entity."""
bot = Mock(spec=Bot)
bot.uuid = bot_uuid or str(uuid.uuid4())
bot.name = name
bot.description = description
bot.adapter = adapter
bot.adapter_config = adapter_config or {'token': 'test_token'}
bot.enable = enable
bot.use_pipeline_uuid = use_pipeline_uuid
bot.use_pipeline_name = use_pipeline_name
bot.pipeline_routing_rules = []
return bot
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
class TestBotServiceGetBots:
"""Tests for get_bots method."""
async def test_get_bots_empty_list(self):
"""Returns empty list when no bots exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity, masked_columns=None: {
'uuid': entity.uuid,
'name': entity.name,
'adapter': entity.adapter,
}
)
service = BotService(ap)
# Execute
result = await service.get_bots()
# Verify
assert result == []
async def test_get_bots_returns_list_with_secrets(self):
"""Returns bot list including adapter_config by default."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
bot1 = _create_mock_bot(bot_uuid='uuid-1', name='Bot 1')
bot2 = _create_mock_bot(bot_uuid='uuid-2', name='Bot 2')
mock_result = _create_mock_result([bot1, bot2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity, masked_columns=None: {
'uuid': entity.uuid,
'name': entity.name,
'adapter': entity.adapter,
'adapter_config': entity.adapter_config if 'adapter_config' not in (masked_columns or []) else None,
}
)
service = BotService(ap)
# Execute
result = await service.get_bots(include_secret=True)
# Verify
assert len(result) == 2
assert result[0]['name'] == 'Bot 1'
assert result[0]['adapter_config'] is not None
async def test_get_bots_masks_secrets(self):
"""Returns bot list without adapter_config when include_secret=False."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
bot1 = _create_mock_bot(bot_uuid='uuid-1', name='Bot 1')
mock_result = _create_mock_result([bot1])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity, masked_columns=None: {
'uuid': entity.uuid,
'name': entity.name,
'adapter': entity.adapter,
'adapter_config': entity.adapter_config if 'adapter_config' not in (masked_columns or []) else None,
}
)
service = BotService(ap)
# Execute
result = await service.get_bots(include_secret=False)
# Verify - adapter_config should be masked
assert result[0]['adapter_config'] is None
class TestBotServiceGetBot:
"""Tests for get_bot method."""
async def test_get_bot_by_uuid_found(self):
"""Returns bot when found by UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
bot = _create_mock_bot(bot_uuid='test-uuid', name='Found Bot')
mock_result = _create_mock_result(first_item=bot)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Found Bot',
'adapter': 'telegram',
}
)
service = BotService(ap)
# Execute
result = await service.get_bot('test-uuid')
# Verify
assert result is not None
assert result['uuid'] == 'test-uuid'
assert result['name'] == 'Found Bot'
async def test_get_bot_by_uuid_not_found(self):
"""Returns None when bot not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = BotService(ap)
# Execute
result = await service.get_bot('nonexistent-uuid')
# Verify
assert result is None
class TestBotServiceGetRuntimeBotInfo:
"""Tests for get_runtime_bot_info method."""
async def test_get_runtime_bot_info_bot_not_found_raises(self):
"""Raises Exception when bot not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = BotService(ap)
# Mock get_bot to return None
service.get_bot = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(Exception, match='Bot not found'):
await service.get_runtime_bot_info('nonexistent-uuid')
async def test_get_runtime_bot_info_returns_webhook_for_wecom(self):
"""Returns webhook URL for wecom adapter."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'api': {
'webhook_prefix': 'http://127.0.0.1:5300',
'extra_webhook_prefix': 'http://extra.example.com',
}
}
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=None)
bot_data = {
'uuid': 'wecom-uuid',
'name': 'WeCom Bot',
'adapter': 'wecom',
'adapter_config': {'token': 'test'},
}
service = BotService(ap)
service.get_bot = AsyncMock(return_value=bot_data)
# Execute
result = await service.get_runtime_bot_info('wecom-uuid')
# Verify
assert result['adapter_runtime_values']['webhook_url'] == '/bots/wecom-uuid'
assert result['adapter_runtime_values']['webhook_full_url'] == 'http://127.0.0.1:5300/bots/wecom-uuid'
async def test_get_runtime_bot_info_no_webhook_for_telegram(self):
"""Returns no webhook URL for non-webhook adapters like telegram."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'api': {}}
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=None)
bot_data = {
'uuid': 'telegram-uuid',
'name': 'Telegram Bot',
'adapter': 'telegram',
'adapter_config': {'token': 'test'},
}
service = BotService(ap)
service.get_bot = AsyncMock(return_value=bot_data)
# Execute
result = await service.get_runtime_bot_info('telegram-uuid')
# Verify - no webhook for telegram
assert result['adapter_runtime_values']['webhook_url'] is None
assert result['adapter_runtime_values']['webhook_full_url'] is None
async def test_get_runtime_bot_info_with_runtime_bot(self):
"""Returns bot_account_id when runtime bot exists."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'api': {}}
ap.platform_mgr = SimpleNamespace()
# Mock runtime bot with adapter
runtime_bot = SimpleNamespace()
runtime_bot.adapter = SimpleNamespace()
runtime_bot.adapter.bot_account_id = 'runtime-account-123'
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
bot_data = {
'uuid': 'runtime-uuid',
'name': 'Runtime Bot',
'adapter': 'telegram',
'adapter_config': {},
}
service = BotService(ap)
service.get_bot = AsyncMock(return_value=bot_data)
# Execute
result = await service.get_runtime_bot_info('runtime-uuid')
# Verify
assert result['adapter_runtime_values']['bot_account_id'] == 'runtime-account-123'
class TestBotServiceCreateBot:
"""Tests for create_bot method."""
async def test_create_bot_max_limit_reached_raises(self):
"""Raises ValueError when max_bots limit reached."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_bots': 2}}}
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.load_bot = AsyncMock()
# Mock get_bots to return 2 bots already
bot1 = _create_mock_bot(bot_uuid='uuid-1')
bot2 = _create_mock_bot(bot_uuid='uuid-2')
mock_result = _create_mock_result([bot1, bot2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Bot 1'})
service = BotService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='Maximum number of bots'):
await service.create_bot({'name': 'New Bot'})
async def test_create_bot_no_limit(self):
"""Creates bot without limit check when max_bots=-1."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'system': {
'limitation': {
'max_bots': -1 # No limit
}
}
}
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.load_bot = AsyncMock()
# Mock pipeline query
pipeline_result = Mock()
pipeline_result.first = Mock(return_value=None)
# Mock bot query after insert
bot_result = Mock()
bot_result.first = Mock(return_value=_create_mock_bot())
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count <= 2:
return pipeline_result # First call: check pipeline
elif call_count == 3:
return Mock() # Insert
return bot_result # Get bot
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid', 'name': 'New Bot'})
service = BotService(ap)
# Execute
bot_uuid = await service.create_bot({'name': 'New Bot', 'adapter': 'telegram', 'adapter_config': {}})
# Verify
assert bot_uuid is not None
assert len(bot_uuid) == 36 # UUID format
async def test_create_bot_sets_default_pipeline(self):
"""Sets default pipeline when one exists."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_bots': -1}}}
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.load_bot = AsyncMock()
# Mock default pipeline
mock_pipeline = SimpleNamespace()
mock_pipeline.uuid = 'default-pipeline-uuid'
mock_pipeline.name = 'Default Pipeline'
pipeline_result = Mock()
pipeline_result.first = Mock(return_value=mock_pipeline)
# Mock bot after insert
bot_result = Mock()
bot_result.first = Mock(return_value=_create_mock_bot())
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return pipeline_result # Check default pipeline
elif call_count == 2:
return Mock() # Insert
return bot_result # Get bot
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'new-uuid',
'name': 'New Bot',
'use_pipeline_uuid': 'default-pipeline-uuid',
'use_pipeline_name': 'Default Pipeline',
}
)
service = BotService(ap)
# Execute
bot_data = {'name': 'New Bot', 'adapter': 'telegram', 'adapter_config': {}}
bot_uuid = await service.create_bot(bot_data)
# Verify - pipeline uuid and name were set
assert 'use_pipeline_uuid' in bot_data
assert 'use_pipeline_name' in bot_data
assert bot_uuid is not None # Verify UUID was returned
class TestBotServiceUpdateBot:
"""Tests for update_bot method."""
async def test_update_bot_removes_uuid_from_data(self):
"""Does not persist caller-provided uuid in update payload."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.remove_bot = AsyncMock()
# Mock pipeline query - not updating pipeline
ap.persistence_mgr.execute_async = AsyncMock()
ap.sess_mgr = SimpleNamespace()
ap.sess_mgr.session_list = []
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'test-uuid', 'name': 'Updated'})
# Create mock runtime bot
runtime_bot = SimpleNamespace()
runtime_bot.enable = False
ap.platform_mgr.load_bot = AsyncMock(return_value=runtime_bot)
# Execute
update_data = {'uuid': 'should-be-removed', 'name': 'Updated Name'}
await service.update_bot('test-uuid', update_data)
update_params = ap.persistence_mgr.execute_async.await_args_list[0].args[0].compile().params
assert update_params['name'] == 'Updated Name'
assert 'should-be-removed' not in update_params.values()
async def test_update_bot_pipeline_not_found_raises(self):
"""Raises Exception when updating with nonexistent pipeline UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Mock pipeline query returns None
pipeline_result = Mock()
pipeline_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=pipeline_result)
service = BotService(ap)
# Execute & Verify
with pytest.raises(Exception, match='Pipeline not found'):
await service.update_bot('test-uuid', {'use_pipeline_uuid': 'nonexistent-pipeline'})
async def test_update_bot_sets_pipeline_name(self):
"""Sets use_pipeline_name when updating use_pipeline_uuid."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.remove_bot = AsyncMock()
# Mock pipeline query
mock_pipeline = SimpleNamespace()
mock_pipeline.name = 'Updated Pipeline'
pipeline_result = Mock()
pipeline_result.first = Mock(return_value=mock_pipeline)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return pipeline_result
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.sess_mgr = SimpleNamespace()
ap.sess_mgr.session_list = []
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'test-uuid'})
runtime_bot = SimpleNamespace()
runtime_bot.enable = False
ap.platform_mgr.load_bot = AsyncMock(return_value=runtime_bot)
# Execute
await service.update_bot('test-uuid', {'use_pipeline_uuid': 'pipeline-uuid'})
update_params = ap.persistence_mgr.execute_async.await_args_list[1].args[0].compile().params
assert update_params['use_pipeline_uuid'] == 'pipeline-uuid'
assert update_params['use_pipeline_name'] == 'Updated Pipeline'
class TestBotServiceDeleteBot:
"""Tests for delete_bot method."""
async def test_delete_bot_calls_remove_and_delete(self):
"""Calls both platform_mgr.remove_bot and persistence delete."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.remove_bot = AsyncMock()
service = BotService(ap)
# Execute
await service.delete_bot('test-uuid')
# Verify
ap.platform_mgr.remove_bot.assert_called_once_with('test-uuid')
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_bot_nonexistent_uuid(self):
"""Delete operation completes even for nonexistent UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.remove_bot = AsyncMock()
service = BotService(ap)
# Execute - should not raise
await service.delete_bot('nonexistent-uuid')
# Verify - both called regardless
ap.platform_mgr.remove_bot.assert_called_once()
class TestBotServiceListEventLogs:
"""Tests for list_event_logs method."""
async def test_list_event_logs_bot_not_found_raises(self):
"""Raises Exception when runtime bot not found."""
# Setup
ap = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=None)
service = BotService(ap)
# Execute & Verify
with pytest.raises(Exception, match='Bot not found'):
await service.list_event_logs('nonexistent-uuid', 0, 10)
async def test_list_event_logs_returns_logs(self):
"""Returns logs from runtime bot logger."""
# Setup
ap = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
# Mock runtime bot with logger
runtime_bot = SimpleNamespace()
runtime_bot.logger = SimpleNamespace()
runtime_bot.logger.get_logs = AsyncMock(
return_value=([SimpleNamespace(to_json=Mock(return_value={'msg': 'log1'}))], 5)
)
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
service = BotService(ap)
# Execute
logs, total = await service.list_event_logs('bot-uuid', 0, 10)
# Verify
assert len(logs) == 1
assert logs[0] == {'msg': 'log1'}
assert total == 5
class TestBotServiceSendMessage:
"""Tests for send_message method."""
async def test_send_message_bot_not_found_raises(self):
"""Raises Exception when bot not found."""
# Setup
ap = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=None)
service = BotService(ap)
# Execute & Verify
with pytest.raises(Exception, match='Bot not found'):
await service.send_message('nonexistent-uuid', 'group', '123', {'test': 'data'})
async def test_send_message_invalid_message_chain_raises(self):
"""Raises Exception when message_chain_data is invalid."""
# Setup
ap = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
runtime_bot = SimpleNamespace()
runtime_bot.adapter = SimpleNamespace()
runtime_bot.adapter.send_message = AsyncMock()
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
service = BotService(ap)
# Execute & Verify - invalid format should raise
with pytest.raises(Exception, match='Invalid message_chain format'):
await service.send_message('bot-uuid', 'group', '123', {'invalid': 'format'})
async def test_send_message_valid_call(self):
"""Sends message through adapter when all valid."""
# Setup
ap = SimpleNamespace()
ap.platform_mgr = SimpleNamespace()
runtime_bot = SimpleNamespace()
runtime_bot.adapter = SimpleNamespace()
runtime_bot.adapter.send_message = AsyncMock()
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
service = BotService(ap)
# Execute with valid message chain format
message_chain_data = {'messages': [{'type': 'text', 'data': {'text': 'Hello'}}]}
# Patch the import location - the module imports inside the function
with patch('langbot_plugin.api.entities.builtin.platform.message.MessageChain') as MockMessageChain:
mock_chain = Mock()
MockMessageChain.model_validate = Mock(return_value=mock_chain)
await service.send_message('bot-uuid', 'group', '123', message_chain_data)
# Verify adapter.send_message was called
runtime_bot.adapter.send_message.assert_called_once_with('group', '123', mock_chain)
@@ -0,0 +1,389 @@
"""Unit tests for API knowledge service.
Tests cover:
- Knowledge base CRUD operations
- Capability checking
- Knowledge engine discovery
- File operations
"""
from __future__ import annotations
import pytest
from unittest.mock import Mock, AsyncMock
from importlib import import_module
def get_knowledge_service_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.api.http.service.knowledge')
def create_mock_app():
"""Create mock Application for testing."""
mock_app = Mock()
mock_app.logger = Mock()
mock_app.rag_mgr = AsyncMock()
mock_app.persistence_mgr = AsyncMock()
mock_app.persistence_mgr.execute_async = AsyncMock()
mock_app.persistence_mgr.serialize_model = Mock(return_value={})
mock_app.plugin_connector = AsyncMock()
mock_app.plugin_connector.is_enable_plugin = True
return mock_app
class TestKnowledgeServiceInit:
"""Tests for KnowledgeService initialization."""
def test_init_stores_app_reference(self):
"""Test that __init__ stores Application reference."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
service = knowledge_module.KnowledgeService(mock_app)
assert service.ap is mock_app
class TestGetKnowledgeBases:
"""Tests for get_knowledge_bases method."""
@pytest.mark.asyncio
async def test_returns_all_kb_details(self):
"""Test that it returns all knowledge base details."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_all_knowledge_base_details = AsyncMock(return_value=[{'uuid': 'kb1', 'name': 'KB1'}])
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_bases()
assert len(result) == 1
assert result[0]['uuid'] == 'kb1'
@pytest.mark.asyncio
async def test_returns_empty_list_when_no_kbs(self):
"""Test that it returns empty list when no knowledge bases."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_all_knowledge_base_details = AsyncMock(return_value=[])
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_bases()
assert result == []
class TestGetKnowledgeBase:
"""Tests for get_knowledge_base method."""
@pytest.mark.asyncio
async def test_returns_kb_details_by_uuid(self):
"""Test that it returns specific KB details."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value={'uuid': 'kb1', 'name': 'KB1'})
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_base('kb1')
assert result['uuid'] == 'kb1'
@pytest.mark.asyncio
async def test_returns_none_when_not_found(self):
"""Test that it returns None when KB not found."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value=None)
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_base('nonexistent')
assert result is None
class TestCreateKnowledgeBase:
"""Tests for create_knowledge_base method."""
@pytest.mark.asyncio
async def test_creates_kb_with_required_fields(self):
"""Test creating KB with required plugin ID."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_kb = Mock()
mock_kb.uuid = 'new_kb_uuid'
mock_app.rag_mgr.create_knowledge_base = AsyncMock(return_value=mock_kb)
service = knowledge_module.KnowledgeService(mock_app)
kb_data = {
'name': 'Test KB',
'knowledge_engine_plugin_id': 'author/engine',
'description': 'Test description',
}
result = await service.create_knowledge_base(kb_data)
assert result == 'new_kb_uuid'
mock_app.rag_mgr.create_knowledge_base.assert_called_once()
@pytest.mark.asyncio
async def test_raises_when_missing_plugin_id(self):
"""Test that ValueError is raised when plugin ID missing."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
service = knowledge_module.KnowledgeService(mock_app)
with pytest.raises(ValueError) as exc_info:
await service.create_knowledge_base({'name': 'Test'})
assert 'knowledge_engine_plugin_id is required' in str(exc_info.value)
@pytest.mark.asyncio
async def test_creates_with_default_name(self):
"""Test that KB is created with default name if not provided."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_kb = Mock()
mock_kb.uuid = 'new_kb_uuid'
mock_app.rag_mgr.create_knowledge_base = AsyncMock(return_value=mock_kb)
service = knowledge_module.KnowledgeService(mock_app)
await service.create_knowledge_base({'knowledge_engine_plugin_id': 'author/engine'})
# Check that default name 'Untitled' was used
call_args = mock_app.rag_mgr.create_knowledge_base.call_args
assert call_args.kwargs['name'] == 'Untitled'
class TestUpdateKnowledgeBase:
"""Tests for update_knowledge_base method."""
@pytest.mark.asyncio
async def test_updates_mutable_fields_only(self):
"""Test that only mutable fields are updated."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value={'uuid': 'kb1', 'name': 'Updated'})
mock_app.rag_mgr.remove_knowledge_base_from_runtime = AsyncMock()
mock_app.rag_mgr.load_knowledge_base = AsyncMock()
service = knowledge_module.KnowledgeService(mock_app)
# Pass both mutable and immutable fields
await service.update_knowledge_base(
'kb1',
{
'name': 'New Name',
'description': 'New desc',
'uuid': 'should_be_filtered', # immutable
},
)
# Check that only mutable fields were passed to update
call_args = mock_app.persistence_mgr.execute_async.call_args
assert call_args is not None
@pytest.mark.asyncio
async def test_returns_early_when_no_mutable_fields(self):
"""Test that update returns early when no mutable fields provided."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
service = knowledge_module.KnowledgeService(mock_app)
# Pass only immutable fields
await service.update_knowledge_base('kb1', {'uuid': 'should_be_filtered'})
# No DB update should be called
mock_app.persistence_mgr.execute_async.assert_not_called()
class TestCheckDocCapability:
"""Tests for _check_doc_capability method."""
@pytest.mark.asyncio
async def test_passes_when_capability_supported(self):
"""Test that check passes when doc_ingestion capability exists."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(
return_value={'knowledge_engine': {'capabilities': ['doc_ingestion']}}
)
service = knowledge_module.KnowledgeService(mock_app)
await service._check_doc_capability('kb1', 'document upload')
# No exception raised means success
@pytest.mark.asyncio
async def test_raises_when_kb_not_found(self):
"""Test that Exception is raised when KB not found."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value=None)
service = knowledge_module.KnowledgeService(mock_app)
with pytest.raises(Exception) as exc_info:
await service._check_doc_capability('nonexistent', 'test operation')
assert 'Knowledge base not found' in str(exc_info.value)
@pytest.mark.asyncio
async def test_raises_when_capability_not_supported(self):
"""Test that Exception is raised when doc_ingestion not in capabilities."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(
return_value={'knowledge_engine': {'capabilities': ['other_capability']}}
)
service = knowledge_module.KnowledgeService(mock_app)
with pytest.raises(Exception) as exc_info:
await service._check_doc_capability('kb1', 'document upload')
assert 'does not support document upload' in str(exc_info.value)
class TestListKnowledgeEngines:
"""Tests for list_knowledge_engines method."""
@pytest.mark.asyncio
async def test_returns_engines_from_plugin_connector(self):
"""Test that it returns knowledge engines from plugin connector."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_knowledge_engines = AsyncMock(
return_value=[{'id': 'engine1', 'name': 'Engine 1'}]
)
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_knowledge_engines()
assert len(result) == 1
assert result[0]['id'] == 'engine1'
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test that it returns empty list when plugin disabled."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.is_enable_plugin = False
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_knowledge_engines()
assert result == []
@pytest.mark.asyncio
async def test_returns_empty_on_exception(self):
"""Test that it returns empty list and logs warning on exception."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_knowledge_engines = AsyncMock(side_effect=Exception('Connection error'))
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_knowledge_engines()
assert result == []
mock_app.logger.warning.assert_called_once()
class TestListParsers:
"""Tests for list_parsers method."""
@pytest.mark.asyncio
async def test_returns_all_parsers(self):
"""Test that it returns all parsers when no MIME type filter."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_parsers = AsyncMock(
return_value=[
{'id': 'parser1', 'supported_mime_types': ['text/plain']},
{'id': 'parser2', 'supported_mime_types': ['application/pdf']},
]
)
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_parsers()
assert len(result) == 2
@pytest.mark.asyncio
async def test_filters_by_mime_type(self):
"""Test that it filters parsers by MIME type."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_parsers = AsyncMock(
return_value=[
{'id': 'parser1', 'supported_mime_types': ['text/plain']},
{'id': 'parser2', 'supported_mime_types': ['application/pdf']},
]
)
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_parsers(mime_type='application/pdf')
assert len(result) == 1
assert result[0]['id'] == 'parser2'
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test that it returns empty list when plugin disabled."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.is_enable_plugin = False
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_parsers()
assert result == []
class TestGetEngineSchemas:
"""Tests for get_engine_creation_schema and get_engine_retrieval_schema."""
@pytest.mark.asyncio
async def test_returns_creation_schema(self):
"""Test that it returns creation schema for engine."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.get_rag_creation_schema = AsyncMock(
return_value={'properties': {'name': {'type': 'string'}}}
)
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_engine_creation_schema('author/engine')
assert 'properties' in result
@pytest.mark.asyncio
async def test_returns_retrieval_schema(self):
"""Test that it returns retrieval schema for engine."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.get_rag_retrieval_schema = AsyncMock(
return_value={'properties': {'top_k': {'type': 'integer'}}}
)
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_engine_retrieval_schema('author/engine')
assert 'properties' in result
@pytest.mark.asyncio
async def test_returns_empty_dict_on_exception(self):
"""Test that it returns empty dict and logs warning on exception."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.get_rag_creation_schema = AsyncMock(side_effect=Exception('Plugin error'))
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_engine_creation_schema('author/engine')
assert result == {}
mock_app.logger.warning.assert_called_once()
@@ -0,0 +1,820 @@
"""
Unit tests for MaintenanceService.
Tests storage maintenance and diagnostics including:
- Cleanup expired files
- Storage analysis
- File counting and sizing
- Monitoring counts
- Binary storage stats
Source: src/langbot/pkg/api/http/service/maintenance.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from types import SimpleNamespace
import datetime
from pathlib import Path
from langbot.pkg.api.http.service.maintenance import MaintenanceService
pytestmark = pytest.mark.asyncio
def _create_mock_result(scalar_value=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.scalar = Mock(return_value=scalar_value)
return result
class TestMaintenanceServiceCleanupExpiredFiles:
"""Tests for cleanup_expired_files method."""
async def test_cleanup_expired_files_default_retention(self):
"""Uses default retention days when config not set."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.storage_mgr = SimpleNamespace()
# Create a proper mock object with __class__.__name__
storage_provider = MagicMock()
storage_provider.__class__.__name__ = 'LocalStorageProvider'
ap.storage_mgr.storage_provider = storage_provider
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Mock the internal cleanup methods - one is async, one is not
service._cleanup_expired_uploaded_files = AsyncMock(return_value=0)
service._cleanup_expired_log_files = Mock(return_value=0) # NOT async!
# Execute
result = await service.cleanup_expired_files()
# Verify - returns counts
assert 'uploaded_files' in result
assert 'log_files' in result
assert result['uploaded_files'] == 0
assert result['log_files'] == 0
async def test_cleanup_expired_files_custom_retention(self):
"""Uses custom retention days from config."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'storage': {
'cleanup': {
'uploaded_file_retention_days': 14,
'log_retention_days': 7,
}
}
}
ap.storage_mgr = SimpleNamespace()
storage_provider = MagicMock()
storage_provider.__class__.__name__ = 'LocalStorageProvider'
ap.storage_mgr.storage_provider = storage_provider
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Mock the internal cleanup methods
service._cleanup_expired_uploaded_files = AsyncMock(return_value=2)
service._cleanup_expired_log_files = Mock(return_value=3) # NOT async
# Execute
result = await service.cleanup_expired_files()
# Verify
assert result['uploaded_files'] == 2
assert result['log_files'] == 3
async def test_cleanup_expired_files_s3_provider(self):
"""Handles S3StorageProvider correctly."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.storage_mgr = SimpleNamespace()
# Mock S3 provider
s3_provider = MagicMock()
s3_provider.__class__.__name__ = 'S3StorageProvider'
s3_provider.delete = AsyncMock()
ap.storage_mgr.storage_provider = s3_provider
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Mock the internal cleanup methods
service._cleanup_expired_uploaded_files = AsyncMock(return_value=1)
service._cleanup_expired_log_files = Mock(return_value=0) # NOT async
# Execute
result = await service.cleanup_expired_files()
# Verify
assert result['uploaded_files'] == 1
assert result['log_files'] == 0
async def test_cleanup_expired_files_invalid_retention(self):
"""Uses default for invalid retention config."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'storage': {
'cleanup': {
'uploaded_file_retention_days': 'invalid', # Invalid
'log_retention_days': 0, # Invalid (less than 1)
}
}
}
ap.storage_mgr = SimpleNamespace()
storage_provider = MagicMock()
storage_provider.__class__.__name__ = 'LocalStorageProvider'
ap.storage_mgr.storage_provider = storage_provider
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Mock the internal cleanup methods
service._cleanup_expired_uploaded_files = AsyncMock(return_value=0)
service._cleanup_expired_log_files = Mock(return_value=0) # NOT async
# Execute
result = await service.cleanup_expired_files()
# Verify - warning logged, defaults used
assert ap.logger.warning.called
assert 'uploaded_files' in result
class TestMaintenanceServiceGetStorageAnalysis:
"""Tests for get_storage_analysis method."""
async def test_get_storage_analysis_basic(self):
"""Returns basic storage analysis."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'database': {'use': 'sqlite', 'sqlite': {'path': 'data/langbot.db'}}}
ap.persistence_mgr = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
ap.task_mgr = SimpleNamespace()
ap.task_mgr.get_stats = Mock(return_value={'running': 0})
# Mock monitoring counts
count_result = _create_mock_result(scalar_value=10)
ap.persistence_mgr.execute_async = AsyncMock(return_value=count_result)
service = MaintenanceService(ap)
# Mock file operations
service._path_size = Mock(return_value=1000)
service._file_count = Mock(return_value=5)
service._monitoring_counts = AsyncMock(return_value={'messages': 10, 'errors': 0})
service._binary_storage_stats = AsyncMock(return_value={'count': 5, 'size_bytes': 500})
service._expired_uploaded_candidates = AsyncMock(return_value=[])
service._expired_log_candidates = Mock(return_value=[])
# Execute
result = await service.get_storage_analysis()
# Verify
assert 'generated_at' in result
assert 'cleanup_policy' in result
assert 'sections' in result
assert 'database' in result
assert 'cleanup_candidates' in result
async def test_get_storage_analysis_sections(self):
"""Returns all storage sections."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'database': {'use': 'postgresql'}}
ap.persistence_mgr = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
ap.task_mgr = None
count_result = _create_mock_result(scalar_value=0)
ap.persistence_mgr.execute_async = AsyncMock(return_value=count_result)
service = MaintenanceService(ap)
service._path_size = Mock(return_value=0)
service._file_count = Mock(return_value=0)
service._monitoring_counts = AsyncMock(return_value={})
service._binary_storage_stats = AsyncMock(return_value={'count': 0, 'size_bytes': 0})
service._expired_uploaded_candidates = AsyncMock(return_value=[])
service._expired_log_candidates = Mock(return_value=[])
# Execute
result = await service.get_storage_analysis()
# Verify - all sections present
sections = {s['key'] for s in result['sections']}
assert 'database' in sections
assert 'logs' in sections
assert 'storage' in sections
assert 'vector_store' in sections
assert 'plugins' in sections
assert 'mcp' in sections
assert 'temp' in sections
async def test_get_storage_analysis_postgresql(self):
"""Handles PostgreSQL database type."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'database': {'use': 'postgresql'}}
ap.persistence_mgr = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
ap.task_mgr = None
count_result = _create_mock_result(scalar_value=0)
ap.persistence_mgr.execute_async = AsyncMock(return_value=count_result)
service = MaintenanceService(ap)
service._path_size = Mock(return_value=0)
service._file_count = Mock(return_value=0)
service._monitoring_counts = AsyncMock(return_value={})
service._binary_storage_stats = AsyncMock(return_value={'count': 0, 'size_bytes': None})
service._expired_uploaded_candidates = AsyncMock(return_value=[])
service._expired_log_candidates = Mock(return_value=[])
# Execute
result = await service.get_storage_analysis()
# Verify
assert result['database']['type'] == 'postgresql'
async def test_get_storage_analysis_with_cleanup_candidates(self):
"""Returns cleanup candidates in analysis."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
ap.task_mgr = None
count_result = _create_mock_result(scalar_value=0)
ap.persistence_mgr.execute_async = AsyncMock(return_value=count_result)
service = MaintenanceService(ap)
service._path_size = Mock(return_value=0)
service._file_count = Mock(return_value=0)
service._monitoring_counts = AsyncMock(return_value={})
service._binary_storage_stats = AsyncMock(return_value={'count': 0, 'size_bytes': 0})
service._expired_uploaded_candidates = AsyncMock(return_value=[{'key': 'old_file', 'size_bytes': 100}])
service._expired_log_candidates = Mock(return_value=[{'name': 'old_log', 'size_bytes': 50}])
# Execute
result = await service.get_storage_analysis()
# Verify
assert len(result['cleanup_candidates']['uploaded_files']) == 1
assert len(result['cleanup_candidates']['log_files']) == 1
class TestMaintenanceServiceMonitoringCounts:
"""Tests for _monitoring_counts method."""
async def test_monitoring_counts_returns_counts(self):
"""Returns counts for all monitoring tables."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
count_result = _create_mock_result(scalar_value=42)
ap.persistence_mgr.execute_async = AsyncMock(return_value=count_result)
service = MaintenanceService(ap)
# Execute
result = await service._monitoring_counts()
# Verify - all table keys present
assert 'messages' in result
assert 'llm_calls' in result
assert 'embedding_calls' in result
assert 'errors' in result
assert 'sessions' in result
assert 'feedback' in result
async def test_monitoring_counts_zero_results(self):
"""Returns zero counts when tables empty."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
count_result = _create_mock_result(scalar_value=0)
ap.persistence_mgr.execute_async = AsyncMock(return_value=count_result)
service = MaintenanceService(ap)
# Execute
result = await service._monitoring_counts()
# Verify - all zero
assert all(v == 0 for v in result.values())
class TestMaintenanceServiceBinaryStorageStats:
"""Tests for _binary_storage_stats method."""
async def test_binary_storage_stats_returns_stats(self):
"""Returns count and size for binary storage."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
# Mock count result
count_result = _create_mock_result(scalar_value=10)
# Mock size result
size_result = _create_mock_result(scalar_value=5000)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return count_result
return size_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MaintenanceService(ap)
# Execute
result = await service._binary_storage_stats()
# Verify
assert result['count'] == 10
assert result['size_bytes'] == 5000
async def test_binary_storage_stats_size_error(self):
"""Handles error when calculating size."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
count_result = _create_mock_result(scalar_value=5)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return count_result
raise Exception('Size calculation error')
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MaintenanceService(ap)
# Execute
result = await service._binary_storage_stats()
# Verify - warning logged, size_bytes None or 0
assert ap.logger.warning.called
assert result['count'] == 5
class TestMaintenanceServicePathSize:
"""Tests for _path_size method."""
def test_path_size_nonexistent_path(self):
"""Returns 0 for nonexistent path."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
# Execute
result = service._path_size(Path('/nonexistent/path'))
# Verify
assert result == 0
def test_path_size_single_file(self):
"""Returns size for single file."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
# Mock file
mock_stat = Mock()
mock_stat.st_size = 100
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'is_file', return_value=True):
with patch.object(Path, 'stat', return_value=mock_stat):
result = service._path_size(Path('test.txt'))
# Verify
assert result == 100
def test_path_size_directory(self):
"""Returns total size for directory."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
# Mock os.walk
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'is_file', return_value=False):
with patch('os.walk') as mock_walk:
mock_walk.return_value = [
('/test_dir', [], ['file1.txt', 'file2.txt']),
]
# Mock file stat
mock_stat = Mock()
mock_stat.st_size = 50
with patch.object(Path, 'stat', return_value=mock_stat):
result = service._path_size(Path('/test_dir'))
# Verify - 2 files * 50 bytes
assert result == 100
class TestMaintenanceServiceFileCount:
"""Tests for _file_count method."""
def test_file_count_nonexistent_path(self):
"""Returns 0 for nonexistent path."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
# Execute
result = service._file_count(Path('/nonexistent/path'))
# Verify
assert result == 0
def test_file_count_single_file(self):
"""Returns 1 for single file."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'is_file', return_value=True):
result = service._file_count(Path('test.txt'))
# Verify
assert result == 1
def test_file_count_directory(self):
"""Returns file count for directory."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'is_file', return_value=False):
with patch('os.walk') as mock_walk:
mock_walk.return_value = [
('/test_dir', [], ['file1.txt', 'file2.txt', 'file3.txt']),
]
result = service._file_count(Path('/test_dir'))
# Verify
assert result == 3
class TestMaintenanceServicePositiveInt:
"""Tests for _positive_int helper method."""
def test_positive_int_valid_value(self):
"""Returns valid positive integer."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Execute
result = service._positive_int(7, 5, 'test_param')
# Verify
assert result == 7
assert not ap.logger.warning.called
def test_positive_int_invalid_string(self):
"""Returns default for invalid string."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Execute
result = service._positive_int('invalid', 5, 'test_param')
# Verify
assert result == 5
assert ap.logger.warning.called
def test_positive_int_invalid_none(self):
"""Returns default for None."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Execute
result = service._positive_int(None, 5, 'test_param')
# Verify
assert result == 5
assert ap.logger.warning.called
def test_positive_int_negative_value(self):
"""Returns default for negative value."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Execute
result = service._positive_int(-1, 5, 'test_param')
# Verify
assert result == 5
assert ap.logger.warning.called
def test_positive_int_zero_value(self):
"""Returns default for zero value."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.logger.warning = Mock()
service = MaintenanceService(ap)
# Execute
result = service._positive_int(0, 5, 'test_param')
# Verify
assert result == 5
assert ap.logger.warning.called
class TestMaintenanceServiceIsUploadedFileKey:
"""Tests for _is_uploaded_file_key helper method."""
def test_is_uploaded_file_key_valid(self):
"""Returns True for valid upload file key."""
# Setup
ap = SimpleNamespace()
service = MaintenanceService(ap)
# Execute - simple filename without path
result = service._is_uploaded_file_key('uploaded_file.txt')
# Verify
assert result is True
def test_is_uploaded_file_key_with_path(self):
"""Returns False for key with path separator."""
# Setup
ap = SimpleNamespace()
service = MaintenanceService(ap)
# Execute - key with path
result = service._is_uploaded_file_key('path/to/file.txt')
# Verify
assert result is False
def test_is_uploaded_file_key_plugin_config(self):
"""Returns False for plugin config prefix."""
# Setup
ap = SimpleNamespace()
service = MaintenanceService(ap)
# Execute - plugin config file
result = service._is_uploaded_file_key('plugin_config_some_plugin.json')
# Verify
assert result is False
class TestMaintenanceServiceExpiredLogCandidates:
"""Tests for _expired_log_candidates method."""
def test_expired_log_candidates_nonexistent_dir(self):
"""Returns empty list when logs dir not exists."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
with patch.object(Path, 'exists', return_value=False):
result = service._expired_log_candidates(3)
# Verify
assert result == []
def test_expired_log_candidates_matches_pattern(self):
"""Matches log file pattern correctly."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
# Mock directory with log files
old_date = datetime.date.today() - datetime.timedelta(days=10)
old_log_name = f'langbot-{old_date.isoformat()}.log'
recent_log_name = f'langbot-{datetime.date.today().isoformat()}.log'
mock_entry_old = Mock(spec=Path)
mock_entry_old.is_file = Mock(return_value=True)
mock_entry_old.name = old_log_name
mock_stat = Mock()
mock_stat.st_size = 1000
mock_entry_old.stat = Mock(return_value=mock_stat)
mock_entry_recent = Mock(spec=Path)
mock_entry_recent.is_file = Mock(return_value=True)
mock_entry_recent.name = recent_log_name
mock_stat2 = Mock()
mock_stat2.st_size = 500
mock_entry_recent.stat = Mock(return_value=mock_stat2)
# Non-log file
mock_entry_other = Mock(spec=Path)
mock_entry_other.is_file = Mock(return_value=True)
mock_entry_other.name = 'other_file.txt'
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'iterdir') as mock_iterdir:
mock_iterdir.return_value = [mock_entry_old, mock_entry_recent, mock_entry_other]
result = service._expired_log_candidates(3)
# Verify - only old log included
assert len(result) == 1
assert result[0]['name'] == old_log_name
def test_expired_log_candidates_includes_path(self):
"""Includes path when include_paths=True."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
old_date = datetime.date.today() - datetime.timedelta(days=10)
old_log_name = f'langbot-{old_date.isoformat()}.log'
mock_entry = Mock(spec=Path)
mock_entry.is_file = Mock(return_value=True)
mock_entry.name = old_log_name
mock_entry.__str__ = Mock(return_value='/data/logs/' + old_log_name)
mock_stat = Mock()
mock_stat.st_size = 1000
mock_entry.stat = Mock(return_value=mock_stat)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'iterdir') as mock_iterdir:
mock_iterdir.return_value = [mock_entry]
result = service._expired_log_candidates(3, include_paths=True)
# Verify - path included
assert 'path' in result[0]
class TestMaintenanceServiceExpiredLocalUploadCandidates:
"""Tests for _expired_local_upload_candidates method."""
def test_expired_local_upload_candidates_nonexistent_dir(self):
"""Returns empty list when storage dir not exists."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
with patch.object(Path, 'exists', return_value=False):
result = service._expired_local_upload_candidates(7)
# Verify
assert result == []
def test_expired_local_upload_candidates_filters_uploaded(self):
"""Only returns uploaded files matching pattern."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
# Mock _is_uploaded_file_key
service._is_uploaded_file_key = Mock(side_effect=lambda key: 'plugin_config_' not in key and '/' not in key)
# Create mock files - one valid, one plugin config
mock_entry_valid = Mock(spec=Path)
mock_entry_valid.is_file = Mock(return_value=True)
mock_entry_valid.name = 'valid_upload.txt'
mock_stat = Mock()
mock_stat.st_size = 100
mock_stat.st_mtime = 0 # Very old
mock_entry_valid.stat = Mock(return_value=mock_stat)
mock_entry_plugin = Mock(spec=Path)
mock_entry_plugin.is_file = Mock(return_value=True)
mock_entry_plugin.name = 'plugin_config_test.json'
mock_stat2 = Mock()
mock_stat2.st_size = 200
mock_stat2.st_mtime = 0
mock_entry_plugin.stat = Mock(return_value=mock_stat2)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'iterdir') as mock_iterdir:
mock_iterdir.return_value = [mock_entry_valid, mock_entry_plugin]
result = service._expired_local_upload_candidates(7)
# Verify - only valid upload included
assert len(result) == 1
assert result[0]['key'] == 'valid_upload.txt'
def test_expired_local_upload_candidates_includes_path(self):
"""Includes path when include_paths=True."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
service = MaintenanceService(ap)
service._is_uploaded_file_key = Mock(return_value=True)
mock_entry = Mock(spec=Path)
mock_entry.is_file = Mock(return_value=True)
mock_entry.name = 'old_file.txt'
mock_entry.__str__ = Mock(return_value='/data/storage/old_file.txt')
mock_stat = Mock()
mock_stat.st_size = 100
mock_stat.st_mtime = 0
mock_entry.stat = Mock(return_value=mock_stat)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'iterdir') as mock_iterdir:
mock_iterdir.return_value = [mock_entry]
result = service._expired_local_upload_candidates(7, include_paths=True)
# Verify - path included
assert 'path' in result[0]
@@ -0,0 +1,715 @@
"""
Unit tests for MCPService.
Tests MCP server CRUD operations including:
- MCP server listing with runtime info
- MCP server creation with limitations
- MCP server update with enable/disable
- MCP server deletion
- MCP server connection testing
Source: src/langbot/pkg/api/http/service/mcp.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, MagicMock
from types import SimpleNamespace
import uuid
from langbot.pkg.api.http.service.mcp import MCPService
from langbot.pkg.entity.persistence.mcp import MCPServer
pytestmark = pytest.mark.asyncio
def _create_mock_mcp_server(
server_uuid: str = None,
name: str = 'Test MCP Server',
enable: bool = True,
mode: str = 'stdio',
extra_args: dict = None,
) -> Mock:
"""Helper to create mock MCPServer entity."""
server = Mock(spec=MCPServer)
server.uuid = server_uuid or str(uuid.uuid4())
server.name = name
server.enable = enable
server.mode = mode
server.extra_args = extra_args or {}
return server
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
class TestMCPServiceGetRuntimeInfo:
"""Tests for get_runtime_info method."""
async def test_get_runtime_info_session_exists(self):
"""Returns runtime info when session exists."""
# Setup
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
mock_session = SimpleNamespace()
mock_session.get_runtime_info_dict = Mock(return_value={'status': 'running', 'tools': 5})
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=mock_session)
service = MCPService(ap)
# Execute
result = await service.get_runtime_info('test-server')
# Verify
assert result is not None
assert result['status'] == 'running'
async def test_get_runtime_info_session_not_exists(self):
"""Returns None when session not exists."""
# Setup
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)
service = MCPService(ap)
# Execute
result = await service.get_runtime_info('nonexistent-server')
# Verify
assert result is None
class TestMCPServiceResources:
"""Tests for MCP resource helpers."""
async def test_get_resource_templates_delegates_to_loader(self):
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_resource_templates = AsyncMock(
return_value=[{'uri_template': 'file:///{path}', 'name': 'files'}]
)
service = MCPService(ap)
result = await service.get_mcp_server_resource_templates('docs')
assert result == [{'uri_template': 'file:///{path}', 'name': 'files'}]
ap.tool_mgr.mcp_tool_loader.get_resource_templates.assert_awaited_once_with('docs')
async def test_read_resource_envelope_uses_ui_preview_source(self):
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.read_resource_envelope = AsyncMock(
return_value={
'server_name': 'docs',
'uri': 'file:///README.md',
'contents': [],
'source': 'ui_preview',
}
)
service = MCPService(ap)
result = await service.read_mcp_server_resource_envelope(
'docs',
'file:///README.md',
max_bytes=4096,
include_blob=True,
)
assert result['source'] == 'ui_preview'
ap.tool_mgr.mcp_tool_loader.read_resource_envelope.assert_awaited_once_with(
'docs',
'file:///README.md',
include_blob=True,
source='ui_preview',
max_bytes=4096,
)
class TestMCPServiceGetMCPServers:
"""Tests for get_mcp_servers method."""
async def test_get_mcp_servers_empty_list(self):
"""Returns empty list when no MCP servers exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
}
)
ap.tool_mgr = None
service = MCPService(ap)
# Execute
result = await service.get_mcp_servers()
# Verify
assert result == []
async def test_get_mcp_servers_returns_serialized_list(self):
"""Returns serialized list of MCP servers."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
server1 = _create_mock_mcp_server(server_uuid='uuid-1', name='Server 1')
server2 = _create_mock_mcp_server(server_uuid='uuid-2', name='Server 2')
mock_result = _create_mock_result([server1, server2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
'enable': entity.enable,
'mode': entity.mode,
}
)
ap.tool_mgr = None
service = MCPService(ap)
# Execute
result = await service.get_mcp_servers()
# Verify
assert len(result) == 2
assert result[0]['name'] == 'Server 1'
assert result[1]['name'] == 'Server 2'
async def test_get_mcp_servers_with_runtime_info(self):
"""Returns MCP servers with runtime info when requested."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
server1 = _create_mock_mcp_server(server_uuid='uuid-1', name='Server 1')
mock_result = _create_mock_result([server1])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
}
)
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)
service = MCPService(ap)
service.get_runtime_info = AsyncMock(return_value={'status': 'connected'})
# Execute
result = await service.get_mcp_servers(contain_runtime_info=True)
# Verify - runtime info included
assert result[0]['runtime_info'] == {'status': 'connected'}
class TestMCPServiceCreateMCPServer:
"""Tests for create_mcp_server method."""
async def test_create_mcp_server_max_extensions_reached_raises(self):
"""Raises ValueError when max_extensions limit reached."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_extensions': 2}}}
ap.plugin_connector = SimpleNamespace()
ap.plugin_connector.list_plugins = AsyncMock(return_value=[Mock(), Mock()]) # 2 plugins
# Mock get_mcp_servers to return 0 servers (2 plugins already)
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={})
ap.tool_mgr = None
service = MCPService(ap)
# Execute & Verify - 2 plugins + new server would exceed limit
with pytest.raises(ValueError, match='Maximum number of extensions'):
await service.create_mcp_server({'name': 'New Server'})
async def test_create_mcp_server_no_limit(self):
"""Creates MCP server without limit when max_extensions=-1."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'system': {
'limitation': {
'max_extensions': -1 # No limit
}
}
}
ap.tool_mgr = None
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid'})
service = MCPService(ap)
# Execute
server_uuid = await service.create_mcp_server({'name': 'New Server'})
# Verify
assert server_uuid is not None
assert len(server_uuid) == 36 # UUID format
async def test_create_mcp_server_duplicate_name_raises(self):
"""Rejects duplicate MCP server names."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_extensions': -1}}}
ap.tool_mgr = None
existing_server = _create_mock_mcp_server(name='Existing Server')
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=existing_server))
ap.persistence_mgr.serialize_model = Mock(return_value={})
service = MCPService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='MCP server already exists: Existing Server'):
await service.create_mcp_server({'name': 'Existing Server'})
async def test_create_mcp_server_loads_server(self):
"""Loads server into tool_mgr when enabled."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_extensions': -1}}}
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
# Create mock server entity
server_entity = _create_mock_mcp_server(server_uuid='new-uuid', enable=True)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result([]) # Empty result for duplicate-name check
elif call_count == 2:
return Mock() # Insert
return _create_mock_result(first_item=server_entity) # Select created
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
)
service = MCPService(ap)
# Execute
await service.create_mcp_server({'name': 'New Server', 'enable': True})
# Verify - host_mcp_server was called
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
async def test_create_mcp_server_disabled_no_load(self):
"""Does not load server when disabled."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_extensions': -1}}}
ap.tool_mgr = None
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid'})
service = MCPService(ap)
# Execute with enable=False
server_uuid = await service.create_mcp_server({'name': 'New Server', 'enable': False})
# Verify - no tool_mgr load attempt
assert server_uuid is not None
class TestMCPServiceGetMCPServerByName:
"""Tests for get_mcp_server_by_name method."""
async def test_get_mcp_server_by_name_found(self):
"""Returns MCP server when found by name."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
server = _create_mock_mcp_server(name='Found Server')
mock_result = _create_mock_result(first_item=server)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Found Server',
'runtime_info': None,
}
)
ap.tool_mgr = None
service = MCPService(ap)
service.get_runtime_info = AsyncMock(return_value=None)
# Execute
result = await service.get_mcp_server_by_name('Found Server')
# Verify
assert result is not None
assert result['name'] == 'Found Server'
async def test_get_mcp_server_by_name_not_found(self):
"""Returns None when MCP server not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = MCPService(ap)
# Execute
result = await service.get_mcp_server_by_name('Nonexistent Server')
# Verify
assert result is None
class TestMCPServiceUpdateMCPServer:
"""Tests for update_mcp_server method."""
async def test_update_mcp_server_disable_enabled_server(self):
"""Removes server when disabling previously enabled server."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {'Old Server': Mock()}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
old_server = _create_mock_mcp_server(name='Old Server', enable=True)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=old_server)
return Mock() # Update
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MCPService(ap)
# Execute - disable server
await service.update_mcp_server('test-uuid', {'enable': False})
# Verify - server was removed
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once()
async def test_update_mcp_server_enable_disabled_server(self):
"""Loads server when enabling previously disabled server."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {}
ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
old_server = _create_mock_mcp_server(name='Old Server', enable=False)
updated_server = _create_mock_mcp_server(name='Old Server', enable=True)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=old_server)
elif call_count == 2:
return Mock() # Update
return _create_mock_result(first_item=updated_server) # Select updated
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'test-uuid', 'name': 'Old Server', 'enable': True}
)
service = MCPService(ap)
# Execute - enable server
await service.update_mcp_server('test-uuid', {'enable': True})
# Verify - server was loaded
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
async def test_update_mcp_server_update_enabled_server(self):
"""Removes and reloads server when updating enabled server."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {'Old Server': Mock()}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
old_server = _create_mock_mcp_server(name='Old Server', enable=True)
# Mock for: first select -> update -> second select (for updated server)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
# All selects return the server
return _create_mock_result(first_item=old_server)
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'test-uuid', 'name': 'Old Server', 'enable': True}
)
service = MCPService(ap)
# Execute - update enabled server (keep enabled, update extra_args)
await service.update_mcp_server('test-uuid', {'enable': True, 'extra_args': {'new': 'args'}})
# Verify - remove and reload
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with('Old Server')
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
async def test_update_mcp_server_no_tool_mgr(self):
"""Updates persistence without tool_mgr operations."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Set mcp_tool_loader to None, not tool_mgr itself
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = None
old_server = _create_mock_mcp_server(name='Server', enable=True)
# Mock execute for select and update
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=old_server)
return Mock() # Update
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MCPService(ap)
# Execute - should not raise
await service.update_mcp_server('test-uuid', {'name': 'New Name'})
# Verify - persistence was called
assert ap.persistence_mgr.execute_async.call_count >= 2
class TestMCPServiceDeleteMCPServer:
"""Tests for delete_mcp_server method."""
async def test_delete_mcp_server_calls_remove_and_delete(self):
"""Calls both persistence delete and tool_mgr remove."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {'Server to Delete': Mock()}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
server = _create_mock_mcp_server(name='Server to Delete')
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=server)
return Mock() # Delete
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MCPService(ap)
# Execute
await service.delete_mcp_server('test-uuid')
# Verify
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with('Server to Delete')
ap.persistence_mgr.execute_async.assert_called()
async def test_delete_mcp_server_not_in_sessions(self):
"""Does not attempt remove if server not in sessions."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {} # Server not in sessions
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
server = _create_mock_mcp_server(name='Not in Sessions')
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=server)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MCPService(ap)
# Execute
await service.delete_mcp_server('test-uuid')
# Verify - remove not called (server not in sessions)
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_not_called()
async def test_delete_mcp_server_nonexistent_uuid(self):
"""Delete operation completes even for nonexistent UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
# No server found
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=None)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MCPService(ap)
# Execute - should not raise
await service.delete_mcp_server('nonexistent-uuid')
# Verify - delete was called regardless
ap.persistence_mgr.execute_async.assert_called()
class TestMCPServiceTestMCPServer:
"""Tests for test_mcp_server method."""
async def test_test_mcp_server_existing_server(self):
"""Tests existing MCP server connection."""
# Setup
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
from langbot.pkg.provider.tools.loaders.mcp import MCPSessionStatus
mock_session = MagicMock()
mock_session.status = MCPSessionStatus.ERROR
mock_session.start = AsyncMock()
mock_session.refresh = AsyncMock()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=mock_session)
ap.task_mgr = SimpleNamespace()
ap.task_mgr.create_user_task = Mock(return_value=SimpleNamespace(id=123))
service = MCPService(ap)
# Execute
task_id = await service.test_mcp_server('existing-server', {})
# Verify - returns task ID
assert task_id == 123
async def test_test_mcp_server_not_found_raises(self):
"""Raises ValueError when server not found."""
# Setup
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)
service = MCPService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='Server not found'):
await service.test_mcp_server('nonexistent-server', {})
async def test_test_mcp_server_new_server(self):
"""Tests new MCP server with underscore name."""
# Setup
ap = SimpleNamespace()
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
mock_session = MagicMock()
mock_session.start = AsyncMock()
ap.tool_mgr.mcp_tool_loader.load_mcp_server = AsyncMock(return_value=mock_session)
ap.task_mgr = SimpleNamespace()
ap.task_mgr.create_user_task = Mock(return_value=SimpleNamespace(id=456))
service = MCPService(ap)
# Execute with '_' name (new server)
task_id = await service.test_mcp_server('_', {'name': 'New Server'})
# Verify - load_mcp_server called
ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_called_once()
assert task_id == 456
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,871 @@
"""
Unit tests for PipelineService.
Tests pipeline CRUD operations including:
- Pipeline listing with sorting
- Pipeline creation with default config
- Pipeline update with bot sync
- Pipeline copy functionality
- Extensions preferences management
Source: src/langbot/pkg/api/http/service/pipeline.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, patch, mock_open
from types import SimpleNamespace
import uuid
import json
from langbot.pkg.api.http.service.pipeline import PipelineService, default_stage_order
from langbot.pkg.entity.persistence.pipeline import LegacyPipeline
pytestmark = pytest.mark.asyncio
def _create_mock_pipeline(
pipeline_uuid: str = None,
name: str = 'Test Pipeline',
description: str = 'Test Description',
is_default: bool = False,
stages: list = None,
config: dict = None,
extensions_preferences: dict = None,
) -> Mock:
"""Helper to create mock LegacyPipeline entity."""
pipeline = Mock(spec=LegacyPipeline)
pipeline.uuid = pipeline_uuid or str(uuid.uuid4())
pipeline.name = name
pipeline.description = description
pipeline.emoji = '⚙️'
pipeline.is_default = is_default
pipeline.for_version = '1.0.0'
pipeline.stages = stages or default_stage_order.copy()
pipeline.config = config or {}
pipeline.extensions_preferences = extensions_preferences or {
'enable_all_plugins': True,
'enable_all_mcp_servers': True,
'plugins': [],
'mcp_servers': [],
}
return pipeline
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
class TestPipelineServiceGetPipelineMetadata:
"""Tests for get_pipeline_metadata method."""
async def test_get_pipeline_metadata_returns_list(self):
"""Returns list of pipeline metadata configs."""
# Setup
ap = SimpleNamespace()
ap.pipeline_config_meta_trigger = {'trigger': {}}
ap.pipeline_config_meta_safety = {'safety': {}}
ap.pipeline_config_meta_ai = {'ai': {}}
ap.pipeline_config_meta_output = {'output': {}}
service = PipelineService(ap)
# Execute
result = await service.get_pipeline_metadata()
# Verify
assert len(result) == 4
assert 'trigger' in result[0]
assert 'safety' in result[1]
assert 'ai' in result[2]
assert 'output' in result[3]
class TestPipelineServiceGetPipelines:
"""Tests for get_pipelines method."""
async def test_get_pipelines_empty_list(self):
"""Returns empty list when no pipelines exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
}
)
service = PipelineService(ap)
# Execute
result = await service.get_pipelines()
# Verify
assert result == []
async def test_get_pipelines_returns_sorted_by_created_at_desc(self):
"""Returns pipelines sorted by created_at descending by default."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
pipeline1 = _create_mock_pipeline(pipeline_uuid='uuid-1', name='Pipeline 1')
pipeline2 = _create_mock_pipeline(pipeline_uuid='uuid-2', name='Pipeline 2')
mock_result = _create_mock_result([pipeline1, pipeline2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
}
)
service = PipelineService(ap)
# Execute
result = await service.get_pipelines()
# Verify
assert len(result) == 2
ap.persistence_mgr.execute_async.assert_called_once()
async def test_get_pipelines_sort_by_updated_at_asc(self):
"""Returns pipelines sorted by updated_at ascending."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={})
service = PipelineService(ap)
# Execute
await service.get_pipelines(sort_by='updated_at', sort_order='ASC')
# Verify - execute was called with sort parameters
ap.persistence_mgr.execute_async.assert_called_once()
class TestPipelineServiceGetPipeline:
"""Tests for get_pipeline method."""
async def test_get_pipeline_by_uuid_found(self):
"""Returns pipeline when found by UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
pipeline = _create_mock_pipeline(pipeline_uuid='test-uuid', name='Found Pipeline')
mock_result = _create_mock_result(first_item=pipeline)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Found Pipeline',
'stages': default_stage_order,
}
)
service = PipelineService(ap)
# Execute
result = await service.get_pipeline('test-uuid')
# Verify
assert result is not None
assert result['uuid'] == 'test-uuid'
assert result['name'] == 'Found Pipeline'
async def test_get_pipeline_by_uuid_not_found(self):
"""Returns None when pipeline not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = PipelineService(ap)
# Execute
result = await service.get_pipeline('nonexistent-uuid')
# Verify
assert result is None
class TestPipelineServiceCreatePipeline:
"""Tests for create_pipeline method."""
async def test_create_pipeline_max_limit_reached_raises(self):
"""Raises ValueError when max_pipelines limit reached."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': 2}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
mock_result = _create_mock_result([_create_mock_pipeline(), _create_mock_pipeline()])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Pipeline 1'})
service = PipelineService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='Maximum number of pipelines'):
await service.create_pipeline({'name': 'New Pipeline'})
async def test_create_pipeline_no_limit(self):
"""Creates pipeline without limit when max_pipelines=-1."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': -1}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
service = PipelineService(ap)
# Override get_pipelines to return empty list (no limit check issue)
service.get_pipelines = AsyncMock(return_value=[])
service.get_pipeline = AsyncMock(return_value={'uuid': 'new-uuid', 'name': 'New Pipeline'})
# Mock persistence for insert
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid', 'name': 'New Pipeline'})
# Mock the file read for default config - patch at the utils module level
default_config = {'trigger': {}, 'safety': {}, 'ai': {}, 'output': {}}
with patch('builtins.open', mock_open(read_data=json.dumps(default_config))):
with patch(
'langbot.pkg.utils.paths.get_resource_path', return_value='templates/default-pipeline-config.json'
):
bot_uuid = await service.create_pipeline({'name': 'New Pipeline'})
# Verify
assert bot_uuid is not None
assert len(bot_uuid) == 36 # UUID format
async def test_create_pipeline_as_default(self):
"""Creates pipeline with is_default=True."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': -1}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
service = PipelineService(ap)
service.get_pipelines = AsyncMock(return_value=[])
service.get_pipeline = AsyncMock(
return_value={'uuid': 'new-uuid', 'name': 'Default Pipeline', 'is_default': True}
)
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'new-uuid', 'name': 'Default Pipeline', 'is_default': True}
)
# Mock the file read
default_config = {}
with patch('builtins.open', mock_open(read_data=json.dumps(default_config))):
with patch(
'langbot.pkg.utils.paths.get_resource_path', return_value='templates/default-pipeline-config.json'
):
await service.create_pipeline({'name': 'Default Pipeline'}, default=True)
# Verify - execute was called
ap.persistence_mgr.execute_async.assert_called()
async def test_create_pipeline_sets_default_extensions_preferences(self):
"""Sets default extensions_preferences when not provided."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': -1}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
service = PipelineService(ap)
service.get_pipelines = AsyncMock(return_value=[])
service.get_pipeline = AsyncMock(
return_value={
'uuid': 'new-uuid',
'extensions_preferences': {},
}
)
insert_params = []
async def mock_execute(query):
params = query.compile().params
if 'extensions_preferences' in params:
insert_params.append(params)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'new-uuid',
'extensions_preferences': {},
}
)
default_config = {}
with patch('builtins.open', mock_open(read_data=json.dumps(default_config))):
with patch(
'langbot.pkg.utils.paths.get_resource_path', return_value='templates/default-pipeline-config.json'
):
await service.create_pipeline({'name': 'New Pipeline'})
assert len(insert_params) == 1
assert insert_params[0]['extensions_preferences'] == {
'enable_all_plugins': True,
'enable_all_mcp_servers': True,
'plugins': [],
'mcp_servers': [],
'mcp_resources': [],
'mcp_resource_agent_read_enabled': True,
}
class _MockResultWithBots:
"""Helper class to mock SQLAlchemy result with iterable .all() method."""
def __init__(self, bots_list):
self._bots_list = bots_list
def all(self):
return self._bots_list
def first(self):
return self._bots_list[0] if self._bots_list else None
class TestPipelineServiceUpdatePipeline:
"""Tests for update_pipeline method."""
async def test_update_pipeline_removes_protected_fields(self):
"""Does not persist protected fields from update data."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.sess_mgr = SimpleNamespace()
ap.sess_mgr.session_list = []
ap.bot_service = None # No bot_service when not updating name
ap.persistence_mgr.execute_async = AsyncMock()
service = PipelineService(ap)
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid', 'name': 'Updated'})
# Execute with protected fields - no name change, so no bot sync
pipeline_data = {
'uuid': 'should-be-removed',
'for_version': 'should-be-removed',
'stages': ['should-be-removed'],
'is_default': True,
'description': 'New description', # Not name change, so no bot_service needed
}
await service.update_pipeline('test-uuid', pipeline_data)
update_params = ap.persistence_mgr.execute_async.await_args_list[0].args[0].compile().params
assert update_params['description'] == 'New description'
assert 'should-be-removed' not in update_params.values()
assert ['should-be-removed'] not in update_params.values()
assert not any(value is True for value in update_params.values())
async def test_update_pipeline_syncs_bot_names(self):
"""Updates bot use_pipeline_name when pipeline name changes."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.sess_mgr = SimpleNamespace()
ap.sess_mgr.session_list = []
ap.bot_service = SimpleNamespace()
ap.bot_service.update_bot = AsyncMock()
# Create proper mock Bot entities with uuid attribute
mock_bot1 = Mock()
mock_bot1.uuid = 'bot-uuid-1'
mock_bot2 = Mock()
mock_bot2.uuid = 'bot-uuid-2'
# Create bot list
bot_list = [mock_bot1, mock_bot2]
# Create mock result using helper class
bot_result = _MockResultWithBots(bot_list)
# The order of calls in update_pipeline:
# 1. UPDATE (line 125) - returns Mock (no result needed)
# 2. SELECT bots (line 136) - returns bot_result with .all()
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
# First call is the UPDATE - just return a Mock
return Mock()
elif call_count == 2:
# Second call is the SELECT bots - return proper result
return bot_result
return Mock() # Any additional calls
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(return_value={})
service = PipelineService(ap)
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid', 'name': 'New Name'})
# Execute with name change
await service.update_pipeline('test-uuid', {'name': 'New Name'})
# Verify - bot_service.update_bot was called for each bot
assert ap.bot_service.update_bot.call_count == 2
async def test_update_pipeline_clears_conversations(self):
"""Clears session conversations using this pipeline."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.sess_mgr = SimpleNamespace()
# Mock session with conversation using this pipeline
session = SimpleNamespace()
session.using_conversation = SimpleNamespace()
session.using_conversation.pipeline_uuid = 'test-uuid'
ap.sess_mgr.session_list = [session]
ap.bot_service = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = PipelineService(ap)
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid'})
# Execute
await service.update_pipeline('test-uuid', {'description': 'Updated'})
# Verify - conversation was cleared
assert session.using_conversation is None
class TestPipelineServiceDeletePipeline:
"""Tests for delete_pipeline method."""
async def test_delete_pipeline_calls_remove_and_delete(self):
"""Calls both pipeline_mgr.remove_pipeline and persistence delete."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
service = PipelineService(ap)
# Execute
await service.delete_pipeline('test-uuid')
# Verify
ap.pipeline_mgr.remove_pipeline.assert_called_once_with('test-uuid')
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_pipeline_nonexistent_uuid(self):
"""Delete operation completes even for nonexistent UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
service = PipelineService(ap)
# Execute - should not raise
await service.delete_pipeline('nonexistent-uuid')
# Verify
ap.pipeline_mgr.remove_pipeline.assert_called_once()
class TestPipelineServiceCopyPipeline:
"""Tests for copy_pipeline method."""
async def test_copy_pipeline_max_limit_reached_raises(self):
"""Raises ValueError when max_pipelines limit reached."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': 2}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
service = PipelineService(ap)
# Mock get_pipelines to return 2 pipelines
service.get_pipelines = AsyncMock(
return_value=[
{'uuid': 'uuid-1', 'name': 'Pipeline 1'},
{'uuid': 'uuid-2', 'name': 'Pipeline 2'},
]
)
# Execute & Verify
with pytest.raises(ValueError, match='Maximum number of pipelines'):
await service.copy_pipeline('original-uuid')
async def test_copy_pipeline_not_found_raises(self):
"""Raises ValueError when original pipeline not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': -1}}}
ap.pipeline_mgr = SimpleNamespace()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
service = PipelineService(ap)
service.get_pipelines = AsyncMock(return_value=[]) # No limit check issue
ap.persistence_mgr.execute_async = AsyncMock(
return_value=_create_mock_result(first_item=None) # Original not found
)
ap.persistence_mgr.serialize_model = Mock(return_value={})
# Execute & Verify
with pytest.raises(ValueError, match='Pipeline original-uuid not found'):
await service.copy_pipeline('original-uuid')
async def test_copy_pipeline_creates_copy(self):
"""Creates a copy with (Copy) suffix."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': -1}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
original = _create_mock_pipeline(
pipeline_uuid='original-uuid',
name='Original Pipeline',
description='Original description',
stages=['Stage1', 'Stage2'],
config={'key': 'value'},
extensions_preferences={'enable_all_plugins': False, 'plugins': ['plugin1']},
)
service = PipelineService(ap)
service.get_pipelines = AsyncMock(return_value=[]) # No limit check issue
# Mock persistence - get original, then insert, then get new
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=original))
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'new-copy-uuid',
'name': 'Original Pipeline (Copy)',
}
)
service.get_pipeline = AsyncMock(
return_value={
'uuid': 'new-copy-uuid',
'name': 'Original Pipeline (Copy)',
}
)
# Execute
new_uuid = await service.copy_pipeline('original-uuid')
# Verify
assert new_uuid is not None
assert len(new_uuid) == 36 # UUID format
async def test_copy_pipeline_is_not_default(self):
"""Copy is never set as default."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'limitation': {'max_pipelines': -1}}}
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.load_pipeline = AsyncMock()
ap.ver_mgr = SimpleNamespace()
ap.ver_mgr.get_current_version = Mock(return_value='1.0.0')
# Original is default
original = _create_mock_pipeline(
pipeline_uuid='original-uuid',
name='Default Pipeline',
is_default=True,
)
service = PipelineService(ap)
service.get_pipelines = AsyncMock(return_value=[])
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=original))
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'copy-uuid', 'is_default': False})
service.get_pipeline = AsyncMock(return_value={'uuid': 'copy-uuid', 'is_default': False})
# Execute
await service.copy_pipeline('original-uuid')
# Verify - pipeline_mgr.load_pipeline called (copy created)
ap.pipeline_mgr.load_pipeline.assert_called_once()
class TestPipelineServiceUpdatePipelineExtensions:
"""Tests for update_pipeline_extensions method."""
async def test_update_extensions_pipeline_not_found_raises(self):
"""Raises ValueError when pipeline not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = PipelineService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='Pipeline nonexistent-uuid not found'):
await service.update_pipeline_extensions('nonexistent-uuid', [])
async def test_update_extensions_sets_plugins(self):
"""Updates plugins in extensions_preferences."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
original_pipeline = _create_mock_pipeline(extensions_preferences={'enable_all_plugins': True, 'plugins': []})
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=original_pipeline)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'extensions_preferences': {
'enable_all_plugins': False,
'plugins': [{'plugin_uuid': 'plugin-1'}],
},
}
)
service = PipelineService(ap)
service.get_pipeline = AsyncMock(
return_value={
'uuid': 'test-uuid',
'extensions_preferences': {
'enable_all_plugins': False,
'plugins': [{'plugin_uuid': 'plugin-1'}],
},
}
)
# Execute
bound_plugins = [{'plugin_uuid': 'plugin-1'}]
await service.update_pipeline_extensions(
'test-uuid',
bound_plugins=bound_plugins,
enable_all_plugins=False,
)
# Verify
ap.persistence_mgr.execute_async.assert_called()
async def test_update_extensions_sets_mcp_servers(self):
"""Updates MCP servers in extensions_preferences."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
original_pipeline = _create_mock_pipeline()
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=original_pipeline)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'extensions_preferences': {
'enable_all_mcp_servers': False,
'mcp_servers': ['mcp-server-1'],
},
}
)
service = PipelineService(ap)
service.get_pipeline = AsyncMock(
return_value={
'uuid': 'test-uuid',
'extensions_preferences': {'mcp_servers': ['mcp-server-1']},
}
)
# Execute
await service.update_pipeline_extensions(
'test-uuid',
bound_plugins=[],
bound_mcp_servers=['mcp-server-1'],
enable_all_mcp_servers=False,
)
# Verify
ap.persistence_mgr.execute_async.assert_called()
async def test_update_extensions_none_mcp_servers_keeps_existing(self):
"""Does not modify mcp_servers when bound_mcp_servers is None."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
original_pipeline = _create_mock_pipeline(
extensions_preferences={
'enable_all_plugins': True,
'enable_all_mcp_servers': True,
'plugins': [],
'mcp_servers': ['existing-server'],
}
)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=original_pipeline)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'test-uuid', 'extensions_preferences': {'mcp_servers': ['existing-server']}}
)
service = PipelineService(ap)
service.get_pipeline = AsyncMock(
return_value={'uuid': 'test-uuid', 'extensions_preferences': {'mcp_servers': ['existing-server']}}
)
# Execute - bound_mcp_servers is None (not provided)
await service.update_pipeline_extensions('test-uuid', bound_plugins=[])
# Verify - persistence was called
ap.persistence_mgr.execute_async.assert_called()
async def test_update_extensions_preserves_mcp_resource_agent_read_when_omitted(self):
"""Does not reset mcp_resource_agent_read_enabled when omitted by older clients."""
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.pipeline_mgr = SimpleNamespace()
ap.pipeline_mgr.remove_pipeline = AsyncMock()
ap.pipeline_mgr.load_pipeline = AsyncMock()
original_pipeline = _create_mock_pipeline(
extensions_preferences={
'enable_all_plugins': True,
'enable_all_mcp_servers': True,
'plugins': [],
'mcp_servers': [],
'mcp_resources': [{'server_uuid': 'srv-1', 'uri': 'file:///README.md'}],
'mcp_resource_agent_read_enabled': False,
}
)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=original_pipeline)
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'test-uuid'})
service = PipelineService(ap)
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid'})
await service.update_pipeline_extensions('test-uuid', bound_plugins=[])
assert original_pipeline.extensions_preferences['mcp_resource_agent_read_enabled'] is False
assert original_pipeline.extensions_preferences['mcp_resources'] == [
{'server_uuid': 'srv-1', 'uri': 'file:///README.md'}
]
class TestDefaultStageOrder:
"""Tests for default_stage_order constant."""
def test_default_stage_order_not_empty(self):
"""Default stage order is not empty."""
assert len(default_stage_order) > 0
def test_default_stage_order_contains_key_stages(self):
"""Default stage order contains key processing stages."""
assert 'MessageProcessor' in default_stage_order
assert 'SendResponseBackStage' in default_stage_order
@@ -0,0 +1,870 @@
"""
Unit tests for ModelProviderService.
Tests model provider management operations including:
- Provider CRUD operations
- Provider model count checking
- Find or create provider logic
- Space model provider API key updates
- Provider model scanning
Source: src/langbot/pkg/api/http/service/provider.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.service.provider import ModelProviderService
from langbot.pkg.entity.persistence.model import ModelProvider, LLMModel, EmbeddingModel, RerankModel
pytestmark = pytest.mark.asyncio
def _create_mock_provider(
provider_uuid: str = 'test-provider-uuid',
name: str = 'Test Provider',
requester: str = 'openai',
base_url: str = 'https://api.openai.com',
api_keys: list = None,
) -> Mock:
"""Helper to create mock ModelProvider entity."""
provider = Mock(spec=ModelProvider)
provider.uuid = provider_uuid
provider.name = name
provider.requester = requester
provider.base_url = base_url
provider.api_keys = api_keys or ['test-key']
return provider
def _create_mock_llm_model(
model_uuid: str = 'test-llm-uuid',
name: str = 'Test LLM',
provider_uuid: str = 'test-provider-uuid',
) -> Mock:
"""Helper to create mock LLMModel entity."""
model = Mock(spec=LLMModel)
model.uuid = model_uuid
model.name = name
model.provider_uuid = provider_uuid
return model
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
result.scalar = Mock(return_value=len(items) if items else 0)
return result
class TestModelProviderServiceGetProviders:
"""Tests for get_providers method."""
async def test_get_providers_empty_list(self):
"""Returns empty list when no providers exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
'requester': entity.requester,
'base_url': entity.base_url,
'api_keys': entity.api_keys,
}
)
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
# Verify
assert result == []
async def test_get_providers_returns_serialized_list(self):
"""Returns serialized list of providers."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
provider1 = _create_mock_provider(provider_uuid='provider-1', name='Provider 1')
provider2 = _create_mock_provider(provider_uuid='provider-2', name='Provider 2')
mock_result = _create_mock_result([provider1, provider2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
'requester': entity.requester,
'base_url': entity.base_url,
'api_keys': entity.api_keys,
}
)
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
# Verify
assert len(result) == 2
assert result[0]['name'] == 'Provider 1'
assert result[1]['name'] == 'Provider 2'
async def test_get_providers_parse_api_keys_json_string(self):
"""Parses api_keys from JSON string if needed."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='provider-1', api_keys='["key1", "key2"]')
mock_result = _create_mock_result([provider])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
'api_keys': entity.api_keys, # Returns string
}
)
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
# Verify - api_keys should be parsed from string
assert result[0]['api_keys'] == ['key1', 'key2']
async def test_get_providers_invalid_json_api_keys_returns_empty(self):
"""Returns empty list for invalid JSON api_keys."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='provider-1', api_keys='invalid-json')
mock_result = _create_mock_result([provider])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'uuid': entity.uuid,
'name': entity.name,
'api_keys': entity.api_keys, # Returns invalid string
}
)
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
# Verify - invalid JSON returns empty list
assert result[0]['api_keys'] == []
class TestModelProviderServiceGetProvider:
"""Tests for get_provider method."""
async def test_get_provider_by_uuid_found(self):
"""Returns provider when found by UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='found-uuid', name='Found Provider')
mock_result = _create_mock_result([], first_item=provider)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'found-uuid',
'name': 'Found Provider',
'api_keys': ['key'],
}
)
service = ModelProviderService(ap)
# Execute
result = await service.get_provider('found-uuid')
# Verify
assert result is not None
assert result['uuid'] == 'found-uuid'
async def test_get_provider_by_uuid_not_found(self):
"""Returns None when provider not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([], first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ModelProviderService(ap)
# Execute
result = await service.get_provider('nonexistent-uuid')
# Verify
assert result is None
class TestModelProviderServiceCreateProvider:
"""Tests for create_provider method."""
async def test_create_provider_generates_uuid(self):
"""Creates provider with generated UUID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.provider_dict = {}
# Mock load_provider to return runtime provider
runtime_provider = Mock()
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = 'generated-uuid'
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.persistence_mgr.execute_async = AsyncMock()
service = ModelProviderService(ap)
# Execute
provider_uuid = await service.create_provider(
{
'name': 'New Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
)
# Verify - UUID is generated
assert provider_uuid is not None
assert len(provider_uuid) == 36 # UUID format
async def test_create_provider_loads_to_runtime(self):
"""Loads provider to runtime model_mgr."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.provider_dict = {}
runtime_provider = Mock()
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = 'runtime-uuid'
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.persistence_mgr.execute_async = AsyncMock()
service = ModelProviderService(ap)
# Execute
result_uuid = await service.create_provider(
{
'name': 'Runtime Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
)
# Verify - provider added to runtime dict and UUID generated
ap.model_mgr.load_provider.assert_called_once()
assert result_uuid is not None
class TestModelProviderServiceUpdateProvider:
"""Tests for update_provider method."""
async def test_update_provider_removes_uuid_from_data(self):
"""Removes uuid from update data before persisting."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.reload_provider = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
service = ModelProviderService(ap)
# Execute
await service.update_provider(
'existing-uuid',
{
'uuid': 'should-be-removed', # Will be removed
'name': 'Updated Name',
},
)
# Verify - reload called
ap.model_mgr.reload_provider.assert_called_once_with('existing-uuid')
async def test_update_provider_reloads_runtime(self):
"""Reloads provider in runtime after update."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.reload_provider = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
service = ModelProviderService(ap)
# Execute
await service.update_provider('update-uuid', {'name': 'New Name'})
# Verify
ap.model_mgr.reload_provider.assert_called_once()
class TestModelProviderServiceDeleteProvider:
"""Tests for delete_provider method."""
async def test_delete_provider_with_llm_models_raises_error(self):
"""Raises ValueError when LLM models reference provider."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Mock LLM model exists - only return LLM result since that's first check
llm_result = _create_mock_result([], first_item=_create_mock_llm_model())
ap.persistence_mgr.execute_async = AsyncMock(return_value=llm_result)
service = ModelProviderService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='Cannot delete provider: LLM models'):
await service.delete_provider('provider-with-llm')
async def test_delete_provider_with_embedding_models_raises_error(self):
"""Raises ValueError when Embedding models reference provider."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Create results for each check type
llm_result = Mock()
llm_result.first = Mock(return_value=None) # No LLM models
embedding_result = Mock()
embedding_result.first = Mock(return_value=Mock(spec=EmbeddingModel)) # Has embedding model
rerank_result = Mock()
rerank_result.first = Mock(return_value=None)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return llm_result
elif call_count == 2:
return embedding_result
return rerank_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = ModelProviderService(ap)
# Execute & Verify - should raise embedding error (LLM check passes, embedding check fails)
with pytest.raises(ValueError, match='Cannot delete provider: Embedding models'):
await service.delete_provider('provider-with-embedding')
async def test_delete_provider_with_rerank_models_raises_error(self):
"""Raises ValueError when Rerank models reference provider."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Create results for each check type
llm_result = Mock()
llm_result.first = Mock(return_value=None) # No LLM models
embedding_result = Mock()
embedding_result.first = Mock(return_value=None) # No embedding models
rerank_result = Mock()
rerank_result.first = Mock(return_value=Mock(spec=RerankModel)) # Has rerank model
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return llm_result
elif call_count == 2:
return embedding_result
return rerank_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = ModelProviderService(ap)
# Execute & Verify - should raise rerank error (LLM and embedding checks pass, rerank check fails)
with pytest.raises(ValueError, match='Cannot delete provider: Rerank models'):
await service.delete_provider('provider-with-rerank')
async def test_delete_provider_no_models_success(self):
"""Deletes provider when no models reference it."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.remove_provider = AsyncMock()
# Mock no models reference provider
empty_result = Mock()
empty_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=empty_result)
service = ModelProviderService(ap)
# Execute
await service.delete_provider('provider-no-models')
# Verify - delete and remove called
ap.model_mgr.remove_provider.assert_called_once_with('provider-no-models')
class TestModelProviderServiceGetProviderModelCounts:
"""Tests for get_provider_model_counts method."""
async def test_get_model_counts_returns_correct_counts(self):
"""Returns correct counts for each model type."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Mock scalar results for counts
llm_result = Mock()
llm_result.scalar = Mock(return_value=3)
embedding_result = Mock()
embedding_result.scalar = Mock(return_value=2)
rerank_result = Mock()
rerank_result.scalar = Mock(return_value=1)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return llm_result
elif call_count == 2:
return embedding_result
return rerank_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = ModelProviderService(ap)
# Execute
result = await service.get_provider_model_counts('provider-uuid')
# Verify
assert result['llm_count'] == 3
assert result['embedding_count'] == 2
assert result['rerank_count'] == 1
async def test_get_model_counts_zero_counts(self):
"""Returns zero counts when no models."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
zero_result = Mock()
zero_result.scalar = Mock(return_value=0)
ap.persistence_mgr.execute_async = AsyncMock(return_value=zero_result)
service = ModelProviderService(ap)
# Execute
result = await service.get_provider_model_counts('empty-provider')
# Verify
assert result['llm_count'] == 0
assert result['embedding_count'] == 0
assert result['rerank_count'] == 0
class TestModelProviderServiceFindOrCreateProvider:
"""Tests for find_or_create_provider method."""
async def test_find_existing_provider_matching_config(self):
"""Returns existing provider UUID when config matches."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
existing_provider = _create_mock_provider(
provider_uuid='existing-uuid',
requester='openai',
base_url='https://api.openai.com',
api_keys=['key1', 'key2'],
)
mock_result = _create_mock_result([existing_provider])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ModelProviderService(ap)
# Execute
result = await service.find_or_create_provider(
requester='openai',
base_url='https://api.openai.com',
api_keys=['key1', 'key2'], # Same keys (sorted)
)
# Verify - returns existing UUID
assert result == 'existing-uuid'
async def test_find_existing_provider_keys_order_mismatch(self):
"""Returns existing provider when keys match but order differs."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
existing_provider = _create_mock_provider(
provider_uuid='existing-uuid',
requester='openai',
base_url='https://api.openai.com',
api_keys=['key1', 'key2'],
)
mock_result = _create_mock_result([existing_provider])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ModelProviderService(ap)
# Execute with reversed key order
result = await service.find_or_create_provider(
requester='openai',
base_url='https://api.openai.com',
api_keys=['key2', 'key1'], # Different order, should still match
)
# Verify - returns existing UUID (keys are sorted in comparison)
assert result == 'existing-uuid'
async def test_create_new_provider_no_match(self):
"""Creates new provider when no existing match."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.provider_dict = {}
runtime_provider = Mock()
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = None # Will be set by uuid.uuid4()
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
# Mock no existing providers
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ModelProviderService(ap)
# Execute
result = await service.find_or_create_provider(
requester='new-requester',
base_url='https://new.api.com',
api_keys=['new-key'],
)
# Verify - creates new provider with valid UUID format
assert result is not None
assert len(result) == 36 # UUID format
# Verify provider was loaded to runtime
ap.model_mgr.load_provider.assert_called_once()
async def test_create_provider_name_from_url_parse(self):
"""Creates provider with name parsed from URL."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.provider_dict = {}
runtime_provider = Mock()
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = 'parsed-url-uuid'
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ModelProviderService(ap)
# Execute
result_uuid = await service.find_or_create_provider(
requester='custom',
base_url='https://api.example.com/v1',
api_keys=['key'],
)
# Verify - name should be parsed from URL (api.example.com)
ap.model_mgr.load_provider.assert_called_once()
assert result_uuid is not None
class TestModelProviderServiceUpdateSpaceModelProviderApiKeys:
"""Tests for update_space_model_provider_api_keys method."""
async def test_update_space_provider_api_keys(self):
"""Updates Space provider API keys."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr.reload_provider = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
service = ModelProviderService(ap)
# Execute
await service.update_space_model_provider_api_keys('space-api-key')
# Verify - update and reload called for Space provider UUID
ap.model_mgr.reload_provider.assert_called_once_with('00000000-0000-0000-0000-000000000000')
class TestModelProviderServiceScanProviderModels:
"""Tests for scan_provider_models method."""
async def test_scan_provider_not_found_raises_error(self):
"""Raises ValueError when provider not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([], first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ModelProviderService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='provider not found'):
await service.scan_provider_models('nonexistent-uuid')
async def test_scan_provider_returns_models_list(self):
"""Returns scanned models list."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.llm_model_service = SimpleNamespace()
ap.embedding_models_service = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='scan-uuid')
mock_result = _create_mock_result([], first_item=provider)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'scan-uuid',
'name': 'Scan Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
)
# Mock runtime provider with scan capability
runtime_provider = Mock()
runtime_provider.requester = Mock()
runtime_provider.token_mgr = Mock()
runtime_provider.token_mgr.get_token = Mock(return_value='token')
runtime_provider.token_mgr.tokens = ['token']
# Mock scan_models to return models
async def mock_scan_models(token):
return {
'models': [
{'id': 'gpt-4', 'name': 'GPT-4', 'type': 'llm'},
{'id': 'text-embedding', 'name': 'Text Embedding', 'type': 'embedding'},
],
'debug': None,
}
runtime_provider.requester.scan_models = AsyncMock(side_effect=mock_scan_models)
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
# Mock existing model services
ap.llm_model_service.get_llm_models_by_provider = AsyncMock(return_value=[])
ap.embedding_models_service.get_embedding_models_by_provider = AsyncMock(return_value=[])
service = ModelProviderService(ap)
# Execute
result = await service.scan_provider_models('scan-uuid')
# Verify
assert 'models' in result
assert len(result['models']) == 2
async def test_scan_provider_filter_by_model_type(self):
"""Returns filtered models by type."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.llm_model_service = SimpleNamespace()
ap.embedding_models_service = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='filter-uuid')
mock_result = _create_mock_result([], first_item=provider)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'filter-uuid',
'name': 'Filter Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
)
runtime_provider = Mock()
runtime_provider.requester = Mock()
runtime_provider.token_mgr = Mock()
runtime_provider.token_mgr.get_token = Mock(return_value='token')
runtime_provider.token_mgr.tokens = ['token']
async def mock_scan_models(token):
return {
'models': [
{'id': 'gpt-4', 'name': 'GPT-4', 'type': 'llm'},
{'id': 'text-embedding', 'name': 'Text Embedding', 'type': 'embedding'},
],
'debug': None,
}
runtime_provider.requester.scan_models = AsyncMock(side_effect=mock_scan_models)
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.llm_model_service.get_llm_models_by_provider = AsyncMock(return_value=[])
ap.embedding_models_service.get_embedding_models_by_provider = AsyncMock(return_value=[])
service = ModelProviderService(ap)
# Execute - filter for LLM only
result = await service.scan_provider_models('filter-uuid', model_type='llm')
# Verify - only LLM models returned
assert len(result['models']) == 1
assert result['models'][0]['type'] == 'llm'
async def test_scan_provider_not_implemented_raises_error(self):
"""Raises ValueError when scan not implemented."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='no-scan-uuid')
mock_result = _create_mock_result([], first_item=provider)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'no-scan-uuid',
'name': 'No Scan Provider',
'requester': 'custom',
'base_url': 'https://custom.api.com',
'api_keys': ['key'],
}
)
runtime_provider = Mock()
runtime_provider.requester = Mock()
runtime_provider.token_mgr = Mock()
runtime_provider.token_mgr.get_token = Mock(return_value='token')
runtime_provider.token_mgr.tokens = ['token']
runtime_provider.requester.scan_models = AsyncMock(side_effect=NotImplementedError('scan not supported'))
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
service = ModelProviderService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='current provider does not support model scanning'):
await service.scan_provider_models('no-scan-uuid')
async def test_scan_provider_marks_already_added_models(self):
"""Marks models that are already added."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.llm_model_service = SimpleNamespace()
ap.embedding_models_service = SimpleNamespace()
provider = _create_mock_provider(provider_uuid='already-added-uuid')
mock_result = _create_mock_result([], first_item=provider)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'already-added-uuid',
'name': 'Already Added Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
)
runtime_provider = Mock()
runtime_provider.requester = Mock()
runtime_provider.token_mgr = Mock()
runtime_provider.token_mgr.get_token = Mock(return_value='token')
runtime_provider.token_mgr.tokens = ['token']
async def mock_scan_models(token):
return {
'models': [
{'id': 'existing-model', 'name': 'Existing Model', 'type': 'llm'},
{'id': 'new-model', 'name': 'New Model', 'type': 'llm'},
],
'debug': None,
}
runtime_provider.requester.scan_models = AsyncMock(side_effect=mock_scan_models)
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
# Mock existing LLM model
ap.llm_model_service.get_llm_models_by_provider = AsyncMock(return_value=[{'name': 'Existing Model'}])
ap.embedding_models_service.get_embedding_models_by_provider = AsyncMock(return_value=[])
service = ModelProviderService(ap)
# Execute
result = await service.scan_provider_models('already-added-uuid')
# Verify - existing model marked as already_added
existing_model = next(m for m in result['models'] if m['name'] == 'Existing Model')
assert existing_model['already_added'] is True
new_model = next(m for m in result['models'] if m['name'] == 'New Model')
assert new_model['already_added'] is False
@@ -0,0 +1,788 @@
"""
Unit tests for SpaceService.
Tests LangBot Space API interactions including:
- OAuth URL generation
- Token exchange and refresh
- User info retrieval
- Credits caching
- Model listing
Source: src/langbot/pkg/api/http/service/space.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from types import SimpleNamespace
import datetime
import time
from langbot.pkg.api.http.service.space import SpaceService
from langbot.pkg.entity.persistence.user import User
pytestmark = pytest.mark.asyncio
def _create_mock_user(
email: str = 'test@example.com',
account_type: str = 'space',
space_account_uuid: str = 'space-uuid-123',
space_access_token: str = 'access_token_123',
space_refresh_token: str = 'refresh_token_123',
space_access_token_expires_at: datetime.datetime = None,
) -> Mock:
"""Helper to create mock User entity."""
user = Mock(spec=User)
user.user = email
user.account_type = account_type
user.space_account_uuid = space_account_uuid
user.space_access_token = space_access_token
user.space_refresh_token = space_refresh_token
user.space_access_token_expires_at = space_access_token_expires_at
return user
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
class TestSpaceServiceGetOAuthAuthorizeUrl:
"""Tests for get_oauth_authorize_url method."""
def test_get_oauth_authorize_url_basic(self):
"""Returns OAuth URL with redirect_uri."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'space': {
'oauth_authorize_url': 'https://space.langbot.app/auth/authorize',
}
}
service = SpaceService(ap)
# Execute
result = service.get_oauth_authorize_url('http://localhost/callback')
# Verify
assert 'redirect_uri=http://localhost/callback' in result
assert 'https://space.langbot.app/auth/authorize' in result
def test_get_oauth_authorize_url_with_state(self):
"""Returns OAuth URL with redirect_uri and state."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {
'space': {
'oauth_authorize_url': 'https://space.langbot.app/auth/authorize',
}
}
service = SpaceService(ap)
# Execute
result = service.get_oauth_authorize_url('http://localhost/callback', state='random_state')
# Verify
assert 'redirect_uri=http://localhost/callback' in result
assert 'state=random_state' in result
def test_get_oauth_authorize_url_default_config(self):
"""Uses default OAuth URL when config not set."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Execute
result = service.get_oauth_authorize_url('http://localhost/callback')
# Verify - uses default URL
assert 'https://space.langbot.app/auth/authorize' in result
class TestSpaceServiceGetUserByEmail:
"""Tests for _get_user_by_email internal method."""
async def test_get_user_by_email_found(self):
"""Returns user when found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(email='found@example.com')
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._get_user_by_email('found@example.com')
# Verify
assert result is not None
assert result.user == 'found@example.com'
async def test_get_user_by_email_not_found(self):
"""Returns None when user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._get_user_by_email('notfound@example.com')
# Verify
assert result is None
class TestSpaceServiceEnsureValidToken:
"""Tests for _ensure_valid_token internal method."""
async def test_ensure_valid_token_user_not_found(self):
"""Returns None when user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._ensure_valid_token('notfound@example.com')
# Verify
assert result is None
async def test_ensure_valid_token_not_space_account(self):
"""Returns None when user is not a space account."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(email='local@example.com', account_type='local')
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._ensure_valid_token('local@example.com')
# Verify
assert result is None
async def test_ensure_valid_token_no_access_token(self):
"""Returns None when user has no access token."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(space_access_token=None)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._ensure_valid_token('test@example.com')
# Verify
assert result is None
async def test_ensure_valid_token_valid_token(self):
"""Returns valid access token when not expired."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Token expires in 1 hour (valid)
mock_user = _create_mock_user(
space_access_token='valid_token',
space_access_token_expires_at=datetime.datetime.now() + datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._ensure_valid_token('test@example.com')
# Verify
assert result == 'valid_token'
async def test_ensure_valid_token_expired_no_refresh(self):
"""Returns None when token expired and no refresh token."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Token expired 1 hour ago
mock_user = _create_mock_user(
space_access_token='expired_token',
space_refresh_token=None,
space_access_token_expires_at=datetime.datetime.now() - datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service._ensure_valid_token('test@example.com')
# Verify
assert result is None
class TestSpaceServiceGetCredits:
"""Tests for get_credits method."""
async def test_get_credits_no_user(self):
"""Returns None when user not found."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service.get_credits('notfound@example.com')
# Verify
assert result is None
async def test_get_credits_returns_cached_value(self):
"""Returns cached credits without API call."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Pre-populate cache
service._credits_cache = {'cached@example.com': (100, time.time())}
# Execute
result = await service.get_credits('cached@example.com')
# Verify - returns cached value without API call
assert result == 100
async def test_get_credits_cache_expired_refreshes(self):
"""Refreshes expired cache."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(
space_access_token='valid_token',
space_access_token_expires_at=datetime.datetime.now() + datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Pre-populate expired cache (70 seconds ago, past 60s TTL)
service._credits_cache = {'test@example.com': (50, time.time() - 70)}
# Mock get_user_info to return new credits
service.get_user_info = AsyncMock(return_value={'credits': 200})
# Execute
result = await service.get_credits('test@example.com')
# Verify - cache was refreshed
assert result == 200
assert service._credits_cache['test@example.com'][0] == 200
async def test_get_credits_force_refresh(self):
"""Force refresh ignores cache."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(
space_access_token='valid_token',
space_access_token_expires_at=datetime.datetime.now() + datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Pre-populate cache
service._credits_cache = {'test@example.com': (100, time.time())}
# Mock get_user_info to return new credits
service.get_user_info = AsyncMock(return_value={'credits': 300})
# Execute with force_refresh=True
result = await service.get_credits('test@example.com', force_refresh=True)
# Verify - fresh value returned
assert result == 300
async def test_get_credits_returns_cached_on_exception(self):
"""Returns cached fallback value when API fails."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(
space_access_token='valid_token',
space_access_token_expires_at=datetime.datetime.now() + datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Pre-populate expired cache - will try to refresh and fail
service._credits_cache = {'test@example.com': (150, time.time() - 70)}
# Mock get_user_info to raise exception
service.get_user_info = AsyncMock(side_effect=Exception('API Error'))
# Execute - should return cached fallback value (even though expired)
result = await service.get_credits('test@example.com')
# Verify - returns cached fallback value (150) because API failed
assert result == 150
class TestSpaceServiceRefreshToken:
"""Tests for refresh_token method."""
async def test_refresh_token_success(self):
"""Refreshes token successfully."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(
return_value={
'code': 0,
'data': {
'access_token': 'new_access_token',
'refresh_token': 'new_refresh_token',
'expires_in': 3600,
},
}
)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.post = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
# Use async context manager mock
mock_session_obj.post.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute
result = await service.refresh_token('old_refresh_token')
# Verify
assert result['access_token'] == 'new_access_token'
async def test_refresh_token_api_error(self):
"""Raises ValueError on API error."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response with error
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(
return_value={
'code': 1,
'msg': 'Invalid refresh token',
}
)
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Invalid refresh token"}')
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.post = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.post.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='Failed to refresh token'):
await service.refresh_token('invalid_refresh_token')
async def test_refresh_token_http_error(self):
"""Raises ValueError on HTTP error."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response with error status
mock_response = MagicMock()
mock_response.status = 500
mock_response.text = AsyncMock(return_value='Internal Server Error')
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.post = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.post.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='Failed to refresh token'):
await service.refresh_token('refresh_token')
class TestSpaceServiceExchangeOAuthCode:
"""Tests for exchange_oauth_code method."""
async def test_exchange_oauth_code_success(self):
"""Exchanges OAuth code successfully."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(
return_value={
'code': 0,
'data': {
'access_token': 'new_access_token',
'refresh_token': 'new_refresh_token',
'expires_in': 3600,
},
}
)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.post = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.post.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute
result = await service.exchange_oauth_code('auth_code')
# Verify
assert result['access_token'] == 'new_access_token'
async def test_exchange_oauth_code_api_error(self):
"""Raises ValueError on API error."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response with error
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Invalid code'})
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Invalid code"}')
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.post = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.post.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='Failed to exchange OAuth code'):
await service.exchange_oauth_code('invalid_code')
class TestSpaceServiceGetUserInfoRaw:
"""Tests for get_user_info_raw method."""
async def test_get_user_info_raw_success(self):
"""Gets user info successfully."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(
return_value={
'code': 0,
'data': {
'email': 'test@example.com',
'credits': 100,
},
}
)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.get = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.get.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.get.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute
result = await service.get_user_info_raw('access_token')
# Verify
assert result['email'] == 'test@example.com'
assert result['credits'] == 100
async def test_get_user_info_raw_api_error(self):
"""Raises ValueError on API error."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response with error
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Unauthorized'})
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Unauthorized"}')
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.get = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.get.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.get.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='Failed to get user info'):
await service.get_user_info_raw('invalid_token')
class TestSpaceServiceGetUserInfo:
"""Tests for get_user_info method (with token validation)."""
async def test_get_user_info_no_token(self):
"""Returns None when no valid token."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Execute
result = await service.get_user_info('notfound@example.com')
# Verify
assert result is None
async def test_get_user_info_with_valid_token(self):
"""Returns user info with valid token."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(
space_access_token='valid_token',
space_access_token_expires_at=datetime.datetime.now() + datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Mock get_user_info_raw
service.get_user_info_raw = AsyncMock(return_value={'email': 'test@example.com', 'credits': 100})
# Execute
result = await service.get_user_info('test@example.com')
# Verify
assert result['email'] == 'test@example.com'
class TestSpaceServiceGetModels:
"""Tests for get_models method."""
async def test_get_models_success(self):
"""Gets models successfully."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response with proper model data matching SpaceModel schema
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(
return_value={
'code': 0,
'data': {
'models': [
{
'uuid': 'uuid-1',
'model_id': 'model-1',
'provider': 'provider-1',
'category': 'chat',
'status': 'active',
},
{
'uuid': 'uuid-2',
'model_id': 'model-2',
'provider': 'provider-2',
'category': 'chat',
'status': 'active',
},
]
},
}
)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.get = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.get.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.get.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute
result = await service.get_models()
# Verify
assert len(result) == 2
async def test_get_models_api_error(self):
"""Raises ValueError on API error."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Mock HTTP response with error
mock_response = MagicMock()
mock_response.status = 200
mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Unauthorized'})
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Unauthorized"}')
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
mock_session_obj.get = MagicMock(return_value=mock_response)
mock_session.return_value = mock_session_obj
mock_session_obj.get.return_value.__aenter__ = AsyncMock(return_value=mock_response)
mock_session_obj.get.return_value.__aexit__ = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='Failed to get models'):
await service.get_models()
class TestSpaceServiceCreditsCache:
"""Tests for credits cache behavior."""
def test_credits_cache_initialized(self):
"""Verify _credits_cache is initialized as empty dict."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
service = SpaceService(ap)
# Verify
assert hasattr(service, '_credits_cache')
assert service._credits_cache == {}
async def test_credits_cache_updates_on_success(self):
"""Cache updates when get_credits succeeds."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(
space_access_token='valid_token',
space_access_token_expires_at=datetime.datetime.now() + datetime.timedelta(hours=1),
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = SpaceService(ap)
# Mock get_user_info
service.get_user_info = AsyncMock(return_value={'credits': 500})
# Execute
result = await service.get_credits('test@example.com')
# Verify - cache updated
assert result == 500
assert 'test@example.com' in service._credits_cache
assert service._credits_cache['test@example.com'][0] == 500
@@ -0,0 +1,610 @@
"""
Unit tests for UserService.
Tests user management operations including:
- User initialization check
- Local user creation and authentication
- JWT token generation and verification
- Password management (reset, change, set)
- Space account management
Source: src/langbot/pkg/api/http/service/user.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.errors.account import AccountEmailMismatchError
pytestmark = pytest.mark.asyncio
def _create_mock_user(
email: str = 'test@example.com',
password: str = 'hashed_password',
account_type: str = 'local',
space_account_uuid: str = None,
) -> Mock:
"""Helper to create mock User entity."""
user = Mock(spec=User)
user.user = email
user.password = password
user.account_type = account_type
user.space_account_uuid = space_account_uuid
return user
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
class TestUserServiceIsInitialized:
"""Tests for is_initialized method."""
async def test_is_initialized_returns_true_when_users_exist(self):
"""Returns True when at least one user exists."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user()
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.is_initialized()
# Verify
assert result is True
async def test_is_initialized_returns_false_when_no_users(self):
"""Returns False when no users exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.is_initialized()
# Verify
assert result is False
async def test_is_initialized_returns_false_on_none_result(self):
"""Returns False when result is None."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.all = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.is_initialized()
# Verify
assert result is False
class TestUserServiceGetUserByEmail:
"""Tests for get_user_by_email method."""
async def test_get_user_by_email_found(self):
"""Returns user when found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(email='found@example.com')
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_user_by_email('found@example.com')
# Verify
assert result is not None
assert result.user == 'found@example.com'
async def test_get_user_by_email_not_found(self):
"""Returns None when user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_user_by_email('notfound@example.com')
# Verify
assert result is None
async def test_get_user_by_email_empty_string(self):
"""Handles empty email string."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_user_by_email('')
# Verify
assert result is None
class TestUserServiceGetUserBySpaceAccountUuid:
"""Tests for get_user_by_space_account_uuid method."""
async def test_get_user_by_space_uuid_found(self):
"""Returns user when Space UUID found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(
email='space@example.com',
account_type='space',
space_account_uuid='space-uuid-123',
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_user_by_space_account_uuid('space-uuid-123')
# Verify
assert result is not None
assert result.space_account_uuid == 'space-uuid-123'
async def test_get_user_by_space_uuid_not_found(self):
"""Returns None when Space UUID not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_user_by_space_account_uuid('nonexistent-uuid')
# Verify
assert result is None
class TestUserServiceAuthenticate:
"""Tests for authenticate method."""
async def test_authenticate_user_not_found_raises_error(self):
"""Raises ValueError when user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 3600}}}
service = UserService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='用户不存在'):
await service.authenticate('nonexistent@example.com', 'password')
async def test_authenticate_space_user_without_password_raises_error(self):
"""Raises ValueError for Space user without local password."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Space user has empty password
mock_user = _create_mock_user(
email='space@example.com',
password='', # Empty password for Space user
account_type='space',
)
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='请使用 Space 账户登录'):
await service.authenticate('space@example.com', 'password')
class TestUserServiceGenerateJwtToken:
"""Tests for generate_jwt_token method."""
async def test_generate_jwt_token_returns_valid_token(self):
"""Generates valid JWT token."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 3600}}}
service = UserService(ap)
# Execute
token = await service.generate_jwt_token('test@example.com')
# Verify - JWT format (base64 encoded parts)
assert token is not None
assert len(token) > 0
parts = token.split('.')
assert len(parts) == 3 # JWT has 3 parts
async def test_generate_jwt_token_custom_expire(self):
"""Generates token with custom expiry."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 7200}}}
service = UserService(ap)
# Execute
token = await service.generate_jwt_token('test@example.com')
# Verify
assert token is not None
class TestUserServiceVerifyJwtToken:
"""Tests for verify_jwt_token method."""
async def test_verify_jwt_token_valid(self):
"""Verifies valid JWT token and returns user email."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 3600}}}
service = UserService(ap)
# First generate a valid token
token = await service.generate_jwt_token('verify@example.com')
# Execute
user_email = await service.verify_jwt_token(token)
# Verify
assert user_email == 'verify@example.com'
async def test_verify_jwt_token_invalid_raises_error(self):
"""Raises error for invalid JWT token."""
# Setup
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 3600}}}
service = UserService(ap)
# Execute & Verify - invalid token should raise JWT error
with pytest.raises(Exception): # jwt.DecodeError or similar
await service.verify_jwt_token('invalid.token.here')
class TestUserServiceResetPassword:
"""Tests for reset_password method."""
async def test_reset_password_updates_password(self):
"""Updates user password."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = UserService(ap)
# Execute
await service.reset_password('test@example.com', 'new_password')
# Verify - execute_async was called with update
ap.persistence_mgr.execute_async.assert_called_once()
class TestUserServiceChangePassword:
"""Tests for change_password method."""
async def test_change_password_user_not_found_raises_error(self):
"""Raises ValueError when user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
service = UserService(ap)
# Mock get_user_by_email to return None
service.get_user_by_email = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='User not found'):
await service.change_password('nonexistent@example.com', 'current', 'new')
async def test_change_password_no_local_password_raises_error(self):
"""Raises ValueError when user has no local password set."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
service = UserService(ap)
# Mock user without password
mock_user = _create_mock_user(email='nopass@example.com', password=None)
service.get_user_by_email = AsyncMock(return_value=mock_user)
# Execute & Verify
with pytest.raises(ValueError, match='No local password set'):
await service.change_password('nopass@example.com', 'current', 'new')
class TestUserServiceGetFirstUser:
"""Tests for get_first_user method."""
async def test_get_first_user_found(self):
"""Returns first user when exists."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_user = _create_mock_user(email='first@example.com')
mock_result = _create_mock_result([mock_user])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_first_user()
# Verify
assert result is not None
assert result.user == 'first@example.com'
async def test_get_first_user_not_found(self):
"""Returns None when no users exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = UserService(ap)
# Execute
result = await service.get_first_user()
# Verify
assert result is None
class TestUserServiceSetPassword:
"""Tests for set_password method."""
async def test_set_password_user_not_found_raises_error(self):
"""Raises ValueError when user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
service = UserService(ap)
# Mock get_user_by_email to return None
service.get_user_by_email = AsyncMock(return_value=None)
# Execute & Verify
with pytest.raises(ValueError, match='User not found'):
await service.set_password('nonexistent@example.com', 'new_password')
async def test_set_password_with_existing_password_requires_current(self):
"""Requires current password when user has existing password."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
service = UserService(ap)
# Mock user with existing password
mock_user = _create_mock_user(email='haspass@example.com', password='hashed_old_password')
service.get_user_by_email = AsyncMock(return_value=mock_user)
# Execute & Verify - should raise when no current_password provided
with pytest.raises(ValueError, match='Current password is required'):
await service.set_password('haspass@example.com', 'new_password')
class TestUserServiceCreateOrUpdateSpaceUser:
"""Tests for create_or_update_space_user method."""
async def test_create_or_update_existing_space_user(self):
"""Updates existing Space user tokens."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.provider_service = SimpleNamespace()
ap.provider_service.update_space_model_provider_api_keys = AsyncMock()
service = UserService(ap)
# Mock existing Space user
existing_user = _create_mock_user(
email='space@example.com',
account_type='space',
space_account_uuid='existing-space-uuid',
)
service.get_user_by_space_account_uuid = AsyncMock(return_value=existing_user)
service.get_user_by_email = AsyncMock(return_value=None)
service.is_initialized = AsyncMock(return_value=True)
ap.persistence_mgr.execute_async = AsyncMock()
# Execute
updated_user = await service.create_or_update_space_user(
space_account_uuid='existing-space-uuid',
email='space@example.com',
access_token='new_access_token',
refresh_token='new_refresh_token',
api_key='new_api_key',
expires_in=3600,
)
# Verify - update was called and user returned
ap.persistence_mgr.execute_async.assert_called()
assert updated_user.space_account_uuid == 'existing-space-uuid'
async def test_create_or_update_new_space_user_first_init(self):
"""Creates new Space user on first initialization."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.provider_service = SimpleNamespace()
ap.provider_service.update_space_model_provider_api_keys = AsyncMock()
service = UserService(ap)
# Mock new user to be returned after creation
new_user = _create_mock_user(
email='newspace@example.com',
account_type='space',
space_account_uuid='new-space-uuid',
)
# First call (line 138) returns None, second call (line 194) returns new_user
call_count = 0
async def mock_get_by_space_uuid(uuid):
nonlocal call_count
call_count += 1
if call_count == 1: # First check for existing user
return None
return new_user # After insert, return the new user
service.get_user_by_space_account_uuid = AsyncMock(side_effect=mock_get_by_space_uuid)
service.get_user_by_email = AsyncMock(return_value=None)
service.is_initialized = AsyncMock(return_value=False) # Not initialized
ap.persistence_mgr.execute_async = AsyncMock()
# Execute
result = await service.create_or_update_space_user(
space_account_uuid='new-space-uuid',
email='newspace@example.com',
access_token='access_token',
refresh_token='refresh_token',
api_key='api_key',
expires_in=3600,
)
# Verify
assert result.space_account_uuid == 'new-space-uuid'
async def test_create_or_update_space_user_already_initialized_raises_error(self):
"""Raises AccountEmailMismatchError when system already initialized and user not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.provider_service = SimpleNamespace()
ap.provider_service.update_space_model_provider_api_keys = AsyncMock()
service = UserService(ap)
# Mock system already initialized, no matching users
service.get_user_by_space_account_uuid = AsyncMock(return_value=None)
service.get_user_by_email = AsyncMock(return_value=None)
service.is_initialized = AsyncMock(return_value=True) # Already initialized
# Execute & Verify
with pytest.raises(AccountEmailMismatchError):
await service.create_or_update_space_user(
space_account_uuid='unknown-space-uuid',
email='unknown@example.com',
access_token='token',
refresh_token='refresh',
api_key='key',
expires_in=3600,
)
async def test_create_or_update_space_user_no_expiry(self):
"""Creates Space user without token expiry."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.provider_service = SimpleNamespace()
ap.provider_service.update_space_model_provider_api_keys = AsyncMock()
service = UserService(ap)
new_user = _create_mock_user(
email='noexpiry@example.com',
account_type='space',
space_account_uuid='noexpiry-uuid',
)
# First call (line 138) returns None, second call (line 194) returns new_user
call_count = 0
async def mock_get_by_space_uuid(uuid):
nonlocal call_count
call_count += 1
if call_count == 1: # First check for existing user
return None
return new_user # After insert, return the new user
service.get_user_by_space_account_uuid = AsyncMock(side_effect=mock_get_by_space_uuid)
service.get_user_by_email = AsyncMock(return_value=None)
service.is_initialized = AsyncMock(return_value=False)
ap.persistence_mgr.execute_async = AsyncMock()
# Execute with expires_in=0 (no expiry)
result = await service.create_or_update_space_user(
space_account_uuid='noexpiry-uuid',
email='noexpiry@example.com',
access_token='token',
refresh_token='refresh',
api_key='key',
expires_in=0, # No expiry
)
# Verify
assert result is not None
assert result.space_account_uuid == 'noexpiry-uuid'
class TestUserServiceCreateUserLock:
"""Tests for create_user_lock attribute."""
def test_create_user_lock_initialized(self):
"""Verify create_user_lock is initialized as asyncio.Lock."""
# Setup
ap = SimpleNamespace()
service = UserService(ap)
# Verify lock exists
assert hasattr(service, '_create_user_lock')
assert service._create_user_lock is not None
@@ -0,0 +1,507 @@
"""
Unit tests for WebhookService.
Tests webhook CRUD operations including:
- Webhook listing
- Webhook creation
- Webhook retrieval by ID
- Webhook updates
- Webhook deletion
- Enabled webhooks filtering
Source: src/langbot/pkg/api/http/service/webhook.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.service.webhook import WebhookService
from langbot.pkg.entity.persistence.webhook import Webhook
pytestmark = pytest.mark.asyncio
def _create_mock_webhook(
webhook_id: int = 1,
name: str = 'Test Webhook',
url: str = 'http://example.com/webhook',
description: str = 'Test Description',
enabled: bool = True,
) -> Mock:
"""Helper to create mock Webhook entity."""
webhook = Mock(spec=Webhook)
webhook.id = webhook_id
webhook.name = name
webhook.url = url
webhook.description = description
webhook.enabled = enabled
return webhook
def _create_mock_result(items: list = None, first_item=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
class TestWebhookServiceGetWebhooks:
"""Tests for get_webhooks method."""
async def test_get_webhooks_empty_list(self):
"""Returns empty list when no webhooks exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'url': entity.url,
}
)
service = WebhookService(ap)
# Execute
result = await service.get_webhooks()
# Verify
assert result == []
async def test_get_webhooks_returns_serialized_list(self):
"""Returns serialized list of webhooks."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
webhook1 = _create_mock_webhook(webhook_id=1, name='Webhook 1')
webhook2 = _create_mock_webhook(webhook_id=2, name='Webhook 2')
mock_result = _create_mock_result([webhook1, webhook2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'url': entity.url,
'description': entity.description,
'enabled': entity.enabled,
}
)
service = WebhookService(ap)
# Execute
result = await service.get_webhooks()
# Verify
assert len(result) == 2
assert result[0]['name'] == 'Webhook 1'
assert result[1]['name'] == 'Webhook 2'
class TestWebhookServiceCreateWebhook:
"""Tests for create_webhook method."""
async def test_create_webhook_full_params(self):
"""Creates webhook with all parameters."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Mock insert result
insert_result = Mock()
# Mock select result for retrieving created webhook
created_webhook = _create_mock_webhook(
webhook_id=1,
name='New Webhook',
url='http://new.example.com/webhook',
description='New Description',
enabled=True,
)
select_result = _create_mock_result(first_item=created_webhook)
# execute_async returns different results
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return insert_result # Insert
return select_result # Select
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'New Webhook',
'url': 'http://new.example.com/webhook',
'description': 'New Description',
'enabled': True,
}
)
service = WebhookService(ap)
# Execute
result = await service.create_webhook(
name='New Webhook',
url='http://new.example.com/webhook',
description='New Description',
enabled=True,
)
# Verify
assert result['name'] == 'New Webhook'
assert result['url'] == 'http://new.example.com/webhook'
assert result['description'] == 'New Description'
assert result['enabled'] is True
async def test_create_webhook_defaults(self):
"""Creates webhook with default description and enabled."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
created_webhook = _create_mock_webhook(
webhook_id=1,
name='Minimal Webhook',
url='http://minimal.example.com',
description='', # Default
enabled=True, # Default
)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return Mock() # Insert
return _create_mock_result(first_item=created_webhook)
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'Minimal Webhook',
'url': 'http://minimal.example.com',
'description': '',
'enabled': True,
}
)
service = WebhookService(ap)
# Execute - only name and url required
result = await service.create_webhook(name='Minimal Webhook', url='http://minimal.example.com')
# Verify defaults
assert result['description'] == ''
assert result['enabled'] is True
async def test_create_webhook_disabled(self):
"""Creates webhook with enabled=False."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
created_webhook = _create_mock_webhook(webhook_id=1, enabled=False)
call_count = 0
async def mock_execute(query):
nonlocal call_count
call_count += 1
if call_count == 1:
return Mock()
return _create_mock_result(first_item=created_webhook)
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(return_value={'id': 1, 'enabled': False})
service = WebhookService(ap)
# Execute
result = await service.create_webhook(name='Disabled', url='http://disabled.com', enabled=False)
# Verify
assert result['enabled'] is False
class TestWebhookServiceGetWebhook:
"""Tests for get_webhook method."""
async def test_get_webhook_by_id_found(self):
"""Returns webhook when found by ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
webhook = _create_mock_webhook(webhook_id=1, name='Found Webhook')
mock_result = _create_mock_result(first_item=webhook)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'Found Webhook',
'url': 'http://example.com/webhook',
}
)
service = WebhookService(ap)
# Execute
result = await service.get_webhook(1)
# Verify
assert result is not None
assert result['id'] == 1
assert result['name'] == 'Found Webhook'
async def test_get_webhook_by_id_not_found(self):
"""Returns None when webhook not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = WebhookService(ap)
# Execute
result = await service.get_webhook(999)
# Verify
assert result is None
async def test_get_webhook_by_id_zero(self):
"""Handles ID=0 (edge case) correctly."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = WebhookService(ap)
# Execute
result = await service.get_webhook(0)
# Verify - should return None (no webhook with ID 0)
assert result is None
class TestWebhookServiceUpdateWebhook:
"""Tests for update_webhook method."""
async def test_update_webhook_name_only(self):
"""Updates only the name field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute
await service.update_webhook(1, name='Updated Name')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_webhook_url_only(self):
"""Updates only the url field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute
await service.update_webhook(1, url='http://updated.example.com')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_webhook_description_only(self):
"""Updates only the description field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute
await service.update_webhook(1, description='Updated description')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_webhook_enabled_only(self):
"""Updates only the enabled field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute
await service.update_webhook(1, enabled=False)
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_webhook_all_fields(self):
"""Updates all fields at once."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute
await service.update_webhook(
1,
name='All Updated',
url='http://all.updated.com',
description='All updated description',
enabled=False,
)
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_webhook_no_fields(self):
"""Does nothing when no fields provided."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute - no update parameters
await service.update_webhook(1)
# Verify - no execute call since no update_data
ap.persistence_mgr.execute_async.assert_not_called()
class TestWebhookServiceDeleteWebhook:
"""Tests for delete_webhook method."""
async def test_delete_webhook_by_id(self):
"""Deletes webhook by ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute
await service.delete_webhook(1)
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_webhook_nonexistent_id(self):
"""Delete operation completes even for nonexistent ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = WebhookService(ap)
# Execute - should not raise
await service.delete_webhook(999)
# Verify - still called
ap.persistence_mgr.execute_async.assert_called_once()
class TestWebhookServiceGetEnabledWebhooks:
"""Tests for get_enabled_webhooks method."""
async def test_get_enabled_webhooks_empty(self):
"""Returns empty list when no enabled webhooks."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={})
service = WebhookService(ap)
# Execute
result = await service.get_enabled_webhooks()
# Verify
assert result == []
async def test_get_enabled_webhooks_filters_enabled(self):
"""Returns only enabled webhooks."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# All returned webhooks should be enabled (SQL filter)
webhook1 = _create_mock_webhook(webhook_id=1, name='Enabled 1', enabled=True)
webhook2 = _create_mock_webhook(webhook_id=2, name='Enabled 2', enabled=True)
mock_result = _create_mock_result([webhook1, webhook2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'enabled': entity.enabled,
}
)
service = WebhookService(ap)
# Execute
result = await service.get_enabled_webhooks()
# Verify
assert len(result) == 2
assert all(w['enabled'] for w in result)
async def test_get_enabled_webhooks_filters_disabled(self):
"""Does not return disabled webhooks."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Empty result because query filters on enabled=True
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={})
service = WebhookService(ap)
# Execute
result = await service.get_enabled_webhooks()
# Verify - should be empty (SQL would filter disabled)
assert result == []
@@ -0,0 +1,42 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.service.apikey import ApiKeyService
@pytest.mark.asyncio
@pytest.mark.parametrize('api_key', [None, 123, b'lbk_bytes', '', 'plain_key', ' LBK_bad', 'sk-lbk_fake'])
async def test_verify_api_key_rejects_non_lbk_keys_without_db_query(api_key):
persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
service = ApiKeyService(SimpleNamespace(persistence_mgr=persistence_mgr, instance_config=instance_config))
result = await service.verify_api_key(api_key)
assert result is False
persistence_mgr.execute_async.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
('db_row', 'expected'),
[
(object(), True),
(None, False),
],
)
async def test_verify_api_key_keeps_db_validation_for_lbk_keys(db_row, expected):
query_result = Mock()
query_result.first.return_value = db_row
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(return_value=query_result))
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
service = ApiKeyService(SimpleNamespace(persistence_mgr=persistence_mgr, instance_config=instance_config))
result = await service.verify_api_key('lbk_valid_format')
assert result is expected
persistence_mgr.execute_async.assert_awaited_once()
@@ -0,0 +1,18 @@
from __future__ import annotations
import pytest
from langbot.pkg.api.http.controller.groups.box_visibility import should_hide_box_runtime_status
@pytest.mark.parametrize(
('edition', 'box_enabled', 'expected'),
[
('cloud', False, True),
('cloud', True, False),
('cloud', None, False),
('community', False, False),
],
)
def test_should_hide_box_runtime_status(edition, box_enabled, expected):
assert should_hide_box_runtime_status(edition, box_enabled) is expected
@@ -0,0 +1,76 @@
from __future__ import annotations
import sys
import types
from importlib import import_module
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
core_app_module = types.ModuleType('langbot.pkg.core.app')
core_app_module.Application = object
sys.modules.setdefault('langbot.pkg.core.app', core_app_module)
pytestmark = pytest.mark.asyncio
async def _create_test_client(mcp_service: SimpleNamespace):
app = quart.Quart(__name__)
user_service = SimpleNamespace(
verify_jwt_token=AsyncMock(return_value='test@example.com'),
get_user_by_email=AsyncMock(return_value=SimpleNamespace(user='test@example.com')),
)
ap = SimpleNamespace(mcp_service=mcp_service, user_service=user_service)
MCPRouterGroup = import_module('langbot.pkg.api.http.controller.groups.resources.mcp').MCPRouterGroup
group = MCPRouterGroup(ap, app)
await group.initialize()
return app.test_client()
async def test_mcp_server_route_accepts_encoded_slash_name():
mcp_service = SimpleNamespace(
get_mcp_server_by_name=AsyncMock(
return_value={
'uuid': 'test-uuid',
'name': 'pab1it0/prometheus',
'enable': True,
'mode': 'stdio',
'extra_args': {},
}
)
)
client = await _create_test_client(mcp_service)
response = await client.get(
'/api/v1/mcp/servers/pab1it0%2Fprometheus',
headers={'Authorization': 'Bearer test-token'},
)
assert response.status_code == 200
mcp_service.get_mcp_server_by_name.assert_awaited_once_with('pab1it0/prometheus')
payload = await response.get_json()
assert payload['data']['server']['name'] == 'pab1it0/prometheus'
async def test_mcp_resource_route_accepts_encoded_slash_name():
mcp_service = SimpleNamespace(
get_mcp_server_by_name=AsyncMock(),
get_mcp_server_resources=AsyncMock(return_value=[]),
get_mcp_server_resource_templates=AsyncMock(return_value=[]),
get_runtime_info=AsyncMock(return_value={'resource_capabilities': {'subscribe': False}}),
)
client = await _create_test_client(mcp_service)
response = await client.get(
'/api/v1/mcp/servers/pab1it0%2Fprometheus/resources',
headers={'Authorization': 'Bearer test-token'},
)
assert response.status_code == 200
mcp_service.get_mcp_server_by_name.assert_not_awaited()
mcp_service.get_mcp_server_resources.assert_awaited_once_with('pab1it0/prometheus')
payload = await response.get_json()
assert payload['data']['resource_capabilities'] == {'subscribe': False}
+106
View File
@@ -0,0 +1,106 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot.pkg.box.connector import BoxRuntimeConnector
def make_app(logger: Mock, runtime_endpoint: str = ''):
return SimpleNamespace(
logger=logger,
instance_config=SimpleNamespace(
data={
'box': {
'backend': 'local',
'runtime': {'endpoint': runtime_endpoint},
'local': {
'profile': 'default',
'allowed_mount_roots': [],
'default_workspace': '',
},
'e2b': {'api_key': '', 'api_url': '', 'template': ''},
}
}
),
)
def test_box_runtime_connector_stdio_when_no_url(monkeypatch: pytest.MonkeyPatch):
"""Without runtime.endpoint, on a non-Docker Unix platform, use stdio."""
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
assert connector._uses_websocket() is False
assert isinstance(connector.client, ActionRPCBoxClient)
def test_box_runtime_connector_ws_when_url_configured(monkeypatch: pytest.MonkeyPatch):
"""With an explicit runtime.endpoint, always use WebSocket."""
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
logger = Mock()
connector = BoxRuntimeConnector(make_app(logger, runtime_endpoint='http://box-runtime:5410'))
assert connector._uses_websocket() is True
assert isinstance(connector.client, ActionRPCBoxClient)
def test_box_runtime_connector_ws_in_docker(monkeypatch: pytest.MonkeyPatch):
"""Inside Docker (no explicit URL), use WebSocket to reach a sibling container."""
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'docker')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
assert connector._uses_websocket() is True
assert connector.ws_relay_base_url == 'http://langbot_box:5410'
def test_box_runtime_connector_ws_with_standalone_flag(monkeypatch: pytest.MonkeyPatch):
"""With --standalone-box flag, use WebSocket even on a local Unix platform."""
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', True)
connector = BoxRuntimeConnector(make_app(Mock()))
assert connector._uses_websocket() is True
def test_box_runtime_connector_ws_relay_url_default(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
assert connector.ws_relay_base_url == 'http://127.0.0.1:5410'
def test_box_runtime_connector_ws_relay_url_explicit(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
assert connector.ws_relay_base_url == 'http://box-runtime:5410'
def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
logger = Mock()
connector = BoxRuntimeConnector(make_app(logger))
subprocess = Mock()
subprocess.returncode = None
handler_task = Mock()
ctrl_task = Mock()
connector._subprocess = subprocess
connector._handler_task = handler_task
connector._ctrl_task = ctrl_task
connector.dispose()
subprocess.terminate.assert_called_once()
handler_task.cancel.assert_called_once()
ctrl_task.cancel.assert_called_once()
assert connector._handler_task is None
assert connector._ctrl_task is None
File diff suppressed because it is too large Load Diff
+149
View File
@@ -0,0 +1,149 @@
from __future__ import annotations
import os
import tempfile
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.box.workspace import (
BoxWorkspaceSession,
classify_python_workspace,
infer_workspace_host_path,
rewrite_mounted_path,
wrap_python_command_with_env,
)
def test_rewrite_mounted_path_translates_host_prefix():
result = rewrite_mounted_path('/tmp/demo/project/app.py', '/tmp/demo/project')
assert result == '/workspace/app.py'
def test_infer_workspace_host_path_unwraps_virtualenv_bin_dir():
with tempfile.TemporaryDirectory() as tmpdir:
project_root = os.path.join(tmpdir, 'project')
os.makedirs(os.path.join(project_root, '.venv', 'bin'))
python_bin = os.path.join(project_root, '.venv', 'bin', 'python')
script = os.path.join(project_root, 'server.py')
with open(python_bin, 'w', encoding='utf-8') as handle:
handle.write('')
with open(script, 'w', encoding='utf-8') as handle:
handle.write('print("ok")\n')
result = infer_workspace_host_path(python_bin, [script])
assert result == os.path.realpath(project_root)
def test_classify_python_workspace_detects_package_and_requirements():
with tempfile.TemporaryDirectory() as tmpdir:
assert classify_python_workspace(tmpdir) is None
with open(os.path.join(tmpdir, 'requirements.txt'), 'w', encoding='utf-8') as handle:
handle.write('requests\n')
assert classify_python_workspace(tmpdir) == 'requirements'
with open(os.path.join(tmpdir, 'pyproject.toml'), 'w', encoding='utf-8') as handle:
handle.write('[project]\nname = "demo"\n')
assert classify_python_workspace(tmpdir) == 'package'
def test_wrap_python_command_with_env_contains_bootstrap_and_command():
command = wrap_python_command_with_env('python script.py')
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
assert 'kill -0 "$_LB_LOCK_OWNER"' in command
assert 'export VIRTUAL_ENV="$_LB_VENV_DIR"' in command
assert command.rstrip().endswith('python script.py')
@pytest.mark.asyncio
async def test_workspace_session_execute_for_query_uses_session_payload():
box_service = SimpleNamespace(execute_spec_payload=AsyncMock(return_value={'ok': True}))
workspace = BoxWorkspaceSession(
box_service,
'skill-person_123-demo',
host_path='/tmp/project',
host_path_mode='rw',
env={'FOO': 'bar'},
)
query = SimpleNamespace(query_id='q1')
result = await workspace.execute_for_query(query, 'python run.py', workdir='/workspace', timeout_sec=30)
assert result == {'ok': True}
payload = box_service.execute_spec_payload.await_args.args[0]
assert payload == {
'session_id': 'skill-person_123-demo',
'workdir': '/workspace',
'env': {'FOO': 'bar'},
'persistent': False,
'host_path': '/tmp/project',
'host_path_mode': 'rw',
'cmd': 'python run.py',
'timeout_sec': 30,
}
@pytest.mark.asyncio
async def test_workspace_session_start_managed_process_rewrites_command_and_args():
box_service = SimpleNamespace(start_managed_process=AsyncMock(return_value={'status': 'running'}))
workspace = BoxWorkspaceSession(
box_service,
'mcp-u1',
host_path='/tmp/project',
host_path_mode='ro',
)
result = await workspace.start_managed_process(
'/tmp/project/.venv/bin/python',
['/tmp/project/server.py', '--config', '/tmp/project/config.json'],
env={'TOKEN': '1'},
)
assert result == {'status': 'running'}
session_id = box_service.start_managed_process.await_args.args[0]
payload = box_service.start_managed_process.await_args.args[1]
assert session_id == 'mcp-u1'
assert payload == {
'command': 'python',
'args': ['/workspace/server.py', '--config', '/workspace/config.json'],
'env': {'TOKEN': '1'},
'cwd': '/workspace',
'process_id': 'default',
}
def test_workspace_session_build_session_payload_keeps_generic_workspace_shape():
workspace = BoxWorkspaceSession(
Mock(),
'workspace-1',
host_path='/tmp/project',
host_path_mode='rw',
env={'FOO': 'bar'},
network='on',
read_only_rootfs=False,
image='python:3.11',
cpus=1.0,
memory_mb=512,
pids_limit=128,
)
assert workspace.build_session_payload() == {
'session_id': 'workspace-1',
'workdir': '/workspace',
'env': {'FOO': 'bar'},
'persistent': False,
'network': 'on',
'read_only_rootfs': False,
'host_path': '/tmp/project',
'host_path_mode': 'rw',
'image': 'python:3.11',
'cpus': 1.0,
'memory_mb': 512,
'pids_limit': 128,
}
+1
View File
@@ -0,0 +1 @@
# Unit tests for command module
+532
View File
@@ -0,0 +1,532 @@
"""
Unit tests for cmdmgr module - REAL imports.
Tests CommandManager initialization, execute, and privilege handling.
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from langbot.pkg.command import operator
from langbot.pkg.command.cmdmgr import CommandManager
from tests.factories import FakeApp, command_query
import langbot_plugin.api.entities.builtin.provider.session as provider_session
class TestCommandManagerInit:
"""Tests for CommandManager initialization."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
@pytest.mark.asyncio
async def test_init_does_not_set_cmd_list(self):
"""CommandManager.__init__ does not set cmd_list (set in initialize())."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
assert mgr.ap is fake_app
assert not hasattr(mgr, 'cmd_list') # Not set until initialize()
@pytest.mark.asyncio
async def test_initialize_sets_path_for_top_level_commands(self):
"""initialize() sets path for top-level commands."""
@operator.operator_class(name='help')
class HelpOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='status')
class StatusOperator(operator.CommandOperator):
async def execute(self, context):
yield None
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
# Check paths are set
help_op = next(op for op in mgr.cmd_list if op.name == 'help')
status_op = next(op for op in mgr.cmd_list if op.name == 'status')
assert help_op.path == 'help'
assert status_op.path == 'status'
@pytest.mark.asyncio
async def test_initialize_sets_path_for_nested_commands(self):
"""initialize() sets path for nested commands."""
@operator.operator_class(name='plugin')
class PluginOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='list', parent_class=PluginOperator)
class PluginListOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='install', parent_class=PluginOperator)
class PluginInstallOperator(operator.CommandOperator):
async def execute(self, context):
yield None
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
plugin_op = next(op for op in mgr.cmd_list if op.name == 'plugin')
list_op = next(op for op in mgr.cmd_list if op.name == 'list')
install_op = next(op for op in mgr.cmd_list if op.name == 'install')
assert plugin_op.path == 'plugin'
assert list_op.path == 'plugin.list'
assert install_op.path == 'plugin.install'
@pytest.mark.asyncio
async def test_initialize_sets_children_for_parent_commands(self):
"""initialize() sets children list for parent commands."""
@operator.operator_class(name='parent')
class ParentOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='child1', parent_class=ParentOperator)
class Child1Operator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='child2', parent_class=ParentOperator)
class Child2Operator(operator.CommandOperator):
async def execute(self, context):
yield None
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
parent_op = next(op for op in mgr.cmd_list if op.name == 'parent')
child_names = [child.name for child in parent_op.children]
assert len(parent_op.children) == 2
assert 'child1' in child_names
assert 'child2' in child_names
@pytest.mark.asyncio
async def test_initialize_instantiates_all_operators(self):
"""initialize() instantiates all preregistered operators."""
@operator.operator_class(name='help')
class HelpOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='status')
class StatusOperator(operator.CommandOperator):
async def execute(self, context):
yield None
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
assert len(mgr.cmd_list) == 2
assert all(isinstance(op, operator.CommandOperator) for op in mgr.cmd_list)
@pytest.mark.asyncio
async def test_initialize_calls_operator_initialize(self):
"""initialize() calls initialize() on each operator."""
init_called = []
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def initialize(self):
init_called.append(self.name)
async def execute(self, context):
yield None
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
assert 'test' in init_called
@pytest.mark.asyncio
async def test_initialize_with_no_operators(self):
"""initialize() handles empty preregistered_operators."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
assert mgr.cmd_list == []
class TestCommandManagerExecute:
"""Tests for CommandManager execute method."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
def _create_session(self, launcher_type=provider_session.LauncherTypes.PERSON, launcher_id=12345):
"""Helper to create a session."""
return provider_session.Session(
launcher_type=launcher_type,
launcher_id=launcher_id,
sender_id=launcher_id,
use_prompt_name='default',
using_conversation=None,
conversations=[],
)
@pytest.mark.asyncio
async def test_execute_returns_generator(self):
"""execute() returns an async generator."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
# Mock plugin_connector.list_commands to return empty list
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('help')
session = self._create_session()
result = mgr.execute('help', '/help', query, session)
assert hasattr(result, '__aiter__')
@pytest.mark.asyncio
async def test_execute_sets_privilege_for_admin(self):
"""execute() sets privilege=2 for admin users."""
fake_app = FakeApp(admins=['person_12345'])
mgr = CommandManager(fake_app)
mgr.cmd_list = []
# Mock plugin_connector
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('status')
query.launcher_type = provider_session.LauncherTypes.PERSON
query.launcher_id = 12345
session = self._create_session()
results = []
async for ret in mgr.execute('status', '/status', query, session):
results.append(ret)
# Verify admin config was checked
assert 'person_12345' in fake_app.instance_config.data['admins']
@pytest.mark.asyncio
async def test_execute_sets_privilege_for_non_admin(self):
"""execute() sets privilege=1 for non-admin users."""
fake_app = FakeApp(admins=['person_12345'])
mgr = CommandManager(fake_app)
mgr.cmd_list = []
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('status')
query.launcher_type = provider_session.LauncherTypes.PERSON
query.launcher_id = 67890 # Not in admins list
session = self._create_session(launcher_id=67890)
results = []
async for ret in mgr.execute('status', '/status', query, session):
results.append(ret)
@pytest.mark.asyncio
async def test_execute_parses_command_text(self):
"""execute() splits command_text into params."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('help arg1 arg2')
session = self._create_session()
results = []
async for ret in mgr.execute('help arg1 arg2', '/help arg1 arg2', query, session):
results.append(ret)
# Command text parsing happens inside execute()
# We verify it doesn't crash
@pytest.mark.asyncio
async def test_execute_passes_bound_plugins(self):
"""execute() passes bound_plugins from query variables."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('help')
query.variables = {'_pipeline_bound_plugins': ['plugin1', 'plugin2']}
session = self._create_session()
results = []
async for ret in mgr.execute('help', '/help', query, session):
results.append(ret)
# Bound plugins are extracted from query.variables
assert query.variables.get('_pipeline_bound_plugins') == ['plugin1', 'plugin2']
class TestCommandManagerInternalExecute:
"""Tests for CommandManager._execute method."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
def _create_context(self, command='help', privilege=1):
"""Helper to create ExecuteContext."""
from langbot_plugin.api.entities.builtin.command import context as cmd_context
session = provider_session.Session(
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=12345,
use_prompt_name='default',
using_conversation=None,
conversations=[],
)
return cmd_context.ExecuteContext(
query_id=1,
session=session,
command_text='help',
full_command_text='/help',
command=command,
crt_command=command,
params=['help'],
crt_params=['help'],
privilege=privilege,
)
@pytest.mark.asyncio
async def test_execute_yields_command_not_found_error(self):
"""_execute yields CommandNotFoundError for unknown commands."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
# Mock plugin_connector.list_commands to return empty list
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
ctx = self._create_context(command='unknown_cmd')
results = []
async for ret in mgr._execute(ctx, mgr.cmd_list):
results.append(ret)
assert len(results) == 1
assert results[0].error is not None
assert '未知命令' in str(results[0].error)
@pytest.mark.asyncio
async def test_execute_calls_plugin_command(self):
"""_execute calls plugin connector for plugin commands."""
from langbot_plugin.api.entities.builtin.command import context as cmd_context
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
# Mock plugin command
mock_command = Mock()
mock_command.metadata.name = 'plugin_cmd'
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[mock_command])
async def mock_plugin_execute(ctx, bound_plugins):
yield cmd_context.CommandReturn(text='plugin response')
fake_app.plugin_connector.execute_command = mock_plugin_execute
ctx = self._create_context(command='plugin_cmd')
results = []
async for ret in mgr._execute(ctx, mgr.cmd_list):
results.append(ret)
assert len(results) == 1
assert results[0].text == 'plugin response'
@pytest.mark.asyncio
async def test_execute_with_bound_plugins(self):
"""_execute passes bound_plugins to plugin connector."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
# Mock plugin command
mock_command = Mock()
mock_command.metadata.name = 'test_cmd'
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[mock_command])
async def mock_execute_command(ctx, bound_plugins):
yield Mock(text='ok')
fake_app.plugin_connector.execute_command = mock_execute_command
ctx = self._create_context(command='test_cmd')
# Execute with bound_plugins parameter
async for _ in mgr._execute(ctx, mgr.cmd_list, bound_plugins=['test_plugin']):
pass
class TestEmptyAndEdgeInputs:
"""Tests for empty and edge inputs."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
def _create_session(self):
"""Helper to create a session."""
return provider_session.Session(
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=12345,
use_prompt_name='default',
using_conversation=None,
conversations=[],
)
@pytest.mark.asyncio
async def test_execute_with_empty_command_text(self):
"""execute() handles empty command_text."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('') # Empty command
session = self._create_session()
results = []
async for ret in mgr.execute('', '/', query, session):
results.append(ret)
# Should yield CommandNotFoundError for empty command
assert len(results) == 1
assert results[0].error is not None
@pytest.mark.asyncio
async def test_execute_with_whitespace_command(self):
"""execute() handles whitespace-only command_text."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query(' ') # Whitespace command
session = self._create_session()
results = []
async for ret in mgr.execute(' ', '/ ', query, session):
results.append(ret)
# Should yield error
assert len(results) >= 1
@pytest.mark.asyncio
async def test_initialize_with_deep_nesting(self):
"""initialize() handles deeply nested commands."""
@operator.operator_class(name='l1')
class L1Operator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='l2', parent_class=L1Operator)
class L2Operator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='l3', parent_class=L2Operator)
class L3Operator(operator.CommandOperator):
async def execute(self, context):
yield None
fake_app = FakeApp()
mgr = CommandManager(fake_app)
await mgr.initialize()
l3_op = next(op for op in mgr.cmd_list if op.name == 'l3')
assert l3_op.path == 'l1.l2.l3'
@pytest.mark.asyncio
async def test_execute_with_special_command_name(self):
"""execute() handles special characters in command name."""
fake_app = FakeApp()
mgr = CommandManager(fake_app)
mgr.cmd_list = []
fake_app.plugin_connector.list_commands = AsyncMock(return_value=[])
query = command_query('test-command_123')
session = self._create_session()
results = []
async for ret in mgr.execute('test-command_123', '/test-command_123', query, session):
results.append(ret)
# Should yield CommandNotFoundError (no such command registered)
assert len(results) == 1
assert results[0].error is not None
+303
View File
@@ -0,0 +1,303 @@
"""
Unit tests for operator module - REAL imports.
Tests the operator_class decorator and CommandOperator base class.
"""
from __future__ import annotations
import pytest
from langbot.pkg.command import operator
class TestOperatorClassDecorator:
"""Tests for operator_class decorator."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
def test_decorator_sets_name(self):
"""Decorator sets command name on class."""
@operator.operator_class(name='test_cmd')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.name == 'test_cmd'
def test_decorator_sets_help(self):
"""Decorator sets help text on class."""
@operator.operator_class(name='test', help='Test help message')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.help == 'Test help message'
def test_decorator_sets_usage(self):
"""Decorator sets usage text on class."""
@operator.operator_class(name='test', usage='!test <arg>')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.usage == '!test <arg>'
def test_decorator_sets_alias(self):
"""Decorator sets alias list on class."""
@operator.operator_class(name='test', alias=['t', 'tst'])
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.alias == ['t', 'tst']
def test_decorator_sets_privilege_default(self):
"""Decorator sets default privilege to 1 (normal user)."""
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.lowest_privilege == 1
def test_decorator_sets_privilege_admin(self):
"""Decorator sets privilege to 2 for admin commands."""
@operator.operator_class(name='admin_cmd', privilege=2)
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.lowest_privilege == 2
def test_decorator_sets_parent_class_none(self):
"""Decorator sets parent_class to None for top-level commands."""
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator.parent_class is None
def test_decorator_sets_parent_class(self):
"""Decorator sets parent_class for sub-commands."""
@operator.operator_class(name='parent')
class ParentOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='child', parent_class=ParentOperator)
class ChildOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert ChildOperator.parent_class is ParentOperator
def test_decorator_registers_to_preregistered_list(self):
"""Decorator appends class to preregistered_operators."""
@operator.operator_class(name='test1')
class TestOperator1(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='test2')
class TestOperator2(operator.CommandOperator):
async def execute(self, context):
yield None
assert TestOperator1 in operator.preregistered_operators
assert TestOperator2 in operator.preregistered_operators
def test_decorator_requires_command_operator_subclass(self):
"""Decorator asserts class is subclass of CommandOperator."""
with pytest.raises(AssertionError):
operator.operator_class(name='invalid')(object)
class TestCommandOperatorBase:
"""Tests for CommandOperator base class."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
def test_init_sets_app(self):
"""__init__ stores application reference."""
class MockApp:
pass
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
app = MockApp()
op = TestOperator(app)
assert op.ap is app
def test_init_sets_empty_children(self):
"""__init__ initializes empty children list."""
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
op = TestOperator(None)
assert op.children == []
def test_class_has_required_attributes(self):
"""CommandOperator has required class attributes."""
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert hasattr(TestOperator, 'name')
assert hasattr(TestOperator, 'alias')
assert hasattr(TestOperator, 'help')
assert hasattr(TestOperator, 'usage')
assert hasattr(TestOperator, 'parent_class')
assert hasattr(TestOperator, 'lowest_privilege')
def test_initialize_is_async_noop(self):
"""Default initialize() is async no-op."""
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
op = TestOperator(None)
# Should not raise
import asyncio
asyncio.get_event_loop().run_until_complete(op.initialize())
def test_execute_is_abstract(self):
"""execute() must be implemented by subclass."""
# Cannot instantiate abstract class
with pytest.raises(TypeError):
operator.CommandOperator(None)
def test_path_not_set_by_decorator(self):
"""path is not set by decorator, set by CommandManager."""
@operator.operator_class(name='test')
class TestOperator(operator.CommandOperator):
async def execute(self, context):
yield None
# path should not exist initially
assert not hasattr(TestOperator, 'path') or TestOperator.path is None
class TestMultipleOperators:
"""Tests for multiple operator registration and hierarchy."""
def setup_method(self):
"""Save and clear preregistered_operators before each test."""
self._saved_operators = operator.preregistered_operators.copy()
operator.preregistered_operators.clear()
def teardown_method(self):
"""Restore preregistered_operators after each test."""
operator.preregistered_operators.clear()
operator.preregistered_operators.extend(self._saved_operators)
def test_multiple_independent_operators(self):
"""Multiple independent operators can be registered."""
@operator.operator_class(name='help')
class HelpOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='status')
class StatusOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='version')
class VersionOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert len(operator.preregistered_operators) == 3
names = [op.name for op in operator.preregistered_operators]
assert 'help' in names
assert 'status' in names
assert 'version' in names
def test_parent_child_hierarchy(self):
"""Parent-child hierarchy can be established."""
@operator.operator_class(name='plugin')
class PluginOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='list', parent_class=PluginOperator)
class PluginListOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='install', parent_class=PluginOperator)
class PluginInstallOperator(operator.CommandOperator):
async def execute(self, context):
yield None
# Both parent and children are in preregistered list
assert len(operator.preregistered_operators) == 3
# Parent-child relationships are established via parent_class
plugin_op = next(op for op in operator.preregistered_operators if op.name == 'plugin')
list_op = next(op for op in operator.preregistered_operators if op.name == 'list')
install_op = next(op for op in operator.preregistered_operators if op.name == 'install')
assert plugin_op.parent_class is None
assert list_op.parent_class is PluginOperator
assert install_op.parent_class is PluginOperator
def test_privilege_inheritance_not_automatic(self):
"""Child operators do not automatically inherit parent privilege."""
@operator.operator_class(name='admin', privilege=2)
class AdminOperator(operator.CommandOperator):
async def execute(self, context):
yield None
@operator.operator_class(name='sub', parent_class=AdminOperator, privilege=1)
class SubOperator(operator.CommandOperator):
async def execute(self, context):
yield None
assert AdminOperator.lowest_privilege == 2
assert SubOperator.lowest_privilege == 1
+1
View File
@@ -0,0 +1 @@
# Config unit tests
@@ -0,0 +1,313 @@
"""
Unit tests for configuration loading and overrides.
Tests cover:
- Valid YAML config loading
- Valid JSON config loading
- Invalid YAML/JSON error behavior
- Missing config file behavior
- Template completion
"""
from __future__ import annotations
import pytest
import json
from langbot.pkg.config.impls.yaml import YAMLConfigFile
from langbot.pkg.config.impls.json import JSONConfigFile
from langbot.pkg.config.manager import ConfigManager
class TestYAMLConfigFile:
"""Tests for YAML config file handling."""
@pytest.mark.asyncio
async def test_valid_yaml_loads(self, tmp_path):
"""Valid YAML config should load correctly."""
config_file = tmp_path / 'test_config.yaml'
# Write valid YAML
config_file.write_text("""
name: test_app
version: 1.0
settings:
debug: true
port: 8080
""")
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'default', 'version': '0.1'},
)
result = await yaml_file.load(completion=False)
assert result['name'] == 'test_app'
assert result['version'] == 1.0
assert result['settings']['debug'] is True
assert result['settings']['port'] == 8080
@pytest.mark.asyncio
async def test_invalid_yaml_raises_error(self, tmp_path):
"""Invalid YAML should raise clear error."""
config_file = tmp_path / 'invalid.yaml'
# Write invalid YAML (unclosed bracket)
config_file.write_text("""
name: test
settings:
- item1
- item2
- [unclosed
""")
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'default'},
)
with pytest.raises(Exception, match='Syntax error'):
await yaml_file.load(completion=False)
@pytest.mark.asyncio
async def test_missing_config_creates_from_template(self, tmp_path):
"""Missing config file should be created from template."""
config_file = tmp_path / 'new_config.yaml'
# File doesn't exist yet
assert not config_file.exists()
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'new_app', 'version': '1.0'},
)
result = await yaml_file.load()
assert config_file.exists()
assert result['name'] == 'new_app'
assert result['version'] == '1.0'
@pytest.mark.asyncio
async def test_template_completion(self, tmp_path):
"""Config should be completed with template defaults."""
config_file = tmp_path / 'partial.yaml'
# Write partial config missing some template keys
config_file.write_text("""
name: custom_name
""")
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'default_name', 'version': '2.0', 'debug': False},
)
result = await yaml_file.load(completion=True)
# Existing key preserved
assert result['name'] == 'custom_name'
# Missing keys filled from template
assert result['version'] == '2.0'
assert result['debug'] is False
@pytest.mark.asyncio
async def test_yaml_save(self, tmp_path):
"""YAML config can be saved."""
config_file = tmp_path / 'save_test.yaml'
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'test'},
)
await yaml_file.save({'name': 'saved_app', 'new_key': 'new_value'})
assert config_file.exists()
content = config_file.read_text()
assert 'saved_app' in content
assert 'new_key' in content
def test_yaml_save_sync(self, tmp_path):
"""YAML config can be saved synchronously."""
config_file = tmp_path / 'sync_save.yaml'
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'test'},
)
yaml_file.save_sync({'name': 'sync_saved'})
assert config_file.exists()
content = config_file.read_text()
assert 'sync_saved' in content
class TestJSONConfigFile:
"""Tests for JSON config file handling."""
@pytest.mark.asyncio
async def test_valid_json_loads(self, tmp_path):
"""Valid JSON config should load correctly."""
config_file = tmp_path / 'test_config.json'
# Write valid JSON
config_file.write_text(
json.dumps(
{
'name': 'json_app',
'version': '1.0',
'settings': {'debug': True, 'port': 8080},
}
)
)
json_file = JSONConfigFile(
str(config_file),
template_data={'name': 'default', 'version': '0.1'},
)
result = await json_file.load(completion=False)
assert result['name'] == 'json_app'
assert result['version'] == '1.0'
assert result['settings']['debug'] is True
@pytest.mark.asyncio
async def test_invalid_json_raises_error(self, tmp_path):
"""Invalid JSON should raise clear error."""
config_file = tmp_path / 'invalid.json'
# Write invalid JSON (missing closing brace)
config_file.write_text('{"name": "test", "unclosed": ')
json_file = JSONConfigFile(
str(config_file),
template_data={'name': 'default'},
)
with pytest.raises(Exception, match='Syntax error'):
await json_file.load(completion=False)
@pytest.mark.asyncio
async def test_missing_json_creates_from_template(self, tmp_path):
"""Missing JSON file should be created from template."""
config_file = tmp_path / 'new_config.json'
json_file = JSONConfigFile(
str(config_file),
template_data={'name': 'new_json_app', 'version': '1.0'},
)
result = await json_file.load()
assert config_file.exists()
assert result['name'] == 'new_json_app'
@pytest.mark.asyncio
async def test_json_save(self, tmp_path):
"""JSON config can be saved."""
config_file = tmp_path / 'save_test.json'
json_file = JSONConfigFile(
str(config_file),
template_data={'name': 'test'},
)
await json_file.save({'name': 'saved_json', 'new_key': 'value'})
assert config_file.exists()
content = config_file.read_text()
data = json.loads(content)
assert data['name'] == 'saved_json'
class TestConfigManager:
"""Tests for ConfigManager."""
@pytest.mark.asyncio
async def test_config_manager_load(self, tmp_path):
"""ConfigManager loads config correctly."""
config_file = tmp_path / 'manager_test.yaml'
config_file.write_text('name: managed_app\nversion: "1.0"\n')
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'default', 'version': '0.1'},
)
manager = ConfigManager(yaml_file)
await manager.load_config()
assert manager.data['name'] == 'managed_app'
assert manager.data['version'] == '1.0'
@pytest.mark.asyncio
async def test_config_manager_dump(self, tmp_path):
"""ConfigManager can dump config."""
config_file = tmp_path / 'dump_test.yaml'
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'default'},
)
manager = ConfigManager(yaml_file)
manager.data = {'name': 'dumped', 'new_field': 'value'}
await manager.dump_config()
content = config_file.read_text()
assert 'dumped' in content
def test_config_manager_dump_sync(self, tmp_path):
"""ConfigManager can dump config synchronously."""
config_file = tmp_path / 'sync_dump.yaml'
yaml_file = YAMLConfigFile(
str(config_file),
template_data={'name': 'default'},
)
manager = ConfigManager(yaml_file)
manager.data = {'name': 'sync_dumped'}
manager.dump_config_sync()
assert config_file.exists()
class TestConfigExists:
"""Tests for config file existence check."""
def test_yaml_exists_true(self, tmp_path):
"""exists() returns True for existing file."""
config_file = tmp_path / 'exists.yaml'
config_file.write_text('name: test')
yaml_file = YAMLConfigFile(str(config_file), template_data={})
assert yaml_file.exists() is True
def test_yaml_exists_false(self, tmp_path):
"""exists() returns False for missing file."""
config_file = tmp_path / 'missing.yaml'
yaml_file = YAMLConfigFile(str(config_file), template_data={})
assert yaml_file.exists() is False
def test_json_exists_true(self, tmp_path):
"""exists() returns True for existing JSON file."""
config_file = tmp_path / 'exists.json'
config_file.write_text('{}')
json_file = JSONConfigFile(str(config_file), template_data={})
assert json_file.exists() is True
def test_json_exists_false(self, tmp_path):
"""exists() returns False for missing JSON file."""
config_file = tmp_path / 'missing.json'
json_file = JSONConfigFile(str(config_file), template_data={})
assert json_file.exists() is False
+1
View File
@@ -0,0 +1 @@
"""Core module unit tests."""
@@ -0,0 +1,192 @@
"""Unit tests for core app config validation methods.
Tests cover:
- _get_positive_int_config() validation
- _get_positive_float_config() validation
"""
from __future__ import annotations
from unittest.mock import Mock
from importlib import import_module
def get_app_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.core.app')
class TestGetPositiveIntConfig:
"""Tests for _get_positive_int_config method."""
def test_returns_value_when_valid_positive_int(self):
"""Test returns parsed int for valid positive value."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_int_config(10, default=30, name='test.config')
assert result == 10
mock_logger.warning.assert_not_called()
def test_returns_value_when_valid_string_int(self):
"""Test returns parsed int for string value."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_int_config('50', default=30, name='test.config')
assert result == 50
mock_logger.warning.assert_not_called()
def test_returns_default_for_zero(self):
"""Test returns default when value is zero."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_int_config(0, default=30, name='test.config')
assert result == 30
mock_logger.warning.assert_called_once()
def test_returns_default_for_negative(self):
"""Test returns default when value is negative."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_int_config(-5, default=30, name='test.config')
assert result == 30
mock_logger.warning.assert_called_once()
def test_returns_default_for_invalid_string(self):
"""Test returns default when value is invalid string."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_int_config('invalid', default=30, name='test.config')
assert result == 30
mock_logger.warning.assert_called_once()
def test_returns_default_for_none(self):
"""Test returns default when value is None."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_int_config(None, default=30, name='test.config')
assert result == 30
mock_logger.warning.assert_called_once()
class TestGetPositiveFloatConfig:
"""Tests for _get_positive_float_config method."""
def test_returns_value_when_valid_positive_float(self):
"""Test returns parsed float for valid positive value."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_float_config(1.5, default=2.0, name='test.config')
assert result == 1.5
mock_logger.warning.assert_not_called()
def test_returns_value_when_valid_int(self):
"""Test returns float for valid int value."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_float_config(2, default=1.0, name='test.config')
assert result == 2.0
mock_logger.warning.assert_not_called()
def test_returns_value_when_valid_string_float(self):
"""Test returns parsed float for string value."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_float_config('0.5', default=1.0, name='test.config')
assert result == 0.5
mock_logger.warning.assert_not_called()
def test_returns_default_for_zero(self):
"""Test returns default when value is zero."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_float_config(0.0, default=1.0, name='test.config')
assert result == 1.0
mock_logger.warning.assert_called_once()
def test_returns_default_for_negative(self):
"""Test returns default when value is negative."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_float_config(-1.0, default=2.0, name='test.config')
assert result == 2.0
mock_logger.warning.assert_called_once()
def test_returns_default_for_invalid_string(self):
"""Test returns default when value is invalid string."""
app_module = get_app_module()
mock_logger = Mock()
app = app_module.Application()
app.logger = mock_logger
result = app._get_positive_float_config('not-a-number', default=1.5, name='test.config')
assert result == 1.5
mock_logger.warning.assert_called_once()
+64
View File
@@ -0,0 +1,64 @@
from __future__ import annotations
import signal
from types import SimpleNamespace
import pytest
from langbot.pkg.core import boot
@pytest.mark.asyncio
async def test_main_signal_handler_handles_sigint_before_app_created(monkeypatch):
captured_handler = {}
def fake_signal(sig, handler):
captured_handler[sig] = handler
async def fake_make_app(loop):
captured_handler[signal.SIGINT](signal.SIGINT, None)
def fake_exit(code):
raise SystemExit(code)
monkeypatch.setattr(signal, 'signal', fake_signal)
monkeypatch.setattr(boot, 'make_app', fake_make_app)
monkeypatch.setattr(boot.os, '_exit', fake_exit)
with pytest.raises(SystemExit) as exc_info:
await boot.main(SimpleNamespace())
assert exc_info.value.code == 0
@pytest.mark.asyncio
async def test_main_signal_handler_disposes_created_app(monkeypatch):
captured_handler = {}
app_inst = SimpleNamespace(disposed=False)
def fake_signal(sig, handler):
captured_handler[sig] = handler
def dispose():
app_inst.disposed = True
async def run():
captured_handler[signal.SIGINT](signal.SIGINT, None)
async def fake_make_app(loop):
app_inst.dispose = dispose
app_inst.run = run
return app_inst
def fake_exit(code):
raise SystemExit(code)
monkeypatch.setattr(signal, 'signal', fake_signal)
monkeypatch.setattr(boot, 'make_app', fake_make_app)
monkeypatch.setattr(boot.os, '_exit', fake_exit)
with pytest.raises(SystemExit) as exc_info:
await boot.main(SimpleNamespace())
assert exc_info.value.code == 0
assert app_inst.disposed is True
@@ -0,0 +1,139 @@
"""Tests for core bootutils dependency checking."""
from __future__ import annotations
import importlib.util
from unittest.mock import MagicMock, patch
from tests.utils.import_isolation import isolated_sys_modules
class TestCheckDeps:
"""Tests for check_deps function."""
def _make_deps_import_mocks(self):
"""Create mocks for deps import."""
return {
'langbot.pkg.utils.pkgmgr': MagicMock(),
}
def test_check_deps_all_present(self):
"""check_deps returns empty list when all deps present."""
mocks = self._make_deps_import_mocks()
with isolated_sys_modules(mocks):
# Mock find_spec to always return a spec (module found)
with patch.object(importlib.util, 'find_spec', return_value=MagicMock()):
from langbot.pkg.core.bootutils.deps import check_deps
import asyncio
result = asyncio.get_event_loop().run_until_complete(check_deps())
assert result == []
def test_check_deps_missing_deps(self):
"""check_deps returns list of missing deps."""
mocks = self._make_deps_import_mocks()
with isolated_sys_modules(mocks):
# Mock find_spec to return None for some deps
def mock_find_spec(name):
if name in ['requests', 'openai']:
return None # Missing
return MagicMock() # Present
with patch.object(importlib.util, 'find_spec', side_effect=mock_find_spec):
from langbot.pkg.core.bootutils.deps import check_deps
import asyncio
result = asyncio.get_event_loop().run_until_complete(check_deps())
assert 'requests' in result
assert 'openai' in result
def test_check_deps_all_missing(self):
"""check_deps returns all deps when none present."""
mocks = self._make_deps_import_mocks()
with isolated_sys_modules(mocks):
# Mock find_spec to always return None
with patch.object(importlib.util, 'find_spec', return_value=None):
from langbot.pkg.core.bootutils.deps import check_deps, required_deps
import asyncio
result = asyncio.get_event_loop().run_until_complete(check_deps())
# Should include all required_deps keys
assert len(result) == len(required_deps)
def test_required_deps_dict_exists(self):
"""required_deps dictionary is defined."""
mocks = self._make_deps_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.bootutils.deps import required_deps
assert isinstance(required_deps, dict)
assert len(required_deps) > 0
# Check some expected deps
assert 'requests' in required_deps
assert 'yaml' in required_deps
def test_required_deps_maps_import_name_to_package_name(self):
"""required_deps maps import name to package name."""
mocks = self._make_deps_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.bootutils.deps import required_deps
# Some import names differ from package names
assert required_deps['PIL'] == 'pillow'
assert required_deps['yaml'] == 'pyyaml'
assert required_deps['jwt'] == 'pyjwt'
class TestPrecheckPluginDeps:
"""Tests for precheck_plugin_deps function."""
def _make_deps_import_mocks(self):
return {
'langbot.pkg.utils.pkgmgr': MagicMock(),
}
def test_precheck_plugin_deps_no_plugins_dir(self):
"""precheck_plugin_deps skips when plugins dir doesn't exist."""
from langbot.pkg.core.bootutils.deps import precheck_plugin_deps
with patch('os.path.exists', return_value=False):
with patch('langbot.pkg.core.bootutils.deps.pkgmgr.install_requirements') as mock_install:
import asyncio
asyncio.get_event_loop().run_until_complete(precheck_plugin_deps())
mock_install.assert_not_called()
def test_precheck_plugin_deps_with_plugins_dir(self):
"""precheck_plugin_deps checks plugins subdirectories."""
from langbot.pkg.core.bootutils.deps import precheck_plugin_deps
def mock_listdir(path):
if path == 'plugins':
return ['plugin1', 'plugin2']
if path == 'plugins/plugin1':
return ['requirements.txt', 'main.py']
if path == 'plugins/plugin2':
return ['main.py']
return []
with patch('os.path.exists', return_value=True):
with patch('os.path.isdir', return_value=True):
with patch('os.listdir', side_effect=mock_listdir):
with patch('langbot.pkg.core.bootutils.deps.pkgmgr.install_requirements') as mock_install:
import asyncio
asyncio.get_event_loop().run_until_complete(precheck_plugin_deps())
mock_install.assert_called_once_with('plugins/plugin1/requirements.txt', extra_params=[])
+120
View File
@@ -0,0 +1,120 @@
"""Tests for the daily-grouped rotating log file handler.
Regression coverage for the bug where a long-running process names its log
file after the *start* day and keeps appending to it across midnight, so no
file ever appears for the current day. See
``langbot.pkg.core.bootutils.log.DailyGroupedRotatingFileHandler``.
"""
from __future__ import annotations
import logging
import os
import re
import langbot.pkg.core.bootutils.log as logmod
from langbot.pkg.core.bootutils.log import DailyGroupedRotatingFileHandler
# Mirror of the cleanup pattern in api/http/service/maintenance.py.
MAINTENANCE_LOG_FILE_PATTERN = re.compile(r'^langbot-(\d{4}-\d{2}-\d{2})\.log(?:\.\d+)?$')
def _listing(directory):
return sorted(os.listdir(directory))
def _make_logger(handler, name):
logger = logging.getLogger(name)
logger.setLevel(logging.INFO)
logger.handlers.clear()
logger.addHandler(handler)
logger.propagate = False
return logger
class TestDailyGroupedRotatingFileHandler:
def _patch_date(self, monkeypatch, box):
"""Make the handler read its current date from ``box['date']``."""
def fake_strftime(fmt, t=None):
if fmt == '%Y-%m-%d':
return box['date']
return '00:00:00'
monkeypatch.setattr(logmod.time, 'strftime', fake_strftime)
def test_initial_file_named_for_current_day(self, tmp_path, monkeypatch):
box = {'date': '2026-06-08'}
self._patch_date(monkeypatch, box)
handler = DailyGroupedRotatingFileHandler(str(tmp_path), max_bytes=10_000, backup_count=3)
logger = _make_logger(handler, 'lb_logtest_initial')
logger.info('hello')
handler.close()
assert _listing(tmp_path) == ['langbot-2026-06-08.log']
def test_same_day_size_rotation_creates_numbered_backups(self, tmp_path, monkeypatch):
box = {'date': '2026-06-08'}
self._patch_date(monkeypatch, box)
handler = DailyGroupedRotatingFileHandler(str(tmp_path), max_bytes=200, backup_count=3)
logger = _make_logger(handler, 'lb_logtest_size')
for i in range(40):
logger.info('padding line to exceed maxBytes %d', i)
handler.close()
files = _listing(tmp_path)
assert 'langbot-2026-06-08.log' in files
assert any(f.startswith('langbot-2026-06-08.log.') for f in files)
def test_rolls_to_new_file_when_day_changes(self, tmp_path, monkeypatch):
box = {'date': '2026-06-08'}
self._patch_date(monkeypatch, box)
handler = DailyGroupedRotatingFileHandler(str(tmp_path), max_bytes=10_000, backup_count=3)
logger = _make_logger(handler, 'lb_logtest_midnight')
logger.info('day1 line')
# Simulate crossing midnight within the same running process.
box['date'] = '2026-06-09'
logger.info('day2 line after midnight')
handler.close()
files = _listing(tmp_path)
assert 'langbot-2026-06-08.log' in files
assert 'langbot-2026-06-09.log' in files
day2 = (tmp_path / 'langbot-2026-06-09.log').read_text(encoding='utf-8')
assert 'day2 line after midnight' in day2
assert 'day1 line' not in day2
def test_rollover_repeats_across_multiple_days(self, tmp_path, monkeypatch):
box = {'date': '2026-06-08'}
self._patch_date(monkeypatch, box)
handler = DailyGroupedRotatingFileHandler(str(tmp_path), max_bytes=10_000, backup_count=3)
logger = _make_logger(handler, 'lb_logtest_multiday')
for day in ('2026-06-08', '2026-06-09', '2026-06-10'):
box['date'] = day
logger.info('line for %s', day)
handler.close()
files = _listing(tmp_path)
for day in ('2026-06-08', '2026-06-09', '2026-06-10'):
assert f'langbot-{day}.log' in files
def test_all_filenames_match_maintenance_cleanup_pattern(self, tmp_path, monkeypatch):
box = {'date': '2026-06-08'}
self._patch_date(monkeypatch, box)
handler = DailyGroupedRotatingFileHandler(str(tmp_path), max_bytes=200, backup_count=3)
logger = _make_logger(handler, 'lb_logtest_pattern')
for i in range(40):
logger.info('padding line %d', i)
box['date'] = '2026-06-09'
logger.info('next day line')
handler.close()
for name in _listing(tmp_path):
assert MAINTENANCE_LOG_FILE_PATTERN.match(name), name
+284
View File
@@ -0,0 +1,284 @@
"""Unit tests for core stages load_config _apply_env_overrides_to_config.
Tests cover:
- Environment variable parsing and path conversion
- Type conversion (bool, int, float, string)
- List handling (comma-separated)
- Dict type skipping
- Missing key creation
"""
from __future__ import annotations
import os
from unittest.mock import patch
from importlib import import_module
def get_load_config_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.core.stages.load_config')
class TestApplyEnvOverridesToConfig:
"""Tests for _apply_env_overrides_to_config function."""
def test_override_string_value(self):
"""Test overriding an existing string config value."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default'}}
env = {'SYSTEM__NAME': 'custom_name'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['name'] == 'custom_name'
def test_override_int_value(self):
"""Test overriding an int value with proper conversion."""
load_config = get_load_config_module()
cfg = {'concurrency': {'pipeline': 5}}
env = {'CONCURRENCY__PIPELINE': '10'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['concurrency']['pipeline'] == 10
assert isinstance(result['concurrency']['pipeline'], int)
def test_override_int_value_invalid_conversion(self):
"""Test that invalid int conversion keeps string value."""
load_config = get_load_config_module()
cfg = {'concurrency': {'pipeline': 5}}
env = {'CONCURRENCY__PIPELINE': 'not_a_number'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
# Falls back to string when conversion fails
assert result['concurrency']['pipeline'] == 'not_a_number'
def test_override_bool_value_true(self):
"""Test overriding bool value with 'true' string."""
load_config = get_load_config_module()
cfg = {'system': {'enable': False}}
env = {'SYSTEM__ENABLE': 'true'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['enable'] is True
def test_override_bool_value_false(self):
"""Test overriding bool value with 'false' string."""
load_config = get_load_config_module()
cfg = {'system': {'enable': True}}
env = {'SYSTEM__ENABLE': 'false'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['enable'] is False
def test_override_bool_value_various_true_forms(self):
"""Test that '1', 'yes', 'on' are treated as true."""
load_config = get_load_config_module()
cfg = {'system': {'flag': False}}
for true_val in ['1', 'yes', 'on', 'TRUE']:
env = {'SYSTEM__FLAG': true_val}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg.copy())
assert result['system']['flag'] is True
def test_override_float_value(self):
"""Test overriding float value with proper conversion."""
load_config = get_load_config_module()
cfg = {'system': {'timeout': 1.5}}
env = {'SYSTEM__TIMEOUT': '2.5'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['timeout'] == 2.5
assert isinstance(result['system']['timeout'], float)
def test_override_list_value(self):
"""Test that comma-separated string converts to list."""
load_config = get_load_config_module()
cfg = {'system': {'disabled_adapters': ['adapter1']}}
env = {'SYSTEM__DISABLED_ADAPTERS': 'aiocqhttp,dingtalk,telegram'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['disabled_adapters'] == ['aiocqhttp', 'dingtalk', 'telegram']
def test_override_list_value_empty_items(self):
"""Test that empty items in comma-separated list are filtered."""
load_config = get_load_config_module()
cfg = {'system': {'disabled_adapters': []}}
env = {'SYSTEM__DISABLED_ADAPTERS': 'a,,b,,,c'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
# Empty items should be filtered out
assert result['system']['disabled_adapters'] == ['a', 'b', 'c']
def test_skip_dict_type_override(self):
"""Test that dict type values are skipped."""
load_config = get_load_config_module()
cfg = {'plugin': {'settings': {'nested': 'value'}}}
env = {'PLUGIN__SETTINGS': 'should_not_apply'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
# Dict type should not be overridden
assert result['plugin']['settings'] == {'nested': 'value'}
def test_create_new_key_when_missing(self):
"""Test that missing keys are created as strings."""
load_config = get_load_config_module()
cfg = {'system': {}}
env = {'SYSTEM__NEW_KEY': 'new_value'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['new_key'] == 'new_value'
def test_create_nested_path(self):
"""Test that intermediate dict is created for nested path."""
load_config = get_load_config_module()
cfg = {}
env = {'NEW__SECTION__KEY': 'value'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['new']['section']['key'] == 'value'
def test_skip_non_uppercase_env_vars(self):
"""Test that non-uppercase env vars are skipped."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default'}}
env = {'system__name': 'should_not_apply'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['name'] == 'default'
def test_skip_env_vars_without_double_underscore(self):
"""Test that env vars without __ are skipped."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default'}}
env = {'SYSTEMNAME': 'should_not_apply'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['name'] == 'default'
def test_nested_config_path(self):
"""Test overriding deeply nested config."""
load_config = get_load_config_module()
cfg = {'level1': {'level2': {'level3': 'original'}}}
env = {'LEVEL1__LEVEL2__LEVEL3': 'overridden'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['level1']['level2']['level3'] == 'overridden'
def test_non_dict_current_breaks(self):
"""Test that path navigation stops when current is not dict."""
load_config = get_load_config_module()
cfg = {'system': 'not_a_dict'}
env = {'SYSTEM__NAME': 'should_not_apply'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
# Should remain unchanged since 'system' is not a dict
assert result == {'system': 'not_a_dict'}
def test_empty_config(self):
"""Test that empty config dict is handled."""
load_config = get_load_config_module()
cfg = {}
env = {'SOME__KEY': 'value'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['some']['key'] == 'value'
def test_no_matching_env_vars(self):
"""Test that config is unchanged when no matching env vars."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default'}}
env = {'OTHER_VAR': 'value'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result == cfg
def test_multiple_env_vars_override(self):
"""Test multiple env vars applied in order."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default', 'enable': True}, 'concurrency': {'pipeline': 5}}
env = {'SYSTEM__NAME': 'custom', 'SYSTEM__ENABLE': 'false', 'CONCURRENCY__PIPELINE': '10'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['system']['name'] == 'custom'
assert result['system']['enable'] is False
assert result['concurrency']['pipeline'] == 10
def test_webhook_prefix_override(self):
"""Test overriding webhook_prefix via environment variable."""
load_config = get_load_config_module()
cfg = {'api': {'port': 5300, 'webhook_prefix': 'http://127.0.0.1:5300', 'extra_webhook_prefix': ''}}
env = {'API__WEBHOOK_PREFIX': 'https://example.com:8080'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['api']['webhook_prefix'] == 'https://example.com:8080'
def test_extra_webhook_prefix_override(self):
"""Test overriding extra_webhook_prefix via environment variable."""
load_config = get_load_config_module()
cfg = {'api': {'port': 5300, 'webhook_prefix': 'http://127.0.0.1:5300', 'extra_webhook_prefix': ''}}
env = {'API__EXTRA_WEBHOOK_PREFIX': 'https://extra.example.com'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['api']['extra_webhook_prefix'] == 'https://extra.example.com'
+178
View File
@@ -0,0 +1,178 @@
"""Tests for core boot stage registration and abstract classes."""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
from tests.utils.import_isolation import isolated_sys_modules
class TestStageClassDecorator:
"""Tests for @stage_class decorator."""
def _make_stage_import_mocks(self):
"""Create mocks for stage import."""
return {
'langbot.pkg.core.app': MagicMock(),
}
def test_stage_class_registers_stage(self):
"""@stage_class registers stage in preregistered_stages."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import stage_class, preregistered_stages
# Clear for clean test
preregistered_stages.clear()
@stage_class('TestStage')
class TestStage:
pass
assert 'TestStage' in preregistered_stages
assert preregistered_stages['TestStage'] == TestStage
def test_stage_class_returns_original_class(self):
"""@stage_class returns the original class unchanged."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import stage_class
@stage_class('TestStage')
class TestStage:
value = 42
# Class attributes should be preserved
assert TestStage.value == 42
def test_stage_class_multiple_stages(self):
"""Multiple stages can be registered."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import stage_class, preregistered_stages
preregistered_stages.clear()
@stage_class('Stage1')
class Stage1:
pass
@stage_class('Stage2')
class Stage2:
pass
assert len(preregistered_stages) == 2
assert preregistered_stages['Stage1'] == Stage1
assert preregistered_stages['Stage2'] == Stage2
class TestBootingStageAbstract:
"""Tests for BootingStage abstract class."""
def _make_stage_import_mocks(self):
return {'langbot.pkg.core.app': MagicMock()}
def test_booting_stage_is_abstract(self):
"""BootingStage is abstract and cannot be instantiated directly."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import BootingStage
with pytest.raises(TypeError):
BootingStage()
def test_booting_stage_requires_run_method(self):
"""Subclass must implement run method."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import BootingStage
class IncompleteStage(BootingStage):
pass
with pytest.raises(TypeError):
IncompleteStage()
def test_booting_stage_subclass_works(self):
"""Complete subclass can be instantiated."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import BootingStage
class CompleteStage(BootingStage):
name = 'CompleteStage'
async def run(self, ap):
pass
stage = CompleteStage()
assert stage.name == 'CompleteStage'
def test_booting_stage_name_attribute(self):
"""BootingStage has name attribute (None by default in abstract)."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import BootingStage
# Abstract class has name attribute defined as None
assert hasattr(BootingStage, 'name')
@pytest.mark.asyncio
async def test_booting_stage_run_signature(self):
"""run method receives Application parameter."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import BootingStage
class TestStage(BootingStage):
name = 'TestStage'
async def run(self, ap):
self.ap_received = ap
stage = TestStage()
mock_ap = MagicMock()
await stage.run(mock_ap)
assert stage.ap_received == mock_ap
class TestPreregisteredStages:
"""Tests for preregistered_stages global registry."""
def _make_stage_import_mocks(self):
return {'langbot.pkg.core.app': MagicMock()}
def test_preregistered_stages_is_dict(self):
"""preregistered_stages is a dictionary."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import preregistered_stages
assert isinstance(preregistered_stages, dict)
def test_preregistered_stages_key_is_string(self):
"""Registry keys are stage names (strings)."""
mocks = self._make_stage_import_mocks()
with isolated_sys_modules(mocks):
from langbot.pkg.core.stage import stage_class, preregistered_stages
preregistered_stages.clear()
@stage_class('MyStage')
class MyStage:
pass
for key in preregistered_stages:
assert isinstance(key, str)
+506
View File
@@ -0,0 +1,506 @@
"""Unit tests for core TaskContext, TaskWrapper, and AsyncTaskManager.
Tests cover:
- TaskContext initialization, state tracking, serialization
- TaskWrapper ID generation, to_dict serialization
- AsyncTaskManager task creation, stats, pruning
Note: Uses import_isolation to break circular import chains.
"""
from __future__ import annotations
import pytest
import asyncio
import sys
from unittest.mock import Mock, MagicMock
from contextlib import contextmanager
from typing import Generator
class MockLifecycleControlScopeEnum:
"""Mock enum value for LifecycleControlScope with .value attribute."""
def __init__(self, value: str):
self.value = value
def __repr__(self):
return f'LifecycleControlScope.{self.value.upper()}'
class MockLifecycleControlScope:
"""Mock enum for LifecycleControlScope."""
APPLICATION = MockLifecycleControlScopeEnum('application')
PLATFORM = MockLifecycleControlScopeEnum('platform')
PIPELINE = MockLifecycleControlScopeEnum('pipeline')
PLUGIN = MockLifecycleControlScopeEnum('plugin')
@contextmanager
def isolated_taskmgr_import() -> Generator[None, None, None]:
"""Context manager to isolate circular imports for taskmgr testing."""
# Mock modules that cause circular imports
mock_entities = MagicMock()
mock_entities.LifecycleControlScope = MockLifecycleControlScope
mock_app = MagicMock()
mock_importutil = MagicMock()
mock_importutil.import_modules_in_pkg = lambda pkg: None
mock_importutil.import_modules_in_pkgs = lambda pkgs: None
mock_http_controller = MagicMock()
mock_rag_mgr = MagicMock()
mocks = {
'langbot.pkg.core.entities': mock_entities,
'langbot.pkg.core.app': mock_app,
'langbot.pkg.api.http.controller.main': mock_http_controller,
'langbot.pkg.rag.knowledge.kbmgr': mock_rag_mgr,
'langbot.pkg.utils.importutil': mock_importutil,
}
# Save original state
saved = {}
for name in mocks:
if name in sys.modules:
saved[name] = sys.modules[name]
# Clear taskmgr to force re-import
taskmgr_name = 'langbot.pkg.core.taskmgr'
if taskmgr_name in sys.modules:
saved[taskmgr_name] = sys.modules[taskmgr_name]
try:
# Apply mocks
for name, module in mocks.items():
sys.modules[name] = module
# Clear taskmgr
sys.modules.pop(taskmgr_name, None)
yield
finally:
# Restore
for name in mocks:
if name in saved:
sys.modules[name] = saved[name]
else:
sys.modules.pop(name, None)
if taskmgr_name in saved:
sys.modules[taskmgr_name] = saved[taskmgr_name]
else:
sys.modules.pop(taskmgr_name, None)
def get_taskmgr_classes():
"""Get TaskContext, TaskWrapper, AsyncTaskManager classes."""
with isolated_taskmgr_import():
from langbot.pkg.core.taskmgr import TaskContext, TaskWrapper, AsyncTaskManager
return TaskContext, TaskWrapper, AsyncTaskManager
def create_mock_app():
"""Create a mock Application for testing."""
mock_app = Mock()
mock_app.event_loop = asyncio.get_running_loop()
mock_app.instance_config = Mock()
mock_app.instance_config.data = {
'system': {
'task_retention': {
'completed_limit': 200,
}
}
}
return mock_app
class TestTaskContext:
"""Tests for TaskContext class."""
def test_init_default_values(self):
"""Test that TaskContext initializes with default values."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
assert ctx.current_action == 'default'
assert ctx.log == ''
assert ctx.metadata == {}
def test_set_current_action(self):
"""Test setting current action."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
ctx.set_current_action('installing_plugin')
assert ctx.current_action == 'installing_plugin'
def test_trace_without_action(self):
"""Test trace method without action override."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
ctx.trace('Starting process')
assert 'Starting process' in ctx.log
assert ctx.current_action == 'default'
def test_trace_with_action_override(self):
"""Test trace method with action override."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
ctx.trace('Downloading', action='download')
assert 'Downloading' in ctx.log
assert ctx.current_action == 'download'
def test_trace_accumulates_logs(self):
"""Test that trace accumulates log entries."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
ctx.trace('Step 1')
ctx.trace('Step 2')
ctx.trace('Step 3')
assert 'Step 1' in ctx.log
assert 'Step 2' in ctx.log
assert 'Step 3' in ctx.log
# Each trace adds a newline
assert ctx.log.count('\n') == 3
def test_to_dict_serialization(self):
"""Test to_dict serialization."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
ctx.set_current_action('test_action')
ctx.trace('Test message')
ctx.metadata['key'] = 'value'
result = ctx.to_dict()
assert result['current_action'] == 'test_action'
assert 'Test message' in result['log']
assert result['metadata'] == {'key': 'value'}
def test_static_new_factory(self):
"""Test TaskContext.new() factory method."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext.new()
assert isinstance(ctx, TaskContext)
assert ctx.current_action == 'default'
def test_static_placeholder_singleton(self):
"""Test TaskContext.placeholder() returns singleton."""
with isolated_taskmgr_import():
from langbot.pkg.core.taskmgr import TaskContext
# Reset global placeholder
import langbot.pkg.core.taskmgr as taskmgr_module
taskmgr_module.placeholder_context = None
ctx1 = TaskContext.placeholder()
ctx2 = TaskContext.placeholder()
assert ctx1 is ctx2
def test_metadata_is_mutable_dict(self):
"""Test that metadata is a mutable dict."""
TaskContext, _, _ = get_taskmgr_classes()
ctx = TaskContext()
ctx.metadata['count'] = 5
ctx.metadata['items'] = ['a', 'b', 'c']
assert ctx.metadata['count'] == 5
assert len(ctx.metadata['items']) == 3
class TestTaskWrapper:
"""Tests for TaskWrapper class."""
@pytest.mark.asyncio
async def test_id_auto_increment(self):
"""Test that task IDs auto-increment."""
TaskContext, TaskWrapper, _ = get_taskmgr_classes()
# Reset ID index
TaskWrapper._id_index = 0
mock_app = create_mock_app()
async def dummy_coro():
await asyncio.sleep(0.01)
return 'done'
wrapper1 = TaskWrapper(mock_app, dummy_coro())
wrapper2 = TaskWrapper(mock_app, dummy_coro())
assert wrapper1.id == 0
assert wrapper2.id == 1
# Clean up
wrapper1.cancel()
wrapper2.cancel()
@pytest.mark.asyncio
async def test_default_task_type_and_kind(self):
"""Test default task_type and kind values."""
_, TaskWrapper, _ = get_taskmgr_classes()
mock_app = create_mock_app()
async def dummy_coro():
return 'done'
wrapper = TaskWrapper(mock_app, dummy_coro())
assert wrapper.task_type == 'system'
assert wrapper.kind == 'system_task'
wrapper.cancel()
@pytest.mark.asyncio
async def test_to_dict_serialization(self):
"""Test TaskWrapper.to_dict serialization."""
_, TaskWrapper, _ = get_taskmgr_classes()
mock_app = create_mock_app()
async def immediate_coro():
return 'result'
wrapper = TaskWrapper(
mock_app,
immediate_coro(),
name='test_task',
label='Test Task',
)
# Wait for task to complete
await wrapper.task
result = wrapper.to_dict()
assert result['name'] == 'test_task'
assert result['label'] == 'Test Task'
assert result['task_type'] == 'system'
assert result['runtime']['done'] == True
assert result['runtime']['result'] == 'result'
@pytest.mark.asyncio
async def test_to_dict_with_exception(self):
"""Test TaskWrapper.to_dict when task has exception."""
_, TaskWrapper, _ = get_taskmgr_classes()
mock_app = create_mock_app()
async def failing_coro():
raise ValueError('Test error')
wrapper = TaskWrapper(mock_app, failing_coro())
# Wait for task to complete
try:
await wrapper.task
except ValueError:
pass
result = wrapper.to_dict()
assert result['runtime']['done'] == True
assert result['runtime']['exception'] == 'Test error'
assert 'exception_traceback' in result['runtime']
@pytest.mark.asyncio
async def test_cancel_task(self):
"""Test cancel method cancels the asyncio task."""
_, TaskWrapper, _ = get_taskmgr_classes()
mock_app = create_mock_app()
async def long_coro():
await asyncio.sleep(10)
return 'done'
wrapper = TaskWrapper(mock_app, long_coro())
# Task should be running
assert not wrapper.task.done()
wrapper.cancel()
# Give it a moment to be cancelled
await asyncio.sleep(0.01)
assert wrapper.task.done()
assert wrapper.task.cancelled()
class TestAsyncTaskManager:
"""Tests for AsyncTaskManager class."""
@pytest.mark.asyncio
async def test_create_task_adds_to_list(self):
"""Test that create_task adds task to tasks list."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def dummy_coro():
await asyncio.sleep(0.01)
return 'done'
wrapper = manager.create_task(dummy_coro())
assert wrapper in manager.tasks
assert len(manager.tasks) == 1
wrapper.cancel()
@pytest.mark.asyncio
async def test_get_stats_counts_correctly(self):
"""Test get_stats returns correct counts."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def immediate_coro():
return 'done'
async def delayed_coro():
await asyncio.sleep(0.1)
return 'done'
# Create tasks
w1 = manager.create_task(immediate_coro())
w2 = manager.create_task(delayed_coro())
# Wait for first to complete
await w1.task
stats = manager.get_stats()
assert stats['total'] == 2
assert stats['completed'] == 1
assert stats['running'] == 1
w2.cancel()
@pytest.mark.asyncio
async def test_get_tasks_dict_filters_by_type(self):
"""Test get_tasks_dict filters by type."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def dummy_coro():
await asyncio.sleep(0.01)
# Create system and user tasks
w1 = manager.create_task(dummy_coro(), task_type='system')
w2 = manager.create_task(dummy_coro(), task_type='user')
w3 = manager.create_task(dummy_coro(), task_type='user')
result = manager.get_tasks_dict(type='user')
assert len(result['tasks']) == 2
for t in result['tasks']:
assert t['task_type'] == 'user'
w1.cancel()
w2.cancel()
w3.cancel()
@pytest.mark.asyncio
async def test_cancel_by_scope(self):
"""Test cancel_by_scope cancels matching tasks."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def long_coro():
await asyncio.sleep(10)
# Create task with APPLICATION scope
w1 = manager.create_task(long_coro(), scopes=[MockLifecycleControlScope.APPLICATION])
# Create task with different scope
w2 = manager.create_task(long_coro(), scopes=[MockLifecycleControlScope.PIPELINE])
manager.cancel_by_scope(MockLifecycleControlScope.APPLICATION)
await asyncio.sleep(0.01)
assert w1.task.cancelled() or w1.task.done()
assert not w2.task.done()
w2.cancel()
@pytest.mark.asyncio
async def test_cancel_task_by_id(self):
"""Test cancel_task cancels specific task by ID."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def long_coro():
await asyncio.sleep(10)
w1 = manager.create_task(long_coro())
w2 = manager.create_task(long_coro())
manager.cancel_task(w1.id)
await asyncio.sleep(0.01)
assert w1.task.done()
assert not w2.task.done()
w2.cancel()
@pytest.mark.asyncio
async def test_create_user_task_sets_user_type(self):
"""Test create_user_task sets task_type to 'user'."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def dummy_coro():
await asyncio.sleep(0.01)
wrapper = manager.create_user_task(dummy_coro())
assert wrapper.task_type == 'user'
wrapper.cancel()
@pytest.mark.asyncio
async def test_get_task_by_id(self):
"""Test get_task_by_id returns correct task."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
async def dummy_coro():
await asyncio.sleep(0.01)
w1 = manager.create_task(dummy_coro())
w2 = manager.create_task(dummy_coro())
found = manager.get_task_by_id(w1.id)
assert found is w1
not_found = manager.get_task_by_id(9999)
assert not_found is None
w1.cancel()
w2.cancel()
+191
View File
@@ -0,0 +1,191 @@
"""
Unit tests for discover engine utilities.
Tests I18nString, Metadata, and Component utilities.
"""
from __future__ import annotations
from langbot.pkg.discover.engine import I18nString, Metadata, Component
class TestI18nString:
"""Tests for I18nString Pydantic model."""
def test_create_with_english_only(self):
"""Create I18nString with only English."""
i18n = I18nString(en_US='Hello')
assert i18n.en_US == 'Hello'
assert i18n.zh_Hans is None
def test_create_with_multiple_languages(self):
"""Create I18nString with multiple languages."""
i18n = I18nString(
en_US='Hello',
zh_Hans='你好',
zh_Hant='你好',
ja_JP='こんにちは',
)
assert i18n.en_US == 'Hello'
assert i18n.zh_Hans == '你好'
assert i18n.zh_Hant == '你好'
assert i18n.ja_JP == 'こんにちは'
def test_to_dict_with_english_only(self):
"""to_dict returns only non-None fields."""
i18n = I18nString(en_US='Hello')
result = i18n.to_dict()
assert result == {'en_US': 'Hello'}
def test_to_dict_with_multiple_languages(self):
"""to_dict returns all non-None fields."""
i18n = I18nString(
en_US='Hello',
zh_Hans='你好',
)
result = i18n.to_dict()
assert result == {'en_US': 'Hello', 'zh_Hans': '你好'}
def test_to_dict_excludes_none(self):
"""to_dict excludes None values."""
i18n = I18nString(
en_US='Hello',
zh_Hans=None,
ja_JP='こんにちは',
)
result = i18n.to_dict()
assert 'zh_Hans' not in result
assert 'en_US' in result
assert 'ja_JP' in result
def test_to_dict_all_languages(self):
"""to_dict with all supported languages."""
i18n = I18nString(
en_US='Hello',
zh_Hans='你好',
zh_Hant='你好',
ja_JP='こんにちは',
th_TH='สวัสดี',
vi_VN='Xin chào',
es_ES='Hola',
)
result = i18n.to_dict()
assert len(result) == 7
class TestMetadata:
"""Tests for Metadata Pydantic model."""
def test_create_minimal(self):
"""Create Metadata with required fields only."""
from langbot.pkg.discover.engine import I18nString
metadata = Metadata(
name='test-component',
label=I18nString(en_US='Test Component'),
)
assert metadata.name == 'test-component'
assert metadata.label.en_US == 'Test Component'
def test_create_with_all_fields(self):
"""Create Metadata with all optional fields."""
from langbot.pkg.discover.engine import I18nString
metadata = Metadata(
name='test-component',
label=I18nString(en_US='Test'),
description=I18nString(en_US='A test component'),
version='1.0.0',
icon='test-icon',
author='Test Author',
repository='https://github.com/test/repo',
)
assert metadata.version == '1.0.0'
assert metadata.icon == 'test-icon'
assert metadata.author == 'Test Author'
class TestComponentManifest:
"""Tests for Component manifest detection."""
def test_is_component_manifest_valid(self):
"""is_component_manifest returns True for valid manifest."""
manifest = {
'apiVersion': 'v1',
'kind': 'Component',
'metadata': {'name': 'test'},
'spec': {},
}
assert Component.is_component_manifest(manifest) is True
def test_is_component_manifest_missing_apiversion(self):
"""is_component_manifest returns False without apiVersion."""
manifest = {
'kind': 'Component',
'metadata': {'name': 'test'},
'spec': {},
}
assert Component.is_component_manifest(manifest) is False
def test_is_component_manifest_missing_kind(self):
"""is_component_manifest returns False without kind."""
manifest = {
'apiVersion': 'v1',
'metadata': {'name': 'test'},
'spec': {},
}
assert Component.is_component_manifest(manifest) is False
def test_is_component_manifest_missing_metadata(self):
"""is_component_manifest returns False without metadata."""
manifest = {
'apiVersion': 'v1',
'kind': 'Component',
'spec': {},
}
assert Component.is_component_manifest(manifest) is False
def test_is_component_manifest_missing_spec(self):
"""is_component_manifest returns False without spec."""
manifest = {
'apiVersion': 'v1',
'kind': 'Component',
'metadata': {'name': 'test'},
}
assert Component.is_component_manifest(manifest) is False
def test_is_component_manifest_empty(self):
"""is_component_manifest returns False for empty dict."""
manifest = {}
assert Component.is_component_manifest(manifest) is False
def test_is_component_manifest_extra_fields_ok(self):
"""is_component_manifest accepts extra fields."""
manifest = {
'apiVersion': 'v1',
'kind': 'Component',
'metadata': {'name': 'test'},
'spec': {},
'extraField': 'ignored',
}
assert Component.is_component_manifest(manifest) is True
@@ -0,0 +1,203 @@
"""Unit tests for persistence database decorators.
Tests cover:
- manager_class decorator registration
- Class attribute setting
- preregistered_managers list population
Note: Uses import isolation to break circular import chains.
"""
from __future__ import annotations
import sys
from unittest.mock import Mock, MagicMock
from contextlib import contextmanager
from typing import Generator
@contextmanager
def isolated_database_import() -> Generator[None, None, None]:
"""Context manager to isolate circular imports for database testing."""
# Mock modules that cause circular imports
mock_app = MagicMock()
mock_importutil = MagicMock()
mock_importutil.import_modules_in_pkg = lambda pkg: None
mock_importutil.import_modules_in_pkgs = lambda pkgs: None
mock_mgr = MagicMock()
mocks = {
'langbot.pkg.core.app': mock_app,
'langbot.pkg.utils.importutil': mock_importutil,
'langbot.pkg.persistence.mgr': mock_mgr,
}
# Save original state
saved = {}
for name in mocks:
if name in sys.modules:
saved[name] = sys.modules[name]
# Clear database module to force re-import
database_name = 'langbot.pkg.persistence.database'
if database_name in sys.modules:
saved[database_name] = sys.modules[database_name]
# Also clear databases submodules
for sub in ['sqlite', 'postgresql']:
full_name = f'langbot.pkg.persistence.databases.{sub}'
if full_name in sys.modules:
saved[full_name] = sys.modules[full_name]
try:
# Apply mocks
for name, module in mocks.items():
sys.modules[name] = module
# Clear database and submodules
sys.modules.pop(database_name, None)
for sub in ['sqlite', 'postgresql']:
sys.modules.pop(f'langbot.pkg.persistence.databases.{sub}', None)
yield
finally:
# Restore
for name in mocks:
if name in saved:
sys.modules[name] = saved[name]
else:
sys.modules.pop(name, None)
if database_name in saved:
sys.modules[database_name] = saved[database_name]
else:
sys.modules.pop(database_name, None)
for sub in ['sqlite', 'postgresql']:
full_name = f'langbot.pkg.persistence.databases.{sub}'
if full_name in saved:
sys.modules[full_name] = saved[full_name]
else:
sys.modules.pop(full_name, None)
def get_database_module():
"""Get database module with import isolation."""
with isolated_database_import():
from langbot.pkg.persistence import database
return database
class TestManagerClassDecorator:
"""Tests for manager_class decorator."""
def test_decorator_sets_name_attribute(self):
"""Test that decorator sets the 'name' attribute on class."""
database = get_database_module()
# Clear preregistered_managers for this test
database.preregistered_managers.clear()
@database.manager_class('test_db')
class TestManager(database.BaseDatabaseManager):
async def initialize(self):
pass
assert TestManager.name == 'test_db'
def test_decorator_adds_to_preregistered_list(self):
"""Test that decorator adds class to preregistered_managers."""
database = get_database_module()
# Clear preregistered_managers for this test
database.preregistered_managers.clear()
@database.manager_class('test_db2')
class TestManager2(database.BaseDatabaseManager):
async def initialize(self):
pass
assert len(database.preregistered_managers) == 1
assert database.preregistered_managers[0] == TestManager2
def test_decorator_returns_original_class(self):
"""Test that decorator returns the same class."""
database = get_database_module()
database.preregistered_managers.clear()
class OriginalClass(database.BaseDatabaseManager):
async def initialize(self):
pass
decorated = database.manager_class('test_db3')(OriginalClass)
assert decorated is OriginalClass
def test_multiple_decorators_register_separately(self):
"""Test that multiple decorated classes register separately."""
database = get_database_module()
database.preregistered_managers.clear()
@database.manager_class('db_a')
class ManagerA(database.BaseDatabaseManager):
async def initialize(self):
pass
@database.manager_class('db_b')
class ManagerB(database.BaseDatabaseManager):
async def initialize(self):
pass
assert len(database.preregistered_managers) == 2
assert database.preregistered_managers[0].name == 'db_a'
assert database.preregistered_managers[1].name == 'db_b'
def test_base_database_manager_has_name_annotation(self):
"""Test that BaseDatabaseManager has name as class annotation."""
database = get_database_module()
# BaseDatabaseManager has name annotation (type hint)
# Check __annotations__ for the type hint
assert 'name' in database.BaseDatabaseManager.__annotations__
def test_decorated_class_inherits_from_base(self):
"""Test that decorated class properly inherits BaseDatabaseManager."""
database = get_database_module()
database.preregistered_managers.clear()
@database.manager_class('test_inherit')
class TestChild(database.BaseDatabaseManager):
async def initialize(self):
pass
assert issubclass(TestChild, database.BaseDatabaseManager)
# Has abstract method requirement satisfied
assert hasattr(TestChild, 'initialize')
def test_decorator_preserves_class_methods(self):
"""Test that decorator preserves existing class methods."""
database = get_database_module()
database.preregistered_managers.clear()
@database.manager_class('preserve_test')
class ManagerWithMethods(database.BaseDatabaseManager):
custom_attr = 'test_value'
async def initialize(self):
pass
def custom_method(self):
return self.custom_attr
assert ManagerWithMethods.custom_attr == 'test_value'
# Create instance to test method (with mock app)
mock_app = Mock()
instance = ManagerWithMethods(mock_app)
assert instance.custom_method() == 'test_value'
@@ -0,0 +1,156 @@
"""Unit tests for persistence manager methods.
Tests cover:
- execute_async() with mock database
- get_db_engine() with mock database manager
"""
from __future__ import annotations
import pytest
from unittest.mock import Mock, AsyncMock, MagicMock
from importlib import import_module
import sqlalchemy
def get_persistence_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.persistence.mgr')
class TestExecuteAsync:
"""Tests for execute_async method."""
@pytest.mark.asyncio
async def test_execute_async_calls_engine_execute(self):
"""Test that execute_async calls engine execute."""
persistence = get_persistence_module()
mock_app = Mock()
mock_app.persistence_mgr = None
mgr = persistence.PersistenceManager(mock_app)
# Mock database manager with async engine
mock_engine = MagicMock()
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock(return_value=Mock())
mock_conn.commit = AsyncMock()
# Setup the async context manager
async_cm = AsyncMock()
async_cm.__aenter__ = AsyncMock(return_value=mock_conn)
async_cm.__aexit__ = AsyncMock(return_value=None)
mock_engine.connect = Mock(return_value=async_cm)
mock_db = Mock()
mock_db.get_engine = Mock(return_value=mock_engine)
mgr.db = mock_db
# Execute a simple select
await mgr.execute_async(sqlalchemy.select(1))
mock_conn.execute.assert_called_once()
mock_conn.commit.assert_called_once()
@pytest.mark.asyncio
async def test_execute_async_returns_result(self):
"""Test that execute_async returns the result from execute.
NOTE: This test verifies the return value chain - that the result
from conn.execute() is properly returned by execute_async().
The mock verifies the value propagation, not the SQL execution.
For real SQL execution tests, see integration tests.
"""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
# Create a mock result with actual attributes to simulate real result
mock_result = Mock(name='query_result')
mock_result.scalar = Mock(return_value=1) # Simulate scalar() method
mock_result.scalars = Mock() # Simulate scalars() method
mock_engine = MagicMock()
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock(return_value=mock_result)
mock_conn.commit = AsyncMock()
async_cm = AsyncMock()
async_cm.__aenter__ = AsyncMock(return_value=mock_conn)
async_cm.__aexit__ = AsyncMock(return_value=None)
mock_engine.connect = Mock(return_value=async_cm)
mock_db = Mock()
mock_db.get_engine = Mock(return_value=mock_engine)
mgr.db = mock_db
result = await mgr.execute_async(sqlalchemy.text('SELECT 1'))
# Verify result is the same object returned by execute
assert result is mock_result
# Verify result has expected methods (simulating real Result object)
assert hasattr(result, 'scalar')
assert result.scalar() == 1
class TestGetDbEngine:
"""Tests for get_db_engine method."""
def test_get_db_engine_returns_engine_from_db_manager(self):
"""Test that get_db_engine returns engine from db manager."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
mock_engine = Mock(name='engine')
mock_db = Mock()
mock_db.get_engine = Mock(return_value=mock_engine)
mgr.db = mock_db
engine = mgr.get_db_engine()
assert engine == mock_engine
mock_db.get_engine.assert_called_once()
def test_get_db_engine_without_db_set_raises(self):
"""Test that get_db_engine raises when db is not set."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
# db is not initialized
mgr.db = None
with pytest.raises(AttributeError):
mgr.get_db_engine()
class TestSerializeModelEdgeCases:
"""Tests for serialize_model edge cases."""
def test_serialize_model_with_all_columns_masked(self):
"""Test serialize_model when all columns are masked."""
persistence = get_persistence_module()
from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import declarative_base
Base = declarative_base()
class SimpleModel(Base):
__tablename__ = 'simple'
id = Column(Integer, primary_key=True)
name = Column(String(50))
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
instance = SimpleModel(id=1, name='test')
result = mgr.serialize_model(SimpleModel, instance, masked_columns=['id', 'name'])
# Result should be empty dict when all columns masked
assert result == {}
@@ -0,0 +1,129 @@
"""Unit tests for persistence serialize_model function.
Tests cover:
- serialize_model() with various column types
- datetime conversion to isoformat
- masked_columns exclusion
"""
from __future__ import annotations
import datetime
from unittest.mock import Mock
from sqlalchemy import Column, Integer, String, DateTime
from sqlalchemy.orm import declarative_base
from importlib import import_module
def get_persistence_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.persistence.mgr')
# Create a simple test model
Base = declarative_base()
class TestModel(Base):
__tablename__ = 'test_model'
id = Column(Integer, primary_key=True)
name = Column(String(50))
created_at = Column(DateTime)
updated_at = Column(DateTime, nullable=True)
class TestSerializeModel:
"""Tests for serialize_model method."""
def test_serialize_string_and_int_columns(self):
"""Test that string and int columns are serialized directly."""
persistence = get_persistence_module()
# Create a mock persistence manager
mock_app = Mock()
mock_app.persistence_mgr = None
mgr = persistence.PersistenceManager(mock_app)
# Create test model instance
instance = TestModel(id=1, name='test_name', created_at=datetime.datetime(2024, 1, 15, 10, 30, 0))
result = mgr.serialize_model(TestModel, instance)
assert result['id'] == 1
assert result['name'] == 'test_name'
def test_serialize_datetime_to_isoformat(self):
"""Test that datetime columns are converted to isoformat string."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
dt = datetime.datetime(2024, 1, 15, 10, 30, 45)
instance = TestModel(id=1, name='test', created_at=dt)
result = mgr.serialize_model(TestModel, instance)
assert result['created_at'] == '2024-01-15T10:30:45'
assert isinstance(result['created_at'], str)
def test_serialize_datetime_with_timezone(self):
"""Test datetime with timezone conversion."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
# datetime with timezone
dt = datetime.datetime(2024, 1, 15, 10, 30, 45, tzinfo=datetime.timezone.utc)
instance = TestModel(id=1, name='test', created_at=dt)
result = mgr.serialize_model(TestModel, instance)
assert '2024-01-15' in result['created_at']
assert isinstance(result['created_at'], str)
def test_serialize_none_datetime(self):
"""Test that None datetime column is serialized as None."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
instance = TestModel(id=1, name='test', created_at=datetime.datetime.now(), updated_at=None)
result = mgr.serialize_model(TestModel, instance)
# None datetime should be None (not converted to isoformat)
assert result['updated_at'] is None
def test_masked_columns_excluded(self):
"""Test that masked columns are excluded from output."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
instance = TestModel(id=1, name='secret_name', created_at=datetime.datetime.now())
result = mgr.serialize_model(TestModel, instance, masked_columns=['name'])
assert 'id' in result
assert 'created_at' in result
assert 'name' not in result
def test_masked_columns_multiple(self):
"""Test that multiple masked columns are excluded."""
persistence = get_persistence_module()
mock_app = Mock()
mgr = persistence.PersistenceManager(mock_app)
instance = TestModel(id=1, name='secret', created_at=datetime.datetime.now())
result = mgr.serialize_model(TestModel, instance, masked_columns=['id', 'name'])
assert 'id' not in result
assert 'name' not in result
assert 'created_at' in result
+265
View File
@@ -0,0 +1,265 @@
"""
Shared test fixtures and configuration
This file provides infrastructure for all pipeline tests, including:
- Mock object factories
- Test fixtures
- Common test helper functions
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
# Preload pipelinemgr so the pipeline.stage module is fully initialised before
# any individual stage test (e.g. preproc, longtext) tries to import it. Without
# this, running a stage test in isolation triggers a circular-import error:
# stage.py → core.app → pipelinemgr → stage.stage_class (not yet bound).
import langbot.pkg.pipeline.pipelinemgr # noqa: F401
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.pipeline import entities as pipeline_entities
class MockApplication:
"""Mock Application object providing all basic dependencies needed by stages"""
def __init__(self):
self.logger = self._create_mock_logger()
self.sess_mgr = self._create_mock_session_manager()
self.model_mgr = self._create_mock_model_manager()
self.tool_mgr = self._create_mock_tool_manager()
self.plugin_connector = self._create_mock_plugin_connector()
self.persistence_mgr = self._create_mock_persistence_manager()
self.query_pool = self._create_mock_query_pool()
self.instance_config = self._create_mock_instance_config()
self.task_mgr = self._create_mock_task_manager()
# Skill manager is optional; PreProcessor only touches it for the
# local-agent runner. None keeps the skill-binding branch inert.
self.skill_mgr = None
def _create_mock_logger(self):
logger = Mock()
logger.debug = Mock()
logger.info = Mock()
logger.error = Mock()
logger.warning = Mock()
return logger
def _create_mock_session_manager(self):
sess_mgr = AsyncMock()
sess_mgr.get_session = AsyncMock()
sess_mgr.get_conversation = AsyncMock()
return sess_mgr
def _create_mock_model_manager(self):
model_mgr = AsyncMock()
model_mgr.get_model_by_uuid = AsyncMock()
return model_mgr
def _create_mock_tool_manager(self):
tool_mgr = AsyncMock()
tool_mgr.get_all_tools = AsyncMock(return_value=[])
return tool_mgr
def _create_mock_plugin_connector(self):
plugin_connector = AsyncMock()
plugin_connector.emit_event = AsyncMock()
return plugin_connector
def _create_mock_persistence_manager(self):
persistence_mgr = AsyncMock()
persistence_mgr.execute_async = AsyncMock()
return persistence_mgr
def _create_mock_query_pool(self):
query_pool = Mock()
query_pool.cached_queries = {}
query_pool.queries = []
query_pool.condition = AsyncMock()
return query_pool
def _create_mock_instance_config(self):
instance_config = Mock()
instance_config.data = {
'command': {'prefix': ['/', '!'], 'enable': True},
'concurrency': {'pipeline': 10},
}
return instance_config
def _create_mock_task_manager(self):
task_mgr = Mock()
task_mgr.create_task = Mock()
return task_mgr
@pytest.fixture
def mock_app():
"""Provides Mock Application instance"""
return MockApplication()
@pytest.fixture
def mock_session():
"""Provides Mock Session object"""
session = Mock()
session.launcher_type = provider_session.LauncherTypes.PERSON
session.launcher_id = 12345
session._semaphore = AsyncMock()
session._semaphore.locked = Mock(return_value=False)
session._semaphore.acquire = AsyncMock()
session._semaphore.release = AsyncMock()
return session
@pytest.fixture
def mock_conversation():
"""Provides Mock Conversation object"""
conversation = Mock()
conversation.uuid = 'test-conversation-uuid'
# Create mock prompt with copy method
mock_prompt = Mock()
mock_prompt.messages = []
mock_prompt.copy = Mock(return_value=Mock(messages=[]))
conversation.prompt = mock_prompt
# Create mock messages list with copy method
mock_messages = Mock()
mock_messages.copy = Mock(return_value=[])
conversation.messages = mock_messages
return conversation
@pytest.fixture
def mock_model():
"""Provides Mock Model object"""
model = Mock()
model.model_entity = Mock()
model.model_entity.uuid = 'test-model-uuid'
model.model_entity.abilities = ['func_call', 'vision']
return model
@pytest.fixture
def mock_adapter():
"""Provides Mock Adapter object"""
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
adapter.reply_message = AsyncMock()
adapter.reply_message_chunk = AsyncMock()
return adapter
@pytest.fixture
def sample_message_chain():
"""Provides sample message chain"""
return platform_message.MessageChain(
[
platform_message.Plain(text='Hello, this is a test message'),
]
)
@pytest.fixture
def sample_message_event(sample_message_chain):
"""Provides sample message event (FriendMessage)"""
sender = platform_entities.Friend(
id=12345,
nickname='TestUser',
remark=None,
)
return platform_events.FriendMessage(
type='FriendMessage',
sender=sender,
message_chain=sample_message_chain,
time=1609459200, # 2021-01-01 00:00:00
)
@pytest.fixture
def sample_query(sample_message_chain, sample_message_event, mock_adapter):
"""Provides sample Query object - using model_construct to bypass validation"""
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
# Use model_construct to bypass Pydantic validation for test purposes
query = pipeline_query.Query.model_construct(
query_id='test-query-id',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=12345,
message_chain=sample_message_chain,
message_event=sample_message_event,
adapter=mock_adapter,
pipeline_uuid='test-pipeline-uuid',
bot_uuid='test-bot-uuid',
pipeline_config={
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {'model': {'primary': 'test-model-uuid', 'fallbacks': []}, 'prompt': 'test-prompt'},
},
'output': {'misc': {'at-sender': False, 'quote-origin': False}},
'trigger': {'misc': {'combine-quote-message': False}},
},
session=None,
prompt=None,
messages=[],
user_message=None,
use_funcs=[],
use_llm_model_uuid=None,
variables={},
resp_messages=[],
resp_message_chain=None,
current_stage_name=None,
)
return query
@pytest.fixture
def sample_pipeline_config():
"""Provides sample pipeline configuration"""
return {
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {'model': {'primary': 'test-model-uuid', 'fallbacks': []}, 'prompt': 'test-prompt'},
},
'output': {'misc': {'at-sender': False, 'quote-origin': False}},
'trigger': {'misc': {'combine-quote-message': False}},
'ratelimit': {'enable': True, 'algo': 'fixwin', 'window': 60, 'limit': 10},
}
def create_stage_result(
result_type: pipeline_entities.ResultType,
query: pipeline_query.Query,
user_notice: str = '',
console_notice: str = '',
debug_notice: str = '',
error_notice: str = '',
) -> pipeline_entities.StageProcessResult:
"""Helper function to create stage process result"""
return pipeline_entities.StageProcessResult(
result_type=result_type,
new_query=query,
user_notice=user_notice,
console_notice=console_notice,
debug_notice=debug_notice,
error_notice=error_notice,
)
def assert_result_continue(result: pipeline_entities.StageProcessResult):
"""Assert result is CONTINUE type"""
assert result.result_type == pipeline_entities.ResultType.CONTINUE
def assert_result_interrupt(result: pipeline_entities.StageProcessResult):
"""Assert result is INTERRUPT type"""
assert result.result_type == pipeline_entities.ResultType.INTERRUPT
@@ -0,0 +1,637 @@
"""
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
def test_merge_messages_preserves_routed_by_rule_if_any_input_matches(self):
"""Merged PendingMessage should keep routed_by_rule when any input was rule-routed."""
aggregator = get_aggregator_module()
app = make_aggregator_app()
agg = aggregator.MessageAggregator(app)
chain1 = text_chain('first')
chain2 = text_chain('second')
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='test-pipeline-uuid',
routed_by_rule=False,
)
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='test-pipeline-uuid',
routed_by_rule=True,
)
merged = agg._merge_messages([pending1, pending2])
assert merged.routed_by_rule is True
assert str(merged.message_chain) == 'first\nsecond'
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
+138
View File
@@ -0,0 +1,138 @@
"""
BanSessionCheckStage unit tests
Tests the actual BanSessionCheckStage implementation from pkg.pipeline.bansess
"""
import pytest
from importlib import import_module
import langbot_plugin.api.entities.builtin.provider.session as provider_session
def get_modules():
"""Lazy import to ensure proper initialization order"""
# Import pipelinemgr first to trigger proper stage registration
bansess = import_module('langbot.pkg.pipeline.bansess.bansess')
entities = import_module('langbot.pkg.pipeline.entities')
return bansess, entities
@pytest.mark.asyncio
async def test_whitelist_allow(mock_app, sample_query):
"""Test whitelist allows matching session"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'whitelist', 'whitelist': ['person_12345']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query == sample_query
@pytest.mark.asyncio
async def test_whitelist_deny(mock_app, sample_query):
"""Test whitelist denies non-matching session"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '99999'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'whitelist', 'whitelist': ['person_12345']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.INTERRUPT
@pytest.mark.asyncio
async def test_blacklist_allow(mock_app, sample_query):
"""Test blacklist allows non-matching session"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'blacklist', 'blacklist': ['person_99999']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
@pytest.mark.asyncio
async def test_blacklist_deny(mock_app, sample_query):
"""Test blacklist denies matching session"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'blacklist', 'blacklist': ['person_12345']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.INTERRUPT
@pytest.mark.asyncio
async def test_wildcard_group(mock_app, sample_query):
"""Test group wildcard matching"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.GROUP
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'whitelist', 'whitelist': ['group_*']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
@pytest.mark.asyncio
async def test_wildcard_person(mock_app, sample_query):
"""Test person wildcard matching"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'whitelist', 'whitelist': ['person_*']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
@pytest.mark.asyncio
async def test_user_id_wildcard(mock_app, sample_query):
"""Test user ID wildcard matching (*_id format)"""
bansess, entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.sender_id = '67890'
sample_query.pipeline_config = {'trigger': {'access-control': {'mode': 'whitelist', 'whitelist': ['*_67890']}}}
stage = bansess.BanSessionCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'BanSessionCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
@@ -0,0 +1,466 @@
"""
Unit tests for ChatMessageHandler - REAL imports.
Tests the actual ChatMessageHandler class from production code.
Uses tests.utils.import_isolation to break circular import chain safely.
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from tests.factories import FakeApp
# ============== FIXTURE USING IMPORT ISOLATION UTILITY ==============
@pytest.fixture(scope='module')
def mock_circular_import_chain():
"""
Break circular import chain using isolated_sys_modules.
Chain: handler → core.app → pipeline.controller → http_controller → groups/plugins → taskmgr
Uses tests.utils.import_isolation for safe, reversible sys.modules manipulation.
"""
from tests.utils.import_isolation import (
isolated_sys_modules,
make_pipeline_handler_import_mocks,
get_handler_modules_to_clear,
)
from langbot_plugin.api.entities.builtin.provider.message import Message
mocks = make_pipeline_handler_import_mocks()
# Create a default runner that yields a simple response
class DefaultRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
yield Message(role='assistant', content='fake response')
mocks['langbot.pkg.provider.runner'].preregistered_runners = [DefaultRunner]
clear = get_handler_modules_to_clear('chat')
with isolated_sys_modules(mocks=mocks, clear=clear):
yield
@pytest.fixture
def fake_app():
"""Create FakeApp instance."""
return FakeApp()
@pytest.fixture
def mock_event_ctx():
"""Create mock event context."""
ctx = Mock()
ctx.is_prevented_default = Mock(return_value=False)
ctx.event = Mock()
ctx.event.user_message_alter = None
ctx.event.reply_message_chain = None
return ctx
@pytest.fixture
def set_runner():
"""Factory fixture to set a custom runner for tests."""
def _set_runner(runner_class):
import sys
sys.modules['langbot.pkg.provider.runner'].preregistered_runners = [runner_class]
return _set_runner
# ============== CACHED LAZY IMPORTS ==============
_chat_handler_module = None
_entities_module = None
def get_chat_handler():
"""Import ChatMessageHandler after circular import chain is mocked."""
global _chat_handler_module
if _chat_handler_module is None:
from importlib import import_module
_chat_handler_module = import_module('langbot.pkg.pipeline.process.handlers.chat')
return _chat_handler_module
def get_entities():
"""Import pipeline entities - uses real module."""
global _entities_module
if _entities_module is None:
from importlib import import_module
_entities_module = import_module('langbot.pkg.pipeline.entities')
return _entities_module
# ============== REAL ChatMessageHandler Tests ==============
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestChatMessageHandlerReal:
"""Tests for real ChatMessageHandler class."""
@pytest.mark.asyncio
async def test_real_import_works(self):
"""Verify we can import the real handler class."""
chat = get_chat_handler()
assert hasattr(chat, 'ChatMessageHandler')
handler_cls = chat.ChatMessageHandler
assert handler_cls.__name__ == 'ChatMessageHandler'
@pytest.mark.asyncio
async def test_handler_creation(self, fake_app):
"""ChatMessageHandler can be instantiated."""
chat = get_chat_handler()
handler = chat.ChatMessageHandler(fake_app)
assert handler.ap is fake_app
@pytest.mark.asyncio
async def test_prevent_default_without_reply_interrupts(self, fake_app, mock_event_ctx):
"""prevent_default without reply chain yields INTERRUPT."""
from tests.factories import text_query
chat = get_chat_handler()
entities = get_entities()
mock_event_ctx.is_prevented_default.return_value = True
mock_event_ctx.event.reply_message_chain = None
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
handler = chat.ChatMessageHandler(fake_app)
query = text_query('hello')
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) == 1
assert results[0].result_type == entities.ResultType.INTERRUPT
@pytest.mark.asyncio
async def test_prevent_default_with_reply_continues(self, fake_app, mock_event_ctx):
"""prevent_default with reply yields CONTINUE and updates resp_messages."""
from tests.factories import text_query, text_chain
chat = get_chat_handler()
entities = get_entities()
reply_chain = text_chain('plugin reply')
mock_event_ctx.is_prevented_default.return_value = True
mock_event_ctx.event.reply_message_chain = reply_chain
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
handler = chat.ChatMessageHandler(fake_app)
query = text_query('hello')
query.resp_messages = []
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) == 1
assert results[0].result_type == entities.ResultType.CONTINUE
assert len(query.resp_messages) == 1
assert query.resp_messages[0] == reply_chain
@pytest.mark.asyncio
async def test_user_message_alter_string(self, fake_app, mock_event_ctx, set_runner):
"""user_message_alter as string updates query.user_message."""
from tests.factories import text_query
from langbot_plugin.api.entities.builtin.provider.message import Message
chat = get_chat_handler()
mock_event_ctx.is_prevented_default.return_value = False
mock_event_ctx.event.user_message_alter = 'altered text'
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
query = text_query('original')
query.adapter = Mock()
query.adapter.is_stream_output_supported = AsyncMock(return_value=False)
query.user_message = Message(role='user', content=[])
class QuickRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
yield Message(role='assistant', content='ok')
set_runner(QuickRunner)
handler = chat.ChatMessageHandler(fake_app)
results = []
async for result in handler.handle(query):
results.append(result)
assert query.user_message.content is not None
@pytest.mark.asyncio
async def test_adapter_without_stream_method_defaults_non_stream(self, fake_app, mock_event_ctx, set_runner):
"""Adapter without is_stream_output_supported defaults to non-stream."""
from tests.factories import text_query
from langbot_plugin.api.entities.builtin.provider.message import Message, ContentElement
chat = get_chat_handler()
mock_event_ctx.is_prevented_default.return_value = False
mock_event_ctx.event.user_message_alter = None
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
query = text_query('test')
query.adapter = Mock(spec=[])
query.user_message = Message(role='user', content=[ContentElement.from_text('test')])
class SingleRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
yield Message(role='assistant', content='response')
set_runner(SingleRunner)
handler = chat.ChatMessageHandler(fake_app)
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) >= 1
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestChatHandlerStreaming:
"""Tests for streaming behavior."""
@pytest.mark.asyncio
async def test_streaming_chunks_collected(self, fake_app, mock_event_ctx, set_runner):
"""Streaming produces multiple results."""
from tests.factories import text_query
from langbot_plugin.api.entities.builtin.provider.message import Message, ContentElement, MessageChunk
chat = get_chat_handler()
mock_event_ctx.is_prevented_default.return_value = False
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
query = text_query('stream test')
query.adapter = Mock()
query.adapter.is_stream_output_supported = AsyncMock(return_value=True)
query.adapter.create_message_card = AsyncMock()
query.user_message = Message(role='user', content=[ContentElement.from_text('test')])
class StreamRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
yield MessageChunk(role='assistant', content='Hello', is_final=False)
yield MessageChunk(role='assistant', content=' World', is_final=True)
set_runner(StreamRunner)
handler = chat.ChatMessageHandler(fake_app)
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) >= 1
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestChatHandlerExceptions:
"""Tests for exception handling."""
@pytest.mark.asyncio
async def test_runner_exception_yields_interrupt(self, fake_app, mock_event_ctx, set_runner):
"""Runner exception yields INTERRUPT with error notices."""
from tests.factories import text_query
from langbot_plugin.api.entities.builtin.provider.message import Message
chat = get_chat_handler()
entities = get_entities()
mock_event_ctx.is_prevented_default.return_value = False
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
query = text_query('fail test')
query.adapter = Mock()
query.adapter.is_stream_output_supported = AsyncMock(return_value=False)
query.user_message = Message(role='user', content=[])
query.pipeline_config = {
'output': {'misc': {'exception-handling': 'show-hint', 'failure-hint': 'Request failed.'}},
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {'prompt': 'default', 'model': {'primary': 'test'}},
},
}
class FailingRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
raise ValueError('API error')
yield
set_runner(FailingRunner)
handler = chat.ChatMessageHandler(fake_app)
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) == 1
assert results[0].result_type == entities.ResultType.INTERRUPT
assert results[0].user_notice == 'Request failed.'
assert results[0].error_notice is not None
@pytest.mark.asyncio
async def test_exception_show_error_mode(self, fake_app, mock_event_ctx, set_runner):
"""show-error mode shows actual exception."""
from tests.factories import text_query
from langbot_plugin.api.entities.builtin.provider.message import Message
chat = get_chat_handler()
mock_event_ctx.is_prevented_default.return_value = False
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
query = text_query('error test')
query.adapter = Mock()
query.adapter.is_stream_output_supported = AsyncMock(return_value=False)
query.user_message = Message(role='user', content=[])
query.pipeline_config = {
'output': {'misc': {'exception-handling': 'show-error'}},
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {'prompt': 'default', 'model': {'primary': 'test'}},
},
}
class ErrorRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
raise ValueError('Custom error')
yield
set_runner(ErrorRunner)
handler = chat.ChatMessageHandler(fake_app)
results = []
async for result in handler.handle(query):
results.append(result)
assert results[0].user_notice == 'Custom error'
@pytest.mark.asyncio
async def test_exception_hide_mode(self, fake_app, mock_event_ctx, set_runner):
"""hide mode shows no user notice."""
from tests.factories import text_query
from langbot_plugin.api.entities.builtin.provider.message import Message
chat = get_chat_handler()
mock_event_ctx.is_prevented_default.return_value = False
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
query = text_query('hide test')
query.adapter = Mock()
query.adapter.is_stream_output_supported = AsyncMock(return_value=False)
query.user_message = Message(role='user', content=[])
query.pipeline_config = {
'output': {'misc': {'exception-handling': 'hide'}},
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {'prompt': 'default', 'model': {'primary': 'test'}},
},
}
class HideErrorRunner:
name = 'local-agent'
def __init__(self, app, config):
self.app = app
self.config = config
async def run(self, query):
raise RuntimeError('hidden')
yield
set_runner(HideErrorRunner)
handler = chat.ChatMessageHandler(fake_app)
results = []
async for result in handler.handle(query):
results.append(result)
assert results[0].user_notice is None
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestChatHandlerHelper:
"""Tests for helper methods."""
def test_cut_str_short(self, fake_app):
"""cut_str returns short string unchanged."""
chat = get_chat_handler()
handler = chat.ChatMessageHandler(fake_app)
result = handler.cut_str('short text')
assert result == 'short text'
def test_cut_str_long(self, fake_app):
"""cut_str truncates long string."""
chat = get_chat_handler()
handler = chat.ChatMessageHandler(fake_app)
result = handler.cut_str('this is a very long string that exceeds twenty characters')
assert '...' in result
assert len(result) <= 23
def test_cut_str_multiline(self, fake_app):
"""cut_str truncates multiline string."""
chat = get_chat_handler()
handler = chat.ChatMessageHandler(fake_app)
result = handler.cut_str('first line\nsecond line')
assert '...' in result
@@ -0,0 +1,78 @@
from __future__ import annotations
from unittest.mock import Mock
import pytest
import langbot_plugin.api.entities.builtin.provider.message as provider_message
# TODO: unskip once the handler ↔ app circular import is resolved
pytest.skip(
'circular import in handler ↔ app; will be unblocked once resolved',
allow_module_level=True,
)
from langbot.pkg.pipeline.process.handler import MessageHandler # noqa: E402
class _StubHandler(MessageHandler):
async def handle(self, query):
raise NotImplementedError
handler = _StubHandler(ap=Mock())
def test_chat_handler_formats_tool_call_request_log():
result = provider_message.Message(
role='assistant',
content='',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(name='exec', arguments='{}'),
)
],
)
summary = handler.format_result_log(result)
assert summary == 'assistant: requested tools: exec'
def test_chat_handler_formats_tool_result_log():
result = provider_message.Message(
role='tool',
content='{"status":"completed","exit_code":0,"backend":"podman","stdout":"42\\n"}',
tool_call_id='call-1',
)
summary = handler.format_result_log(result)
# Tool results use generic cut_str truncation
assert summary is not None
assert summary.startswith('tool: {"status":"com')
assert summary.endswith('...')
def test_chat_handler_formats_tool_error_log():
result = provider_message.MessageChunk(
role='tool',
content='err: host_path must point to an existing directory on the host',
tool_call_id='call-1',
is_final=True,
)
summary = handler.format_result_log(result)
assert summary is not None
assert summary.startswith('tool error: err: host_path must')
assert summary.endswith('...')
def test_chat_handler_skips_empty_assistant_log():
result = provider_message.Message(role='assistant', content='')
summary = handler.format_result_log(result)
assert summary is None
@@ -0,0 +1,110 @@
from __future__ import annotations
from datetime import datetime, timedelta
from importlib import import_module
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import yaml
def _preproc_module():
# Import pipelinemgr first so pipeline stages are registered without tripping
# the stage <-> core.app circular import during isolated test collection.
import_module('langbot.pkg.pipeline.pipelinemgr')
return import_module('langbot.pkg.pipeline.preproc.preproc')
def _entities_module():
return import_module('langbot.pkg.pipeline.entities')
def _conversation(created_at: datetime, updated_at: datetime | None = None):
prompt = Mock()
prompt.messages = []
prompt.copy = Mock(return_value=Mock(messages=[]))
return SimpleNamespace(
uuid='existing-conversation-uuid',
create_time=created_at,
update_time=updated_at,
prompt=prompt,
messages=[],
)
def _prompt_preprocessing_context(default_prompt=None, prompt=None):
ctx = Mock()
ctx.event.default_prompt = default_prompt or []
ctx.event.prompt = prompt or []
return ctx
async def _run_preprocessor(mock_app, sample_query, conversation):
session = SimpleNamespace(launcher_type=sample_query.launcher_type, launcher_id=sample_query.launcher_id)
mock_app.sess_mgr.get_session = AsyncMock(return_value=session)
mock_app.sess_mgr.get_conversation = AsyncMock(return_value=conversation)
mock_app.plugin_connector.emit_event = AsyncMock(return_value=_prompt_preprocessing_context())
sample_query.pipeline_config = {
'ai': {
'runner': {'runner': 'local-agent', 'expire-time': 60},
'local-agent': {'model': {'primary': '', 'fallbacks': []}, 'prompt': []},
},
'trigger': {'misc': {'combine-quote-message': False}},
'output': {'misc': {'exception-handling': 'show-hint'}},
}
return await _preproc_module().PreProcessor(mock_app).process(sample_query, 'PreProcessor')
@pytest.mark.asyncio
async def test_preprocessor_expires_conversation_from_last_update_time(mock_app, sample_query):
conversation = _conversation(
created_at=datetime.now() - timedelta(seconds=10),
updated_at=datetime.now() - timedelta(seconds=120),
)
result = await _run_preprocessor(mock_app, sample_query, conversation)
assert result.result_type == _entities_module().ResultType.CONTINUE
assert conversation.uuid is None
assert conversation.update_time > datetime.now() - timedelta(seconds=5)
assert result.new_query.variables['conversation_id'] is None
@pytest.mark.asyncio
async def test_preprocessor_keeps_conversation_when_last_update_is_not_expired(mock_app, sample_query):
conversation = _conversation(
created_at=datetime.now() - timedelta(seconds=120),
updated_at=datetime.now() - timedelta(seconds=30),
)
result = await _run_preprocessor(mock_app, sample_query, conversation)
assert result.result_type == _entities_module().ResultType.CONTINUE
assert conversation.uuid == 'existing-conversation-uuid'
assert conversation.update_time > datetime.now() - timedelta(seconds=5)
assert result.new_query.variables['conversation_id'] == 'existing-conversation-uuid'
def test_expire_time_metadata_lives_under_ai_runner_not_safety():
# Use path relative to test file location for portability
# test file: tests/unit_tests/pipeline/test_chat_session_limit.py
# project root: 4 levels up
project_root = Path(__file__).parent.parent.parent.parent
metadata_dir = project_root / 'src' / 'langbot' / 'templates' / 'metadata' / 'pipeline'
ai_meta = yaml.safe_load((metadata_dir / 'ai.yaml').read_text())
safety_meta = yaml.safe_load((metadata_dir / 'safety.yaml').read_text())
ai_stage_names = [stage['name'] for stage in ai_meta['stages']]
assert 'session-limit' not in ai_stage_names
assert 'session-limit' not in [stage['name'] for stage in safety_meta['stages']]
runner_stage = next(stage for stage in ai_meta['stages'] if stage['name'] == 'runner')
expire_time = next(item for item in runner_stage['config'] if item['name'] == 'expire-time')
assert 'Conversation expire time' in expire_time['label']['en_US']
assert 'Session validity' not in expire_time['label']['en_US']
+512
View File
@@ -0,0 +1,512 @@
"""
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 [filter_impl.name for filter_impl in stage.filter_chain] == ['content-ignore']
@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)
assert {filter_impl.name for filter_impl in stage.filter_chain} == {
'ban-word-filter',
'content-ignore',
}
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,403 @@
"""
Unit tests for CommandHandler - REAL imports.
Tests the actual CommandHandler class from production code.
Uses tests.utils.import_isolation to break circular import chain safely.
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from tests.factories import FakeApp, command_query
# ============== FIXTURE USING IMPORT ISOLATION UTILITY ==============
@pytest.fixture(scope='module')
def mock_circular_import_chain():
"""
Break circular import chain using isolated_sys_modules.
Chain: handler → core.app → pipeline.controller → http_controller → groups/plugins → taskmgr
Uses tests.utils.import_isolation for safe, reversible sys.modules manipulation.
"""
from tests.utils.import_isolation import (
isolated_sys_modules,
make_pipeline_handler_import_mocks,
get_handler_modules_to_clear,
)
mocks = make_pipeline_handler_import_mocks()
clear = get_handler_modules_to_clear('command')
with isolated_sys_modules(mocks=mocks, clear=clear):
yield
@pytest.fixture
def fake_app():
"""Create FakeApp instance."""
return FakeApp()
@pytest.fixture
def mock_event_ctx():
"""Create mock event context."""
ctx = Mock()
ctx.is_prevented_default = Mock(return_value=False)
ctx.event = Mock()
ctx.event.reply_message_chain = None
return ctx
@pytest.fixture
def mock_execute_factory():
"""Factory fixture to create mock cmd_mgr.execute generators."""
def _create_execute(
text: str | None = 'ok',
error: str | None = None,
image_url: str | None = None,
image_base64: str | None = None,
file_url: str | None = None,
):
async def mock_execute(command_text, full_command_text, query, session):
ret = Mock()
ret.text = text
ret.error = error
ret.image_url = image_url
ret.image_base64 = image_base64
ret.file_url = file_url
yield ret
return mock_execute
return _create_execute
# ============== CACHED LAZY IMPORTS ==============
_command_handler_module = None
_entities_module = None
def get_command_handler():
"""Import CommandHandler after circular import chain is mocked."""
global _command_handler_module
if _command_handler_module is None:
from importlib import import_module
_command_handler_module = import_module('langbot.pkg.pipeline.process.handlers.command')
return _command_handler_module
def get_entities():
"""Import pipeline entities - uses real module."""
global _entities_module
if _entities_module is None:
from importlib import import_module
_entities_module = import_module('langbot.pkg.pipeline.entities')
return _entities_module
# ============== REAL CommandHandler Tests ==============
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestCommandHandlerReal:
"""Tests for real CommandHandler class."""
@pytest.mark.asyncio
async def test_real_import_works(self):
"""Verify we can import the real handler class."""
command = get_command_handler()
assert hasattr(command, 'CommandHandler')
handler_cls = command.CommandHandler
assert handler_cls.__name__ == 'CommandHandler'
@pytest.mark.asyncio
async def test_handler_creation(self, fake_app):
"""CommandHandler can be instantiated."""
command = get_command_handler()
handler = command.CommandHandler(fake_app)
assert handler.ap is fake_app
@pytest.mark.asyncio
async def test_command_parsing_extracts_command_name(self, fake_app, mock_event_ctx):
"""Command text is extracted after prefix."""
command = get_command_handler()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
executed_commands = []
async def track_execute(command_text, full_command_text, query, session):
executed_commands.append(command_text)
ret = Mock()
ret.text = 'ok'
ret.error = None
ret.image_url = None
ret.image_base64 = None
ret.file_url = None
yield ret
fake_app.cmd_mgr.execute = track_execute
handler = command.CommandHandler(fake_app)
query = command_query('help arg1 arg2')
results = []
async for result in handler.handle(query):
results.append(result)
assert executed_commands[0] == 'help arg1 arg2'
@pytest.mark.asyncio
async def test_admin_privilege_check(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Admin users get privilege level 2."""
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
command = get_command_handler()
fake_app.instance_config.data = {'admins': ['person_12345']}
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory()
handler = command.CommandHandler(fake_app)
query = command_query('status')
query.launcher_type = LauncherTypes.PERSON
query.launcher_id = 12345
results = []
async for result in handler.handle(query):
results.append(result)
call_args = fake_app.plugin_connector.emit_event.call_args
event = call_args[0][0]
assert event.is_admin is True
@pytest.mark.asyncio
async def test_non_admin_privilege_check(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Non-admin users get privilege level 1."""
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
command = get_command_handler()
fake_app.instance_config.data = {'admins': ['person_12345']}
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory()
handler = command.CommandHandler(fake_app)
query = command_query('status')
query.launcher_type = LauncherTypes.PERSON
query.launcher_id = 67890
results = []
async for result in handler.handle(query):
results.append(result)
call_args = fake_app.plugin_connector.emit_event.call_args
event = call_args[0][0]
assert event.is_admin is False
@pytest.mark.asyncio
async def test_prevent_default_with_reply_continues(self, fake_app, mock_event_ctx):
"""prevent_default with reply yields CONTINUE."""
from tests.factories.message import text_chain
command = get_command_handler()
entities = get_entities()
reply_chain = text_chain('plugin reply')
mock_event_ctx.is_prevented_default.return_value = True
mock_event_ctx.event.reply_message_chain = reply_chain
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
handler = command.CommandHandler(fake_app)
query = command_query('test')
query.resp_messages = []
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) == 1
assert results[0].result_type == entities.ResultType.CONTINUE
assert len(query.resp_messages) == 1
assert query.resp_messages[0] == reply_chain
@pytest.mark.asyncio
async def test_prevent_default_without_reply_interrupts(self, fake_app, mock_event_ctx):
"""prevent_default without reply yields INTERRUPT."""
command = get_command_handler()
entities = get_entities()
mock_event_ctx.is_prevented_default.return_value = True
mock_event_ctx.event.reply_message_chain = None
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
handler = command.CommandHandler(fake_app)
query = command_query('test')
results = []
async for result in handler.handle(query):
results.append(result)
assert len(results) == 1
assert results[0].result_type == entities.ResultType.INTERRUPT
@pytest.mark.asyncio
async def test_event_type_person_command(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Person launcher creates PersonCommandSent event."""
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
from langbot_plugin.api.entities import events
command = get_command_handler()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory()
handler = command.CommandHandler(fake_app)
query = command_query('help')
query.launcher_type = LauncherTypes.PERSON
results = []
async for result in handler.handle(query):
results.append(result)
call_args = fake_app.plugin_connector.emit_event.call_args
event = call_args[0][0]
assert isinstance(event, events.PersonCommandSent)
@pytest.mark.asyncio
async def test_event_type_group_command(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Group launcher creates GroupCommandSent event."""
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
from langbot_plugin.api.entities import events
command = get_command_handler()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory()
handler = command.CommandHandler(fake_app)
query = command_query('help')
query.launcher_type = LauncherTypes.GROUP
results = []
async for result in handler.handle(query):
results.append(result)
call_args = fake_app.plugin_connector.emit_event.call_args
event = call_args[0][0]
assert isinstance(event, events.GroupCommandSent)
@pytest.mark.asyncio
async def test_command_result_text(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Text result is added to resp_messages."""
command = get_command_handler()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory(text='Command output')
handler = command.CommandHandler(fake_app)
query = command_query('echo')
query.resp_messages = []
results = []
async for result in handler.handle(query):
results.append(result)
assert len(query.resp_messages) == 1
msg = query.resp_messages[0]
assert msg.role == 'command'
assert len(msg.content) == 1
assert msg.content[0].type == 'text'
assert msg.content[0].text == 'Command output'
@pytest.mark.asyncio
async def test_command_result_error(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Error result creates error message."""
command = get_command_handler()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory(text=None, error='Command failed')
handler = command.CommandHandler(fake_app)
query = command_query('fail')
query.resp_messages = []
results = []
async for result in handler.handle(query):
results.append(result)
assert len(query.resp_messages) == 1
msg = query.resp_messages[0]
assert msg.role == 'command'
assert msg.content == 'Command failed'
@pytest.mark.asyncio
async def test_command_result_image_url(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Image URL result is added to content."""
command = get_command_handler()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory(
text='Here is the image:', image_url='https://example.com/image.png'
)
handler = command.CommandHandler(fake_app)
query = command_query('image')
query.resp_messages = []
results = []
async for result in handler.handle(query):
results.append(result)
msg = query.resp_messages[0]
assert len(msg.content) == 2
assert msg.content[0].type == 'text'
assert msg.content[1].type == 'image_url'
@pytest.mark.asyncio
async def test_command_result_empty_interrupts(self, fake_app, mock_event_ctx, mock_execute_factory):
"""Empty result yields INTERRUPT."""
command = get_command_handler()
entities = get_entities()
fake_app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
fake_app.cmd_mgr.execute = mock_execute_factory(text=None)
handler = command.CommandHandler(fake_app)
query = command_query('empty')
results = []
async for result in handler.handle(query):
results.append(result)
assert results[0].result_type == entities.ResultType.INTERRUPT
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestCommandHandlerHelper:
"""Tests for helper methods."""
def test_cut_str_short(self, fake_app):
"""cut_str returns short string unchanged."""
command = get_command_handler()
handler = command.CommandHandler(fake_app)
result = handler.cut_str('short text')
assert result == 'short text'
def test_cut_str_long(self, fake_app):
"""cut_str truncates long string."""
command = get_command_handler()
handler = command.CommandHandler(fake_app)
result = handler.cut_str('this is a very long string that exceeds twenty characters')
assert '...' in result
assert len(result) <= 23
def test_cut_str_multiline(self, fake_app):
"""cut_str truncates multiline string."""
command = get_command_handler()
handler = command.CommandHandler(fake_app)
result = handler.cut_str('first line\nsecond line')
assert '...' in result
@@ -0,0 +1,113 @@
"""Unit tests for config_coercion module"""
from __future__ import annotations
import pytest
from langbot.pkg.pipeline.config_coercion import _coerce_value, coerce_pipeline_config
class TestCoerceValue:
"""Tests for _coerce_value function"""
def test_none_passthrough(self):
assert _coerce_value(None, 'integer') is None
assert _coerce_value(None, 'boolean') is None
def test_string_to_integer(self):
assert _coerce_value('120', 'integer') == 120
assert _coerce_value('0', 'integer') == 0
assert _coerce_value('-5', 'integer') == -5
def test_integer_passthrough(self):
assert _coerce_value(42, 'integer') == 42
def test_string_to_float(self):
assert _coerce_value('3.14', 'number') == 3.14
assert _coerce_value('3.14', 'float') == 3.14
def test_int_to_float(self):
assert _coerce_value(3, 'number') == 3.0
assert isinstance(_coerce_value(3, 'number'), float)
def test_float_passthrough(self):
assert _coerce_value(3.14, 'float') == 3.14
def test_string_to_bool(self):
assert _coerce_value('true', 'boolean') is True
assert _coerce_value('True', 'boolean') is True
assert _coerce_value('false', 'boolean') is False
assert _coerce_value('False', 'boolean') is False
def test_bool_passthrough(self):
assert _coerce_value(True, 'boolean') is True
assert _coerce_value(False, 'boolean') is False
def test_invalid_bool_string_raises(self):
with pytest.raises(ValueError):
_coerce_value('notabool', 'boolean')
def test_unknown_type_passthrough(self):
assert _coerce_value('hello', 'string') == 'hello'
assert _coerce_value('hello', 'unknown') == 'hello'
def test_invalid_integer_raises(self):
with pytest.raises(ValueError):
_coerce_value('abc', 'integer')
class TestCoercePipelineConfig:
"""Tests for coerce_pipeline_config function"""
def _make_meta(self, section_name: str, stage_name: str, fields: list[dict]) -> dict:
return {
'name': section_name,
'stages': [{'name': stage_name, 'config': fields}],
}
def test_coerce_integer_in_config(self):
config = {'trigger': {'misc': {'timeout': '120'}}}
meta = self._make_meta('trigger', 'misc', [{'name': 'timeout', 'type': 'integer'}])
coerce_pipeline_config(config, meta)
assert config['trigger']['misc']['timeout'] == 120
def test_coerce_boolean_in_config(self):
config = {'output': {'misc': {'at-sender': 'true'}}}
meta = self._make_meta('output', 'misc', [{'name': 'at-sender', 'type': 'boolean'}])
coerce_pipeline_config(config, meta)
assert config['output']['misc']['at-sender'] is True
def test_missing_section_skipped(self):
config = {'ai': {}}
meta = self._make_meta('trigger', 'misc', [{'name': 'x', 'type': 'integer'}])
coerce_pipeline_config(config, meta) # should not raise
def test_missing_field_skipped(self):
config = {'trigger': {'misc': {}}}
meta = self._make_meta('trigger', 'misc', [{'name': 'nonexistent', 'type': 'integer'}])
coerce_pipeline_config(config, meta) # should not raise
def test_invalid_value_logs_warning(self, caplog):
config = {'trigger': {'misc': {'timeout': 'abc'}}}
meta = self._make_meta('trigger', 'misc', [{'name': 'timeout', 'type': 'integer'}])
import logging
with caplog.at_level(logging.WARNING):
coerce_pipeline_config(config, meta)
assert config['trigger']['misc']['timeout'] == 'abc' # unchanged
assert 'Failed to coerce' in caplog.text
def test_empty_metadata(self):
config = {'trigger': {'misc': {'timeout': '120'}}}
coerce_pipeline_config(config) # no metadata args, should not raise
def test_multiple_metadata(self):
config = {
'trigger': {'misc': {'timeout': '120'}},
'output': {'misc': {'at-sender': 'false'}},
}
meta_trigger = self._make_meta('trigger', 'misc', [{'name': 'timeout', 'type': 'integer'}])
meta_output = self._make_meta('output', 'misc', [{'name': 'at-sender', 'type': 'boolean'}])
coerce_pipeline_config(config, meta_trigger, meta_output)
assert config['trigger']['misc']['timeout'] == 120
assert config['output']['misc']['at-sender'] is False
+359
View File
@@ -0,0 +1,359 @@
"""
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 AsyncMock, 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
assert len(result.new_query.resp_message_chain) == 1
components = list(result.new_query.resp_message_chain[0])
assert len(components) == 1
assert isinstance(components[0], platform_message.Plain)
assert components[0].text == 'short response'
@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
components = list(result.new_query.resp_message_chain[0])
assert [type(component) for component in components] == [
platform_message.Plain,
platform_message.Image,
]
assert components[0].text == 'short'
assert components[1].url == 'https://example.com/img.png'
@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 = []
result = await stage.process(query, 'LongTextProcessStage')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query is query
@pytest.mark.asyncio
async def test_empty_response_message_chain_does_not_call_strategy(self):
"""Empty response chains should be a no-op for long text processing."""
longtext = get_longtext_module()
entities = get_entities_module()
app = FakeApp()
stage = longtext.LongTextProcessStage(app)
stage.strategy_impl = AsyncMock()
query = text_query('hello')
query.pipeline_config = make_longtext_config(strategy='forward', threshold=1)
query.resp_message_chain = []
result = await stage.process(query, 'LongTextProcessStage')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query is query
stage.strategy_impl.process.assert_not_called()
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
components = list(result.new_query.resp_message_chain[0])
assert len(components) == 1
assert isinstance(components[0], platform_message.Forward)
@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_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
components = list(result.new_query.resp_message_chain[0])
assert len(components) == 1
assert isinstance(components[0], platform_message.Plain)
assert components[0].text == short_text
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)
+321
View File
@@ -0,0 +1,321 @@
"""
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 precisely.
Algorithm: traverse backwards, collect while current_round < max_round, count user messages as rounds.
For max_round=2 with 7 messages (u1, a1, u2, a2, u3, a3, u_current):
- Iterate: u_current(r=0<2, collect, r=1), a3(r=1<2, collect), u3(r=1<2, collect, r=2)
- a2: r=2 not < 2 → break
- Collected reverse: [u_current, a3, u3]
- Reversed: [u3, a3, u_current] = 3 messages
"""
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
# 7 messages = 3 full rounds + 1 current user
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 keep exactly 3 messages: message3, response3, current message
messages = result.new_query.messages
assert len(messages) == 3
# Verify exact message content
assert messages[0].role == 'user'
assert messages[0].content == 'message 3'
assert messages[1].role == 'assistant'
assert messages[1].content == 'response 3'
assert messages[2].role == 'user'
assert messages[2].content == 'current message'
@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
messages = result.new_query.messages
assert [(msg.role, msg.content) for msg in messages] == [
('user', 'user2'),
('assistant', 'asst2'),
('user', 'user3'),
]
@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
messages = result.new_query.messages
assert [(msg.role, msg.content) for msg in messages] == [('user', 'current')]
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')
+353
View File
@@ -0,0 +1,353 @@
"""
Unit tests for N8nServiceAPIRunner._process_response
Tests cover four scenarios:
- Stream adapter + n8n stream format (type:item/end)
- Stream adapter + n8n plain JSON
- Non-stream adapter + n8n stream format
- Non-stream adapter + n8n plain JSON
"""
from __future__ import annotations
import json
import sys
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
import langbot_plugin.api.entities.builtin.provider.message as provider_message
# Break the circular import chain while importing n8nsvapi:
# n8nsvapi → runner → app → pipelinemgr → all runners → runner (partially init)
# The stubs are restored in a ``finally`` block so this module does NOT pollute
# sys.modules for other test modules (e.g. ones importing the real
# LocalAgentRunner, which would otherwise inherit ``object`` and break).
# Mirrors master's intent but uses try/finally so a raised import doesn't
# leave the global namespace in a stubbed state, and includes
# ``langbot.pkg.utils.httpclient`` which master didn't stub.
_runner_stub = MagicMock()
_runner_stub.runner_class = lambda name: (lambda cls: cls) # no-op decorator
_runner_stub.RequestRunner = object
_import_stubs = {
'langbot.pkg.provider.runner': _runner_stub,
'langbot.pkg.core.app': MagicMock(),
'langbot.pkg.utils.httpclient': MagicMock(),
}
_saved_modules = {name: sys.modules.get(name) for name in _import_stubs}
for _name, _stub in _import_stubs.items():
sys.modules[_name] = _stub
try:
from langbot.pkg.provider.runners.n8nsvapi import N8nServiceAPIRunner
finally:
for _name, _original in _saved_modules.items():
if _original is None:
sys.modules.pop(_name, None)
else:
sys.modules[_name] = _original
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def make_runner(output_key: str = 'response') -> N8nServiceAPIRunner:
ap = Mock()
ap.logger = Mock()
pipeline_config = {
'ai': {
'n8n-service-api': {
'webhook-url': 'http://test-n8n/webhook',
'output-key': output_key,
'auth-type': 'none',
}
}
}
return N8nServiceAPIRunner(ap, pipeline_config)
def make_mock_response(chunks: list[bytes | str], status: int = 200):
"""Build a minimal aiohttp.ClientResponse mock with iter_chunked support."""
response = Mock()
response.status = status
async def iter_chunked(size):
for chunk in chunks:
yield chunk
response.content = Mock()
response.content.iter_chunked = iter_chunked
return response
async def collect_chunks(runner: N8nServiceAPIRunner, chunks: list[bytes | str]):
"""Run _process_response and collect all yielded MessageChunks."""
response = make_mock_response(chunks)
result = []
async for chunk in runner._process_response(response):
result.append(chunk)
return result
# ---------------------------------------------------------------------------
# _process_response: stream format (type:item/end)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_stream_format_single_item():
"""Single item + end in one chunk yields final chunk with full content."""
runner = make_runner()
data = b'{"type":"item","content":"hello"}{"type":"end"}'
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == 'hello'
assert chunks[0].msg_sequence == 1
@pytest.mark.asyncio
async def test_stream_format_multi_item_accumulates():
"""Multiple items accumulate into full_content."""
runner = make_runner()
chunks_data = [
b'{"type":"item","content":"foo"}',
b'{"type":"item","content":"bar"}',
b'{"type":"end"}',
]
chunks = await collect_chunks(runner, chunks_data)
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == 'foobar'
assert chunks[0].msg_sequence == 1
@pytest.mark.asyncio
async def test_stream_format_batches_every_8_items():
"""Every 8th item triggers an intermediate yield before the final."""
runner = make_runner()
items = [f'{{"type":"item","content":"{i}"}}' for i in range(8)]
items.append('{"type":"end"}')
data = ''.join(items).encode()
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 2
assert chunks[0].is_final is False
assert chunks[0].content == '01234567'
assert chunks[0].msg_sequence == 1
assert chunks[1].is_final is True
assert chunks[1].content == '01234567'
assert chunks[1].msg_sequence == 2
@pytest.mark.asyncio
async def test_stream_format_split_across_network_chunks():
"""JSON split across multiple network chunks is reassembled correctly."""
runner = make_runner()
part1 = b'{"type":"item","con'
part2 = b'tent":"world"}{"type":"end"}'
chunks = await collect_chunks(runner, [part1, part2])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == 'world'
@pytest.mark.asyncio
async def test_stream_format_no_spurious_empty_yield():
"""chunk_idx==0 guard prevents spurious empty yield before any item is received."""
runner = make_runner()
# Send some non-stream JSON first, then stream
data = b'{"type":"item","content":"x"}{"type":"end"}'
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].content == 'x'
# ---------------------------------------------------------------------------
# _process_response: plain JSON fallback
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_plain_json_with_output_key():
"""Plain JSON with matching output_key extracts value via output_key."""
runner = make_runner(output_key='response')
data = json.dumps({'response': 'hello world'}).encode()
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == 'hello world'
@pytest.mark.asyncio
async def test_plain_json_output_key_not_found():
"""Plain JSON without output_key falls back to entire JSON string."""
runner = make_runner(output_key='response')
payload = {'other_key': 'hello'}
data = json.dumps(payload).encode()
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert json.loads(chunks[0].content) == payload
@pytest.mark.asyncio
async def test_plain_json_output_key_empty_string():
"""output_key present but value is empty string — returns empty string, not whole JSON."""
runner = make_runner(output_key='response')
data = json.dumps({'response': ''}).encode()
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == ''
@pytest.mark.asyncio
async def test_plain_json_non_dict_response():
"""Plain JSON array falls back to raw text."""
runner = make_runner()
data = b'["a", "b"]'
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == '["a", "b"]'
@pytest.mark.asyncio
async def test_invalid_json_returns_raw_text():
"""Non-JSON response returns raw text as-is."""
runner = make_runner()
data = b'plain text response'
chunks = await collect_chunks(runner, [data])
assert len(chunks) == 1
assert chunks[0].is_final is True
assert chunks[0].content == 'plain text response'
# ---------------------------------------------------------------------------
# _call_webhook: output type depends on is_stream
# ---------------------------------------------------------------------------
def make_query(is_stream: bool):
"""Build a minimal Query mock."""
query = Mock()
query.adapter = AsyncMock()
query.adapter.is_stream_output_supported = AsyncMock(return_value=is_stream)
session = Mock()
session.using_conversation = Mock()
session.using_conversation.uuid = 'test-uuid'
session.launcher_type = Mock()
session.launcher_type.value = 'person'
session.launcher_id = '12345'
query.session = session
query.user_message = Mock()
query.user_message.content = 'hi'
query.variables = {}
return query
def make_http_session_mock(response_bytes: bytes, status: int = 200):
"""Mock httpclient.get_session() returning a session whose post() yields response_bytes."""
mock_response = make_mock_response([response_bytes], status=status)
mock_response.status = status
mock_cm = AsyncMock()
mock_cm.__aenter__ = AsyncMock(return_value=mock_response)
mock_cm.__aexit__ = AsyncMock(return_value=False)
mock_session = Mock()
mock_session.post = Mock(return_value=mock_cm)
return mock_session
@pytest.mark.asyncio
async def test_call_webhook_nonstream_adapter_plain_json():
"""Non-stream adapter + plain JSON → single Message with output_key value."""
runner = make_runner(output_key='response')
query = make_query(is_stream=False)
http_session = make_http_session_mock(json.dumps({'response': 'result text'}).encode())
with patch('langbot.pkg.provider.runners.n8nsvapi.httpclient.get_session', return_value=http_session):
results = []
async for msg in runner._call_webhook(query):
results.append(msg)
assert len(results) == 1
assert isinstance(results[0], provider_message.Message)
assert results[0].content == 'result text'
@pytest.mark.asyncio
async def test_call_webhook_stream_adapter_stream_format():
"""Stream adapter + stream format → MessageChunks, last is_final."""
runner = make_runner()
query = make_query(is_stream=True)
data = b'{"type":"item","content":"hi"}{"type":"end"}'
http_session = make_http_session_mock(data)
with patch('langbot.pkg.provider.runners.n8nsvapi.httpclient.get_session', return_value=http_session):
results = []
async for msg in runner._call_webhook(query):
results.append(msg)
assert all(isinstance(r, provider_message.MessageChunk) for r in results)
assert results[-1].is_final is True
assert results[-1].content == 'hi'
@pytest.mark.asyncio
async def test_call_webhook_stream_adapter_plain_json():
"""Stream adapter + plain JSON → single MessageChunk with is_final=True."""
runner = make_runner(output_key='response')
query = make_query(is_stream=True)
data = json.dumps({'response': 'fallback'}).encode()
http_session = make_http_session_mock(data)
with patch('langbot.pkg.provider.runners.n8nsvapi.httpclient.get_session', return_value=http_session):
results = []
async for msg in runner._call_webhook(query):
results.append(msg)
assert all(isinstance(r, provider_message.MessageChunk) for r in results)
assert results[-1].is_final is True
assert results[-1].content == 'fallback'
@pytest.mark.asyncio
async def test_call_webhook_nonstream_adapter_stream_format():
"""Non-stream adapter + stream format → single Message with accumulated content."""
runner = make_runner()
query = make_query(is_stream=False)
data = b'{"type":"item","content":"foo"}{"type":"item","content":"bar"}{"type":"end"}'
http_session = make_http_session_mock(data)
with patch('langbot.pkg.provider.runners.n8nsvapi.httpclient.get_session', return_value=http_session):
results = []
async for msg in runner._call_webhook(query):
results.append(msg)
assert len(results) == 1
assert isinstance(results[0], provider_message.Message)
assert results[0].content == 'foobar'
@@ -0,0 +1,43 @@
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.service.pipeline import PipelineService
@pytest.mark.asyncio
async def test_update_pipeline_filters_protected_fields_without_mutating_input(mock_app):
service = PipelineService(mock_app)
loaded_pipeline = Mock()
service.get_pipeline = AsyncMock(return_value=loaded_pipeline)
bot = Mock(uuid='bot-uuid')
bot_result = Mock(all=Mock(return_value=[bot]))
mock_app.persistence_mgr.execute_async = AsyncMock(side_effect=[None, bot_result])
mock_app.bot_service = Mock(update_bot=AsyncMock())
mock_app.pipeline_mgr = Mock(remove_pipeline=AsyncMock(), load_pipeline=AsyncMock())
mock_app.sess_mgr.session_list = []
pipeline_data = {
'uuid': 'caller-uuid',
'for_version': '1.0.0',
'stages': ['CallerStage'],
'is_default': True,
'name': 'Updated pipeline',
}
original_pipeline_data = pipeline_data.copy()
await service.update_pipeline('pipeline-uuid', pipeline_data)
assert pipeline_data == original_pipeline_data
update_stmt = mock_app.persistence_mgr.execute_async.await_args_list[0].args[0]
updated_fields = {getattr(field, 'key', str(field)) for field in update_stmt._values}
assert updated_fields == {'name'}
mock_app.bot_service.update_bot.assert_awaited_once_with(
'bot-uuid',
{'use_pipeline_name': 'Updated pipeline'},
)
mock_app.pipeline_mgr.remove_pipeline.assert_awaited_once_with('pipeline-uuid')
mock_app.pipeline_mgr.load_pipeline.assert_awaited_once_with(loaded_pipeline)
@@ -0,0 +1,207 @@
"""
PipelineManager unit tests
"""
import pytest
from unittest.mock import AsyncMock, Mock
from importlib import import_module
def get_pipelinemgr_module():
return import_module('langbot.pkg.pipeline.pipelinemgr')
def get_stage_module():
return import_module('langbot.pkg.pipeline.stage')
def get_entities_module():
return import_module('langbot.pkg.pipeline.entities')
def get_persistence_pipeline_module():
return import_module('langbot.pkg.entity.persistence.pipeline')
@pytest.mark.asyncio
async def test_pipeline_manager_initialize(mock_app):
"""Test pipeline manager initialization"""
pipelinemgr = get_pipelinemgr_module()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[])))
manager = pipelinemgr.PipelineManager(mock_app)
await manager.initialize()
assert manager.stage_dict is not None
assert len(manager.pipelines) == 0
@pytest.mark.asyncio
async def test_load_pipeline(mock_app):
"""Test loading a single pipeline"""
pipelinemgr = get_pipelinemgr_module()
persistence_pipeline = get_persistence_pipeline_module()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[])))
manager = pipelinemgr.PipelineManager(mock_app)
await manager.initialize()
# Create test pipeline entity
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.stages = []
pipeline_entity.config = {'test': 'config'}
pipeline_entity.extensions_preferences = {'plugins': []}
await manager.load_pipeline(pipeline_entity)
assert len(manager.pipelines) == 1
assert manager.pipelines[0].pipeline_entity.uuid == 'test-uuid'
@pytest.mark.asyncio
async def test_get_pipeline_by_uuid(mock_app):
"""Test getting pipeline by UUID"""
pipelinemgr = get_pipelinemgr_module()
persistence_pipeline = get_persistence_pipeline_module()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[])))
manager = pipelinemgr.PipelineManager(mock_app)
await manager.initialize()
# Create and add test pipeline
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.stages = []
pipeline_entity.config = {}
pipeline_entity.extensions_preferences = {'plugins': []}
await manager.load_pipeline(pipeline_entity)
# Test retrieval
result = await manager.get_pipeline_by_uuid('test-uuid')
assert result is not None
assert result.pipeline_entity.uuid == 'test-uuid'
# Test non-existent UUID
result = await manager.get_pipeline_by_uuid('non-existent')
assert result is None
@pytest.mark.asyncio
async def test_remove_pipeline(mock_app):
"""Test removing a pipeline"""
pipelinemgr = get_pipelinemgr_module()
persistence_pipeline = get_persistence_pipeline_module()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[])))
manager = pipelinemgr.PipelineManager(mock_app)
await manager.initialize()
# Create and add test pipeline
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.stages = []
pipeline_entity.config = {}
pipeline_entity.extensions_preferences = {'plugins': []}
await manager.load_pipeline(pipeline_entity)
assert len(manager.pipelines) == 1
# Remove pipeline
await manager.remove_pipeline('test-uuid')
assert len(manager.pipelines) == 0
@pytest.mark.asyncio
async def test_runtime_pipeline_execute(mock_app, sample_query):
"""Test runtime pipeline execution with real Pydantic models."""
pipelinemgr = get_pipelinemgr_module()
stage = get_stage_module()
persistence_pipeline = get_persistence_pipeline_module()
entities = get_entities_module()
# Create result using real Pydantic model (not Mock) to ensure validation
real_result = entities.StageProcessResult(
result_type=entities.ResultType.CONTINUE,
new_query=sample_query,
user_notice='',
console_notice='',
debug_notice='',
error_notice='',
)
mock_stage = Mock(spec=stage.PipelineStage)
mock_stage.process = AsyncMock(return_value=real_result)
# Create stage container
stage_container = pipelinemgr.StageInstContainer(inst_name='TestStage', inst=mock_stage)
# Create pipeline entity
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.config = sample_query.pipeline_config
pipeline_entity.extensions_preferences = {'plugins': []}
# Create runtime pipeline
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [stage_container])
# Mock plugin connector
event_ctx = Mock()
event_ctx.is_prevented_default = Mock(return_value=False)
mock_app.plugin_connector.emit_event = AsyncMock(return_value=event_ctx)
# Add query to cached_queries to prevent KeyError in finally block
mock_app.query_pool.cached_queries[sample_query.query_id] = sample_query
# Execute pipeline
await runtime_pipeline.run(sample_query)
# Verify stage was called
mock_stage.process.assert_called_once()
def test_runtime_pipeline_prefers_local_agent_mcp_resources(mock_app):
"""Local Agent resource selection should override legacy extension prefs."""
pipelinemgr = get_pipelinemgr_module()
persistence_pipeline = get_persistence_pipeline_module()
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.config = {
'ai': {
'local-agent': {
'mcp-resources': [{'server_uuid': 'srv-new', 'uri': 'file:///new.md'}],
'mcp-resource-agent-read-enabled': False,
}
}
}
pipeline_entity.extensions_preferences = {
'mcp_resources': [{'server_uuid': 'srv-old', 'uri': 'file:///old.md'}],
'mcp_resource_agent_read_enabled': True,
}
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [])
assert runtime_pipeline.mcp_resource_attachments == [{'server_uuid': 'srv-new', 'uri': 'file:///new.md'}]
assert runtime_pipeline.mcp_resource_agent_read_enabled is False
def test_runtime_pipeline_falls_back_to_extension_mcp_resources(mock_app):
"""Existing extension prefs remain compatible until a Local Agent value exists."""
pipelinemgr = get_pipelinemgr_module()
persistence_pipeline = get_persistence_pipeline_module()
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.config = {'ai': {'local-agent': {}}}
pipeline_entity.extensions_preferences = {
'mcp_resources': [{'server_uuid': 'srv-old', 'uri': 'file:///old.md'}],
'mcp_resource_agent_read_enabled': False,
}
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [])
assert runtime_pipeline.mcp_resource_attachments == [{'server_uuid': 'srv-old', 'uri': 'file:///old.md'}]
assert runtime_pipeline.mcp_resource_agent_read_enabled is False
+290
View File
@@ -0,0 +1,290 @@
"""
Unit tests for QueryPool.
Tests query management, ID generation, and async context handling.
"""
from __future__ import annotations
import pytest
from unittest.mock import Mock, patch
from langbot.pkg.pipeline.pool import QueryPool
pytestmark = pytest.mark.asyncio
class TestQueryPoolInit:
"""Tests for QueryPool initialization."""
def test_init_creates_empty_pool(self):
"""QueryPool initializes with empty lists."""
pool = QueryPool()
assert pool.queries == []
assert pool.cached_queries == {}
assert pool.query_id_counter == 0
assert pool.pool_lock is not None
assert pool.condition is not None
def test_init_counter_starts_at_zero(self):
"""Counter starts at zero."""
pool = QueryPool()
assert pool.query_id_counter == 0
class TestQueryPoolAddQuery:
"""Tests for add_query method."""
async def test_add_query_adds_query_with_id(self):
"""add_query creates, stores, and caches a Query with the correct ID."""
pool = QueryPool()
# Mock Query creation
mock_query = Mock()
mock_query.query_id = 0
mock_query.bot_uuid = 'test-bot-uuid'
mock_query.launcher_id = 12345
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.return_value = mock_query
await pool.add_query(
bot_uuid='test-bot-uuid',
launcher_type=Mock(),
launcher_id=12345,
sender_id=12345,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
# Query is added to list and cache
assert pool.queries[0] is mock_query
assert pool.cached_queries[0] is mock_query
assert mock_query.query_id == 0
async def test_add_query_increments_counter(self):
"""Each add_query increments the counter."""
pool = QueryPool()
mock_query1 = Mock()
mock_query1.query_id = 0
mock_query2 = Mock()
mock_query2.query_id = 1
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.side_effect = [mock_query1, mock_query2]
await pool.add_query(
bot_uuid='bot1',
launcher_type=Mock(),
launcher_id=1,
sender_id=1,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
await pool.add_query(
bot_uuid='bot2',
launcher_type=Mock(),
launcher_id=2,
sender_id=2,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
assert pool.query_id_counter == 2
assert pool.queries[0].query_id == 0
assert pool.queries[1].query_id == 1
async def test_add_query_appends_to_list(self):
"""Query is appended to queries list."""
pool = QueryPool()
mock_query = Mock()
mock_query.query_id = 0
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.return_value = mock_query
await pool.add_query(
bot_uuid='bot1',
launcher_type=Mock(),
launcher_id=1,
sender_id=1,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
assert len(pool.queries) == 1
assert pool.queries[0] is mock_query
async def test_add_query_caches_query(self):
"""Query is cached by query_id."""
pool = QueryPool()
mock_query = Mock()
mock_query.query_id = 0
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.return_value = mock_query
await pool.add_query(
bot_uuid='bot1',
launcher_type=Mock(),
launcher_id=1,
sender_id=1,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
assert 0 in pool.cached_queries
assert pool.cached_queries[0] is mock_query
async def test_add_query_with_pipeline_uuid(self):
"""Query can have pipeline_uuid set."""
pool = QueryPool()
mock_query = Mock()
mock_query.query_id = 0
mock_query.pipeline_uuid = 'test-pipeline-uuid'
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.return_value = mock_query
await pool.add_query(
bot_uuid='bot1',
launcher_type=Mock(),
launcher_id=1,
sender_id=1,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
pipeline_uuid='test-pipeline-uuid',
)
# Verify pipeline_uuid was passed to Query constructor
call_kwargs = MockQuery.call_args[1]
assert call_kwargs['pipeline_uuid'] == 'test-pipeline-uuid'
async def test_add_query_sets_routed_by_rule_variable(self):
"""Query has _routed_by_rule variable."""
pool = QueryPool()
mock_query = Mock()
mock_query.query_id = 0
mock_query.variables = {'_routed_by_rule': True}
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.return_value = mock_query
await pool.add_query(
bot_uuid='bot1',
launcher_type=Mock(),
launcher_id=1,
sender_id=1,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
routed_by_rule=True,
)
# Verify variables includes _routed_by_rule
call_kwargs = MockQuery.call_args[1]
assert call_kwargs['variables']['_routed_by_rule'] is True
async def test_add_query_notifier_condition(self):
"""add_query notifies waiting consumers."""
pool = QueryPool()
mock_query = Mock()
mock_query.query_id = 0
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.return_value = mock_query
# Track if notify_all was called
original_notify = pool.condition.notify_all
notify_called = []
def mock_notify():
notify_called.append(True)
return original_notify()
pool.condition.notify_all = mock_notify
await pool.add_query(
bot_uuid='bot1',
launcher_type=Mock(),
launcher_id=1,
sender_id=1,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
assert len(notify_called) == 1
class TestQueryPoolContext:
"""Tests for async context manager."""
async def test_aenter_acquires_lock(self):
"""__aenter__ acquires the pool lock."""
pool = QueryPool()
async with pool as p:
# Lock is acquired
assert pool.pool_lock.locked()
assert p is pool
async def test_aexit_releases_lock(self):
"""__aexit__ releases the pool lock."""
pool = QueryPool()
async with pool:
pass
# Lock is released after context exit
assert not pool.pool_lock.locked()
class TestQueryPoolEdgeCases:
"""Tests for edge cases."""
async def test_multiple_queries_cached_correctly(self):
"""Multiple queries are cached separately."""
pool = QueryPool()
mock_queries = []
for i in range(5):
q = Mock()
q.query_id = i
mock_queries.append(q)
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query') as MockQuery:
MockQuery.side_effect = mock_queries
for i in range(5):
await pool.add_query(
bot_uuid=f'bot{i}',
launcher_type=Mock(),
launcher_id=i,
sender_id=i,
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
# All cached
assert len(pool.cached_queries) == 5
# Each query is cached by its ID
for i in range(5):
assert pool.cached_queries[i] is mock_queries[i]
+491
View File
@@ -0,0 +1,491 @@
"""
Unit tests for PreProcessor pipeline stage.
Tests cover preprocessing behavior including:
- Normal text message processing
- Empty message handling
- Unsupported message segment handling
- Image/file segment behavior
- Model selection and fallback
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from importlib import import_module
from types import SimpleNamespace
from tests.factories import (
FakeApp,
text_query,
empty_query,
image_query,
group_text_query,
)
def get_preproc_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.pipeline.preproc.preproc')
def get_entities_module():
"""Lazy import for pipeline entities."""
return import_module('langbot.pkg.pipeline.entities')
class TestPreProcessorNormalText:
"""Tests for normal text message preprocessing."""
@pytest.mark.asyncio
async def test_normal_text_continues(self):
"""Normal text message should continue pipeline."""
preproc = get_preproc_module()
entities = get_entities_module()
app = FakeApp()
# Mock session manager to return a session
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
# Mock conversation
mock_conversation = Mock()
mock_conversation.prompt = Mock()
mock_conversation.prompt.messages = []
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.update_time = Mock()
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
# Mock model manager
mock_model = Mock()
mock_model.model_entity = Mock()
mock_model.model_entity.uuid = 'test-model-uuid'
mock_model.model_entity.abilities = ['func_call', 'vision']
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=mock_model)
# Mock tool manager
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
# Mock plugin connector
mock_event_ctx = Mock()
mock_event_ctx.event = Mock()
mock_event_ctx.event.default_prompt = []
mock_event_ctx.event.prompt = []
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = text_query('hello world')
result = await stage.process(query, 'PreProcessor')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query is not None
@pytest.mark.asyncio
async def test_normal_text_sets_user_message(self):
"""PreProcessor should set user_message from text content."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
mock_model = Mock()
mock_model.model_entity = Mock(uuid='test-model', abilities=['func_call'])
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=mock_model)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = text_query('test message')
result = await stage.process(query, 'PreProcessor')
assert result.new_query.user_message is not None
assert result.new_query.user_message.role == 'user'
class TestPreProcessorEmptyMessage:
"""Tests for empty message handling."""
@pytest.mark.asyncio
async def test_empty_message_continues(self):
"""Empty message should follow expected behavior."""
preproc = get_preproc_module()
entities = get_entities_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=None)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = empty_query()
result = await stage.process(query, 'PreProcessor')
# Empty message should still continue with an empty provider content list.
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query.user_message is not None
assert result.new_query.user_message.content == []
class TestPreProcessorImageSegment:
"""Tests for image segment handling."""
@pytest.mark.asyncio
async def test_image_with_vision_model(self):
"""Image should be included when model supports vision."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
# Model with vision support
mock_model = Mock()
mock_model.model_entity = Mock(uuid='vision-model', abilities=['func_call', 'vision'])
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=mock_model)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
# Image query with base64
query = image_query(text='look at this', url=None)
# Set base64 on the image component
import langbot_plugin.api.entities.builtin.platform.message as platform_message
chain = platform_message.MessageChain(
[
platform_message.Plain(text='look at this'),
platform_message.Image(base64='data:image/png;base64,abc123'),
]
)
query.message_chain = chain
result = await stage.process(query, 'PreProcessor')
assert result.result_type == preproc.entities.ResultType.CONTINUE
# User message should have content
assert result.new_query.user_message.content is not None
@pytest.mark.asyncio
async def test_image_without_vision_model(self):
"""Image should be excluded when model doesn't support vision."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
# Model WITHOUT vision support
mock_model = Mock()
mock_model.model_entity = Mock(uuid='text-only-model', abilities=['func_call'])
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=mock_model)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = image_query(text='describe this')
result = await stage.process(query, 'PreProcessor')
assert result.result_type == preproc.entities.ResultType.CONTINUE
class TestPreProcessorModelSelection:
"""Tests for model selection and fallback behavior."""
@pytest.mark.asyncio
async def test_primary_model_selected(self):
"""Primary model UUID should be set in query."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
mock_model = Mock()
mock_model.model_entity = Mock(uuid='primary-model-uuid', abilities=['func_call'])
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=mock_model)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = text_query('hello')
# Set pipeline config with primary model
query.pipeline_config = {
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {
'model': {'primary': 'primary-model-uuid', 'fallbacks': []},
'prompt': 'default',
},
},
'output': {'misc': {'at-sender': False}},
'trigger': {'misc': {}},
}
result = await stage.process(query, 'PreProcessor')
assert result.new_query.use_llm_model_uuid == 'primary-model-uuid'
@pytest.mark.asyncio
async def test_fallback_models_resolved(self):
"""Fallback model UUIDs should be resolved and stored."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
# Primary model
mock_primary = Mock()
mock_primary.model_entity = Mock(uuid='primary-uuid', abilities=['func_call'])
# Fallback model
mock_fallback = Mock()
mock_fallback.model_entity = Mock(uuid='fallback-uuid', abilities=['func_call'])
async def mock_get_model(uuid):
if uuid == 'primary-uuid':
return mock_primary
elif uuid == 'fallback-uuid':
return mock_fallback
raise ValueError(f'Model {uuid} not found')
app.model_mgr.get_model_by_uuid = AsyncMock(side_effect=mock_get_model)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = text_query('hello')
query.pipeline_config = {
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {
'model': {'primary': 'primary-uuid', 'fallbacks': ['fallback-uuid']},
'prompt': 'default',
},
},
'output': {'misc': {'at-sender': False}},
'trigger': {'misc': {}},
}
result = await stage.process(query, 'PreProcessor')
assert '_fallback_model_uuids' in result.new_query.variables
assert 'fallback-uuid' in result.new_query.variables['_fallback_model_uuids']
class TestPreProcessorVariables:
"""Tests for query variable extraction."""
@pytest.mark.asyncio
async def test_variables_set_from_query(self):
"""PreProcessor should set variables from query context."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = 'conv-123'
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=None)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = text_query('hello', sender_id=67890)
result = await stage.process(query, 'PreProcessor')
variables = result.new_query.variables
assert 'launcher_type' in variables
assert 'launcher_id' in variables
assert 'sender_id' in variables
assert variables['sender_id'] == 67890
assert 'user_message_text' in variables
@pytest.mark.asyncio
async def test_group_variables_include_group_name(self):
"""Group messages should include group_name variable."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='group')
mock_session.launcher_id = 99999
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=None)
app.tool_mgr.get_all_tools = AsyncMock(return_value=[])
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = group_text_query('hello', group_id=99999)
result = await stage.process(query, 'PreProcessor')
variables = result.new_query.variables
assert 'group_name' in variables
assert 'sender_name' in variables
class TestPreProcessorToolSelection:
"""Tests for Local Agent tool selection."""
@pytest.mark.asyncio
async def test_local_agent_filters_selected_tools(self):
"""Only selected tools should be exposed when all-tools mode is off."""
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
mock_conversation.prompt = Mock(messages=[])
mock_conversation.prompt.copy = Mock(return_value=Mock(messages=[]))
mock_conversation.messages = []
mock_conversation.uuid = None
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
mock_model = Mock()
mock_model.model_entity = Mock(uuid='primary-model-uuid', abilities=['func_call'])
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=mock_model)
app.tool_mgr.get_all_tools = AsyncMock(
return_value=[
SimpleNamespace(name='exec'),
SimpleNamespace(name='plugin_tool'),
SimpleNamespace(name='mcp_tool'),
]
)
mock_event_ctx = Mock()
mock_event_ctx.event = Mock(default_prompt=[], prompt=[])
app.plugin_connector.emit_event = AsyncMock(return_value=mock_event_ctx)
stage = preproc.PreProcessor(app)
query = text_query('hello')
query.pipeline_config = {
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {
'model': {'primary': 'primary-model-uuid', 'fallbacks': []},
'prompt': 'default',
'enable-all-tools': False,
'tools': ['plugin_tool'],
},
},
'output': {'misc': {'at-sender': False}},
'trigger': {'misc': {}},
}
result = await stage.process(query, 'PreProcessor')
assert [tool.name for tool in result.new_query.use_funcs] == ['plugin_tool']
@@ -0,0 +1,75 @@
"""
QueryPool unit tests
"""
import pytest
import langbot_plugin.api.entities.builtin.provider.session as provider_session
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
from langbot.pkg.pipeline.pool import QueryPool
class DummyEventLogger(abstract_platform_logger.AbstractEventLogger):
async def info(self, text, images=None, message_session_id=None, no_throw=True):
pass
async def debug(self, text, images=None, message_session_id=None, no_throw=True):
pass
async def warning(self, text, images=None, message_session_id=None, no_throw=True):
pass
async def error(self, text, images=None, message_session_id=None, no_throw=True):
pass
class DummyAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def send_message(self, target_type, target_id, message):
pass
async def reply_message(self, message_source, message, quote_origin=False):
pass
def register_listener(self, event_type, callback):
pass
def unregister_listener(self, event_type, callback):
pass
async def run_async(self):
pass
async def kill(self):
return True
@pytest.mark.asyncio
async def test_add_query_returns_created_query_and_preserves_side_effects(
sample_message_chain,
sample_message_event,
):
"""add_query returns the created Query while keeping pool/cache updates."""
query_pool = QueryPool()
adapter = DummyAdapter(config={}, logger=DummyEventLogger())
query = await query_pool.add_query(
bot_uuid='test-bot-uuid',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=67890,
message_event=sample_message_event,
message_chain=sample_message_chain,
adapter=adapter,
pipeline_uuid='test-pipeline-uuid',
routed_by_rule=True,
)
assert query is query_pool.queries[0]
assert query_pool.cached_queries[0] is query
assert query_pool.query_id_counter == 1
assert query.query_id == 0
assert query.bot_uuid == 'test-bot-uuid'
assert query.pipeline_uuid == 'test-pipeline-uuid'
assert query.variables == {'_routed_by_rule': True}
+341
View File
@@ -0,0 +1,341 @@
"""
RateLimit stage unit tests
Tests the actual RateLimit implementation from pkg.pipeline.ratelimit
"""
import pytest
import asyncio
import time
from unittest.mock import AsyncMock, Mock, patch
from importlib import import_module
import langbot_plugin.api.entities.builtin.provider.session as provider_session
def get_modules():
"""Lazy import to ensure proper initialization order"""
# Import pipelinemgr first to trigger proper stage registration
ratelimit = import_module('langbot.pkg.pipeline.ratelimit.ratelimit')
entities = import_module('langbot.pkg.pipeline.entities')
algo_module = import_module('langbot.pkg.pipeline.ratelimit.algo')
return ratelimit, entities, algo_module
def get_fixedwin_module():
"""Lazy import of FixedWindowAlgo"""
return import_module('langbot.pkg.pipeline.ratelimit.algos.fixedwin')
class TestFixedWindowAlgo:
"""Tests for the actual FixedWindowAlgo implementation.
IMPORTANT: These tests verify the real algorithm logic, not mocks.
"""
@pytest.fixture
def mock_app_for_algo(self):
"""Create mock app for algorithm initialization."""
mock_app = Mock()
mock_app.logger = Mock()
return mock_app
@pytest.fixture
def sample_query_with_rate_limit(self, sample_query):
"""Create query with rate limit configuration."""
sample_query.pipeline_config = {
'safety': {
'rate-limit': {
'window-length': 60, # 60 seconds window
'limitation': 10, # 10 requests per window
'strategy': 'drop',
}
}
}
return sample_query
@pytest.mark.asyncio
async def test_fixedwin_algo_initialization(self, mock_app_for_algo):
"""Test that FixedWindowAlgo initializes correctly."""
fixedwin = get_fixedwin_module()
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
assert algo.containers_lock is not None
assert algo.containers == {}
@pytest.mark.asyncio
async def test_fixedwin_within_limit_returns_true(self, mock_app_for_algo, sample_query_with_rate_limit):
"""Test that requests within limit are allowed."""
fixedwin = get_fixedwin_module()
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# Make requests within limit
for i in range(10):
result = await algo.require_access(
sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345'
)
assert result is True, f'Request {i + 1} should be allowed'
@pytest.mark.asyncio
async def test_fixedwin_exceeds_limit_drop_strategy(self, mock_app_for_algo, sample_query_with_rate_limit):
"""Test that exceeding limit with 'drop' strategy returns False."""
fixedwin = get_fixedwin_module()
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# Exhaust the limit
for i in range(10):
await algo.require_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345')
# Next request should be denied
result = await algo.require_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345')
assert result is False, 'Request exceeding limit should be denied'
@pytest.mark.asyncio
async def test_fixedwin_different_sessions_isolated(self, mock_app_for_algo, sample_query_with_rate_limit):
"""Test that different sessions have independent rate limits."""
fixedwin = get_fixedwin_module()
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# Exhaust limit for session 1
for i in range(10):
await algo.require_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, 'session1')
# Session 2 should still have its own limit
result = await algo.require_access(
sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, 'session2'
)
assert result is True, 'Different session should have independent limit'
@pytest.mark.asyncio
async def test_fixedwin_limit_one_request(self, mock_app_for_algo, sample_query):
"""Test with limitation=1 allows only one request."""
fixedwin = get_fixedwin_module()
sample_query.pipeline_config = {
'safety': {
'rate-limit': {
'window-length': 60,
'limitation': 1, # Only 1 request allowed
'strategy': 'drop',
}
}
}
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# First request allowed
result1 = await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, '12345')
assert result1 is True
# Second request denied
result2 = await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, '12345')
assert result2 is False
@pytest.mark.asyncio
async def test_fixedwin_container_persists(self, mock_app_for_algo, sample_query_with_rate_limit):
"""Test that container is created and persists across requests."""
fixedwin = get_fixedwin_module()
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# First request creates container
await algo.require_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345')
# Key format: 'LauncherTypes.PERSON_12345' (enum string representation)
expected_key = 'LauncherTypes.PERSON_12345'
assert expected_key in algo.containers
container = algo.containers[expected_key]
# Container should have records
assert len(container.records) > 0
@pytest.mark.asyncio
async def test_fixedwin_new_window_clears_records(self, mock_app_for_algo, sample_query):
"""Test that a new time window starts fresh records.
This test verifies the window calculation logic:
- Records are keyed by window start timestamp
- When window advances, new key is created
"""
fixedwin = get_fixedwin_module()
# Use a very short window for testing
sample_query.pipeline_config = {
'safety': {
'rate-limit': {
'window-length': 1, # 1 second window for fast test
'limitation': 5,
'strategy': 'drop',
}
}
}
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# Make requests in current window
now = int(time.time())
window_start = now - now % 1
for i in range(5):
await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'test')
# Key format: 'LauncherTypes.PERSON_test'
expected_key = 'LauncherTypes.PERSON_test'
container = algo.containers[expected_key]
assert window_start in container.records
assert container.records[window_start] == 5
# Wait for next window (1 second)
await asyncio.sleep(1.1)
# New request should be allowed (new window)
result = await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'test')
assert result is True, 'New window should allow new requests'
@pytest.mark.asyncio
async def test_fixedwin_wait_strategy_blocks_until_next_window(self, mock_app_for_algo, sample_query):
"""Test that 'wait' strategy blocks until next window.
NOTE: This test is timing-sensitive and may take ~1 second.
"""
fixedwin = get_fixedwin_module()
# Use 1-second window for testability
sample_query.pipeline_config = {
'safety': {
'rate-limit': {
'window-length': 1,
'limitation': 1, # Only 1 request per second
'strategy': 'wait',
}
}
}
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# First request allowed
start_time = time.time()
result1 = await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'wait_test')
assert result1 is True
# Exhaust limit
await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'wait_test')
# Third request should wait and then succeed
result3 = await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'wait_test')
elapsed = time.time() - start_time
assert result3 is True, 'After wait, request should succeed'
# Should have waited approximately until next window
# With 1-second window, elapsed should be > 0.5 second (allowing for timing variance)
# Note: This is a timing-sensitive test, so we use a generous tolerance
assert elapsed >= 0.5, f'Should have waited for next window, elapsed={elapsed:.2f}s'
@pytest.mark.asyncio
async def test_fixedwin_release_access(self, mock_app_for_algo, sample_query_with_rate_limit):
"""Test that release_access does nothing (current implementation)."""
fixedwin = get_fixedwin_module()
algo = fixedwin.FixedWindowAlgo(mock_app_for_algo)
await algo.initialize()
# release_access is empty in current implementation
await algo.release_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345')
# Should not raise or change state
assert 'person_12345' not in algo.containers
# Original mock-based tests for RateLimit stage integration
@pytest.mark.asyncio
async def test_require_access_allowed(mock_app, sample_query):
"""Test RequireRateLimitOccupancy allows access when rate limit is not exceeded"""
ratelimit, entities, algo_module = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {}
# Create mock algorithm that allows access
mock_algo = Mock(spec=algo_module.ReteLimitAlgo)
mock_algo.require_access = AsyncMock(return_value=True)
mock_algo.initialize = AsyncMock()
stage = ratelimit.RateLimit(mock_app)
# Patch the algorithm selection to use our mock
with patch.object(algo_module, 'preregistered_algos', []):
stage.algo = mock_algo
result = await stage.process(sample_query, 'RequireRateLimitOccupancy')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query == sample_query
mock_algo.require_access.assert_called_once_with(sample_query, 'person', '12345')
@pytest.mark.asyncio
async def test_require_access_denied(mock_app, sample_query):
"""Test RequireRateLimitOccupancy denies access when rate limit is exceeded"""
ratelimit, entities, algo_module = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {}
# Create mock algorithm that denies access
mock_algo = Mock(spec=algo_module.ReteLimitAlgo)
mock_algo.require_access = AsyncMock(return_value=False)
mock_algo.initialize = AsyncMock()
stage = ratelimit.RateLimit(mock_app)
# Patch the algorithm selection to use our mock
with patch.object(algo_module, 'preregistered_algos', []):
stage.algo = mock_algo
result = await stage.process(sample_query, 'RequireRateLimitOccupancy')
assert result.result_type == entities.ResultType.INTERRUPT
assert result.user_notice != ''
mock_algo.require_access.assert_called_once()
@pytest.mark.asyncio
async def test_release_access(mock_app, sample_query):
"""Test ReleaseRateLimitOccupancy releases rate limit occupancy"""
ratelimit, entities, algo_module = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {}
# Create mock algorithm
mock_algo = Mock(spec=algo_module.ReteLimitAlgo)
mock_algo.release_access = AsyncMock()
mock_algo.initialize = AsyncMock()
stage = ratelimit.RateLimit(mock_app)
# Patch the algorithm selection to use our mock
with patch.object(algo_module, 'preregistered_algos', []):
stage.algo = mock_algo
result = await stage.process(sample_query, 'ReleaseRateLimitOccupancy')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query == sample_query
mock_algo.release_access.assert_called_once_with(sample_query, 'person', '12345')
+140
View File
@@ -0,0 +1,140 @@
"""
GroupRespondRuleCheckStage unit tests
Tests the actual GroupRespondRuleCheckStage implementation from pkg.pipeline.resprule
"""
import pytest
from unittest.mock import AsyncMock, Mock
from importlib import import_module
import langbot_plugin.api.entities.builtin.provider.session as provider_session
import langbot_plugin.api.entities.builtin.platform.message as platform_message
def get_modules():
"""Lazy import to ensure proper initialization order"""
# Import pipelinemgr first to trigger proper stage registration
# pipelinemgr = import_module('langbot.pkg.pipeline.pipelinemgr')
resprule = import_module('langbot.pkg.pipeline.resprule.resprule')
entities = import_module('langbot.pkg.pipeline.entities')
rule = import_module('langbot.pkg.pipeline.resprule.rule')
rule_entities = import_module('langbot.pkg.pipeline.resprule.entities')
return resprule, entities, rule, rule_entities
@pytest.mark.asyncio
async def test_person_message_skip(mock_app, sample_query):
"""Test person message skips rule check"""
resprule, entities, rule, rule_entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.PERSON
sample_query.pipeline_config = {'trigger': {'group-respond-rules': {}}}
stage = resprule.GroupRespondRuleCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
result = await stage.process(sample_query, 'GroupRespondRuleCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query == sample_query
@pytest.mark.asyncio
async def test_group_message_no_match(mock_app, sample_query):
"""Test group message with no matching rules"""
resprule, entities, rule, rule_entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.GROUP
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'group-respond-rules': {}}}
# Create mock rule matcher that doesn't match
mock_rule = Mock(spec=rule.GroupRespondRule)
mock_rule.match = AsyncMock(
return_value=rule_entities.RuleJudgeResult(matching=False, replacement=sample_query.message_chain)
)
stage = resprule.GroupRespondRuleCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
stage.rule_matchers = [mock_rule]
result = await stage.process(sample_query, 'GroupRespondRuleCheckStage')
assert result.result_type == entities.ResultType.INTERRUPT
assert result.new_query == sample_query
mock_rule.match.assert_called_once()
@pytest.mark.asyncio
async def test_group_message_match(mock_app, sample_query):
"""Test group message with matching rule"""
resprule, entities, rule, rule_entities = get_modules()
sample_query.launcher_type = provider_session.LauncherTypes.GROUP
sample_query.launcher_id = '12345'
sample_query.pipeline_config = {'trigger': {'group-respond-rules': {}}}
# Create new message chain after rule processing
new_chain = platform_message.MessageChain([platform_message.Plain(text='Processed message')])
# Create mock rule matcher that matches
mock_rule = Mock(spec=rule.GroupRespondRule)
mock_rule.match = AsyncMock(return_value=rule_entities.RuleJudgeResult(matching=True, replacement=new_chain))
stage = resprule.GroupRespondRuleCheckStage(mock_app)
await stage.initialize(sample_query.pipeline_config)
stage.rule_matchers = [mock_rule]
result = await stage.process(sample_query, 'GroupRespondRuleCheckStage')
assert result.result_type == entities.ResultType.CONTINUE
assert result.new_query == sample_query
assert sample_query.message_chain == new_chain
mock_rule.match.assert_called_once()
@pytest.mark.asyncio
async def test_atbot_rule_match(mock_app, sample_query):
"""Test AtBotRule removes At component"""
resprule, entities, rule, rule_entities = get_modules()
atbot_module = import_module('langbot.pkg.pipeline.resprule.rules.atbot')
sample_query.launcher_type = provider_session.LauncherTypes.GROUP
sample_query.adapter.bot_account_id = '999'
# Create message chain with At component
message_chain = platform_message.MessageChain(
[platform_message.At(target='999'), platform_message.Plain(text='Hello bot')]
)
sample_query.message_chain = message_chain
atbot_rule = atbot_module.AtBotRule(mock_app)
await atbot_rule.initialize()
result = await atbot_rule.match(str(message_chain), message_chain, {}, sample_query)
assert result.matching is True
# At component should be removed
assert len(result.replacement.root) == 1
assert isinstance(result.replacement.root[0], platform_message.Plain)
@pytest.mark.asyncio
async def test_atbot_rule_no_match(mock_app, sample_query):
"""Test AtBotRule when no At component present"""
resprule, entities, rule, rule_entities = get_modules()
atbot_module = import_module('langbot.pkg.pipeline.resprule.rules.atbot')
sample_query.launcher_type = provider_session.LauncherTypes.GROUP
sample_query.adapter.bot_account_id = '999'
# Create message chain without At component
message_chain = platform_message.MessageChain([platform_message.Plain(text='Hello')])
sample_query.message_chain = message_chain
atbot_rule = atbot_module.AtBotRule(mock_app)
await atbot_rule.initialize()
result = await atbot_rule.match(str(message_chain), message_chain, {}, sample_query)
assert result.matching is False
+616
View File
@@ -0,0 +1,616 @@
"""
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 get_plugin_diagnostics_module():
"""Lazy import for plugin diagnostic attribution helpers."""
return import_module('langbot.pkg.pipeline.plugin_diagnostics')
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
@pytest.mark.asyncio
async def test_message_chain_direct_append_consumes_pending_plugin_source(self):
"""MessageChain replies from earlier plugin events keep attribution."""
wrapper = get_wrapper_module()
app = FakeApp()
stage = wrapper.ResponseWrapper(app)
await stage.initialize(make_wrapper_config())
reply_chain = platform_message.MessageChain([platform_message.Plain(text='response')])
query = text_query('hello')
query.pipeline_config = make_wrapper_config()
query.resp_messages = [reply_chain]
query.resp_message_chain = []
plugin_diagnostics = get_plugin_diagnostics_module()
plugin_diagnostics.record_pending_plugin_response_source(
query,
reply_chain,
[
{
'kind': 'reply_message_chain',
'plugin': {'author': 'tester', 'name': 'demo'},
}
],
[{'manifest': {'metadata': {'author': 'observer', 'name': 'not-reply-source'}}}],
'PersonNormalMessageReceived',
)
results = []
async for result in stage.process(query, 'ResponseWrapper'):
results.append(result)
sources = plugin_diagnostics._get_response_sources(results[0].new_query, 0)
assert sources[0].plugin == {'author': 'tester', 'name': 'demo'}
assert sources[0].event_name == 'PersonNormalMessageReceived'
assert sources[0].is_approximate is False
assert '_plugin_response_sources' not in query.variables
assert '_plugin_pending_response_sources' not in query.variables
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()
app = FakeApp()
app.plugin_connector.emit_event = AsyncMock()
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)
assert results == []
assert query.resp_message_chain == []
app.plugin_connector.emit_event.assert_not_called()
@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)
assert len(results) == 2
for result in results:
assert result.result_type == entities.ResultType.CONTINUE
assert app.plugin_connector.emit_event.await_count == 2
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)
@pytest.mark.asyncio
async def test_custom_reply_records_plugin_source(self):
"""Plugin reply_message_chain should keep emitted plugin attribution."""
wrapper = get_wrapper_module()
app = FakeApp()
app.sess_mgr.get_session = AsyncMock(return_value=make_session())
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
mock_event_ctx._emitted_plugins = [
{
'manifest': {'metadata': {'author': 'observer', 'name': 'not-reply-source'}},
'plugin_config': {'token': 'secret-token'},
},
]
mock_event_ctx._response_sources = [
{
'kind': 'reply_message_chain',
'plugin': {'author': 'tester', 'name': 'demo'},
}
]
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 = []
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)
plugin_diagnostics = get_plugin_diagnostics_module()
sources = plugin_diagnostics._get_response_sources(results[0].new_query, 0)
assert sources[0].plugin == {'author': 'tester', 'name': 'demo'}
assert sources[0].event_name == 'NormalMessageResponded'
assert sources[0].is_approximate is False
assert 'secret-token' not in str(sources)
assert '_plugin_response_sources' not in query.variables
@pytest.mark.asyncio
async def test_custom_reply_falls_back_to_emitted_plugins_for_old_runtime(self):
"""Older plugin runtimes without response_sources keep approximate attribution."""
wrapper = get_wrapper_module()
app = FakeApp()
app.sess_mgr.get_session = AsyncMock(return_value=make_session())
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
mock_event_ctx._emitted_plugins = [
{'manifest': {'metadata': {'author': 'tester', 'name': 'demo'}}},
]
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 = []
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)
plugin_diagnostics = get_plugin_diagnostics_module()
sources = plugin_diagnostics._get_response_sources(results[0].new_query, 0)
assert sources[0].plugin == {'author': 'tester', 'name': 'demo'}
assert sources[0].is_approximate is True
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
@@ -0,0 +1,146 @@
"""Unit tests for ResponseWrapper outbound-attachment helpers.
Covers the sandbox -> user attachment path added for the Box attachment
round-trip:
* ``_is_final_assistant_message`` — only the terminal, tool-call-free assistant
message (or a final MessageChunk) should trigger collection.
* ``_append_outbound_attachments`` — collects sandbox outbox files exactly once
per query and maps each descriptor to the right platform component, swallowing
collection errors.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from langbot.pkg.pipeline.wrapper.wrapper import ResponseWrapper
def _make_wrapper(box_service) -> ResponseWrapper:
app = SimpleNamespace(logger=Mock())
wrapper = ResponseWrapper.__new__(ResponseWrapper)
wrapper.ap = app
return wrapper
def _make_query():
return SimpleNamespace(variables={})
def test_is_final_assistant_message_plain_assistant():
wrapper = _make_wrapper(box_service=None)
msg = provider_message.Message(role='assistant', content='done')
assert wrapper._is_final_assistant_message(msg) is True
def test_is_final_assistant_message_rejects_non_assistant():
wrapper = _make_wrapper(box_service=None)
msg = provider_message.Message(role='tool', content='{}')
assert wrapper._is_final_assistant_message(msg) is False
def test_is_final_assistant_message_rejects_tool_call_round():
wrapper = _make_wrapper(box_service=None)
msg = provider_message.Message(
role='assistant',
content='calling',
tool_calls=[
provider_message.ToolCall(
id='c1',
type='function',
function=provider_message.FunctionCall(name='exec', arguments='{}'),
)
],
)
assert wrapper._is_final_assistant_message(msg) is False
def test_is_final_assistant_message_non_final_chunk():
wrapper = _make_wrapper(box_service=None)
chunk = provider_message.MessageChunk(role='assistant', content='partial', is_final=False)
assert wrapper._is_final_assistant_message(chunk) is False
final_chunk = provider_message.MessageChunk(role='assistant', content='partial', is_final=True)
assert wrapper._is_final_assistant_message(final_chunk) is True
@pytest.mark.asyncio
async def test_append_outbound_attachments_maps_each_type():
box_service = SimpleNamespace(
available=True,
collect_outbound_attachments=AsyncMock(
return_value=[
{'type': 'Image', 'base64': 'data:image/png;base64,iVBORw0K'},
{'type': 'Voice', 'base64': 'data:audio/wav;base64,UklGRg=='},
{'type': 'File', 'name': 'report.xlsx', 'base64': 'data:app;base64,UEsDBA=='},
]
),
)
wrapper = _make_wrapper(box_service)
wrapper.ap.box_service = box_service
query = _make_query()
chain = platform_message.MessageChain([])
await wrapper._append_outbound_attachments(query, chain)
kinds = [type(c).__name__ for c in chain]
assert kinds == ['Image', 'Voice', 'File']
assert query.variables['_sandbox_outbound_collected'] is True
# File keeps its name
file_comp = chain[2]
assert getattr(file_comp, 'name', None) == 'report.xlsx'
@pytest.mark.asyncio
async def test_append_outbound_attachments_runs_once_per_query():
box_service = SimpleNamespace(
available=True,
collect_outbound_attachments=AsyncMock(return_value=[]),
)
wrapper = _make_wrapper(box_service)
wrapper.ap.box_service = box_service
query = _make_query()
query.variables['_sandbox_outbound_collected'] = True
chain = platform_message.MessageChain([])
await wrapper._append_outbound_attachments(query, chain)
box_service.collect_outbound_attachments.assert_not_awaited()
assert len(chain) == 0
@pytest.mark.asyncio
async def test_append_outbound_attachments_noop_without_box_service():
wrapper = _make_wrapper(box_service=None)
wrapper.ap.box_service = None
query = _make_query()
chain = platform_message.MessageChain([])
await wrapper._append_outbound_attachments(query, chain)
assert len(chain) == 0
# not marked collected, since service is unavailable
assert '_sandbox_outbound_collected' not in query.variables
@pytest.mark.asyncio
async def test_append_outbound_attachments_swallows_collection_error():
box_service = SimpleNamespace(
available=True,
collect_outbound_attachments=AsyncMock(side_effect=RuntimeError('boom')),
)
wrapper = _make_wrapper(box_service)
wrapper.ap.box_service = box_service
query = _make_query()
chain = platform_message.MessageChain([])
# must not raise
await wrapper._append_outbound_attachments(query, chain)
assert len(chain) == 0
wrapper.ap.logger.warning.assert_called_once()
@@ -0,0 +1,538 @@
import pytest
import aiocqhttp
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.platform.sources.aiocqhttp import (
AiocqhttpAdapter,
AiocqhttpEventConverter,
AiocqhttpMessageConverter,
)
async def _convert_single(component: platform_message.MessageComponent):
chain = platform_message.MessageChain([component])
message, _, _ = await AiocqhttpMessageConverter.yiri2target(chain)
return message[0]
@pytest.mark.asyncio
@pytest.mark.parametrize(
('payload', 'expected'),
[
('data:image/jpeg;base64,raw-image', 'base64://raw-image'),
('raw-image', 'base64://raw-image'),
('base64://raw-image', 'base64://raw-image'),
],
)
async def test_image_base64_payload_is_normalized(payload, expected):
segment = await _convert_single(platform_message.Image(base64=payload))
assert segment.type == 'image'
assert segment.data['file'] == expected
@pytest.mark.asyncio
async def test_voice_data_uri_base64_payload_is_normalized():
segment = await _convert_single(platform_message.Voice(base64='data:audio/wav;base64,raw-voice'))
assert segment.type == 'record'
assert segment.data['file'] == 'base64://raw-voice'
@pytest.mark.asyncio
@pytest.mark.parametrize(
('component', 'expected'),
[
(
platform_message.File(name='report.txt', base64='data:text/plain;base64,raw-file'),
{'file': 'base64://raw-file', 'name': 'report.txt'},
),
(
platform_message.File(name='report.txt', base64='raw-file'),
{'file': 'base64://raw-file', 'name': 'report.txt'},
),
(
platform_message.File(name='a.txt', url='http://example.com/a.txt'),
{'file': 'http://example.com/a.txt', 'name': 'a.txt'},
),
(
platform_message.File(name='a.txt', path='/tmp/a.txt'),
{'file': '/tmp/a.txt', 'name': 'a.txt'},
),
],
)
async def test_file_message_uses_available_file_source(component, expected):
segment = await _convert_single(component)
assert segment.type == 'file'
assert segment.data == expected
@pytest.mark.asyncio
async def test_forward_image_base64_payload_is_normalized():
forward = platform_message.Forward(
node_list=[
platform_message.ForwardMessageNode(
sender_id='10001',
sender_name='Tester',
message_chain=platform_message.MessageChain(
[platform_message.Image(base64='data:image/png;base64,raw-forward-image')]
),
)
]
)
messages = []
class Logger:
async def info(self, _message):
return None
async def error(self, _message):
return None
class Bot:
async def call_action(self, action, **kwargs):
assert action == 'send_forward_msg'
messages.append(kwargs)
platform = AiocqhttpAdapter.model_construct(
bot_account_id='10000',
config={},
logger=Logger(),
bot=Bot(),
)
await platform._send_forward_message(1000, forward)
assert messages[0]['messages'][0]['data']['content'][0] == {
'type': 'image',
'data': {'file': 'base64://raw-forward-image'},
}
@pytest.mark.asyncio
async def test_group_message_member_name_prefers_group_card():
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
'title': 'Special Title',
},
}
)
class Bot:
async def get_group_info(self, group_id):
assert group_id == 2000
return {'group_id': group_id, 'group_name': 'Test Group'}
converted = await AiocqhttpEventConverter().target2yiri(event, Bot())
assert converted.sender.member_name == 'Group Card'
assert converted.sender.group.id == 2000
assert converted.sender.group.name == 'Test Group'
assert converted.sender.special_title == 'Special Title'
@pytest.mark.asyncio
async def test_group_message_member_name_falls_back_to_nickname():
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': '',
'role': 'member',
},
}
)
converted = await AiocqhttpEventConverter().target2yiri(event)
assert converted.sender.member_name == 'QQ Nickname'
@pytest.mark.asyncio
async def test_group_message_special_title_uses_group_member_info_when_sender_title_is_empty():
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
'title': '',
},
}
)
class Bot:
async def get_group_info(self, group_id):
return {'group_id': group_id, 'group_name': 'Test Group'}
async def get_group_member_info(self, group_id, user_id):
assert group_id == 2000
assert user_id == 3000
return {'group_id': group_id, 'user_id': user_id, 'title': 'Member Title'}
converted = await AiocqhttpEventConverter().target2yiri(event, Bot())
assert converted.sender.special_title == 'Member Title'
@pytest.mark.asyncio
async def test_group_message_special_title_does_not_lookup_when_sender_title_exists():
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
'title': 'Event Title',
},
}
)
class Bot:
async def get_group_info(self, group_id):
return {'group_id': group_id, 'group_name': 'Test Group'}
async def get_group_member_info(self, group_id, user_id):
raise AssertionError('get_group_member_info should not be called')
converted = await AiocqhttpEventConverter().target2yiri(event, Bot())
assert converted.sender.special_title == 'Event Title'
@pytest.mark.asyncio
async def test_group_message_special_title_member_info_failure_is_cached(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
'title': '',
},
}
)
now = 1000.0
class Bot:
member_info_calls = 0
async def get_group_info(self, group_id):
return {'group_id': group_id, 'group_name': 'Test Group'}
async def get_group_member_info(self, group_id, user_id):
self.member_info_calls += 1
raise RuntimeError('api unavailable')
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: now)
bot = Bot()
converter = AiocqhttpEventConverter()
first = await converter.target2yiri(event, bot)
second = await converter.target2yiri(event, bot)
assert first.sender.special_title == ''
assert second.sender.special_title == ''
assert bot.member_info_calls == 1
@pytest.mark.asyncio
async def test_group_message_special_title_member_info_cache_expires(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
'title': '',
},
}
)
now = 1000.0
class Bot:
member_info_calls = 0
async def get_group_info(self, group_id):
return {'group_id': group_id, 'group_name': 'Test Group'}
async def get_group_member_info(self, group_id, user_id):
self.member_info_calls += 1
return {
'group_id': group_id,
'user_id': user_id,
'title': f'Member Title {self.member_info_calls}',
}
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: now)
bot = Bot()
converter = AiocqhttpEventConverter()
first = await converter.target2yiri(event, bot)
now = 87401.0
second = await converter.target2yiri(event, bot)
assert first.sender.special_title == 'Member Title 1'
assert second.sender.special_title == 'Member Title 2'
assert bot.member_info_calls == 2
@pytest.mark.asyncio
async def test_group_message_special_title_retries_after_negative_cache_expires(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
'title': '',
},
}
)
now = 1000.0
class Bot:
member_info_calls = 0
async def get_group_info(self, group_id):
return {'group_id': group_id, 'group_name': 'Test Group'}
async def get_group_member_info(self, group_id, user_id):
self.member_info_calls += 1
if self.member_info_calls == 1:
raise RuntimeError('api unavailable')
return {'group_id': group_id, 'user_id': user_id, 'title': 'Recovered Title'}
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: now)
bot = Bot()
converter = AiocqhttpEventConverter()
failed = await converter.target2yiri(event, bot)
now = 1601.0
recovered = await converter.target2yiri(event, bot)
assert failed.sender.special_title == ''
assert recovered.sender.special_title == 'Recovered Title'
assert bot.member_info_calls == 2
@pytest.mark.asyncio
async def test_group_message_group_name_is_cached(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
},
}
)
class Bot:
calls = 0
async def get_group_info(self, group_id):
self.calls += 1
assert group_id == 2000
return {'group_id': group_id, 'group_name': 'Cached Group'}
monotonic = 1000.0
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: monotonic)
bot = Bot()
converter = AiocqhttpEventConverter()
first = await converter.target2yiri(event, bot)
second = await converter.target2yiri(event, bot)
assert first.sender.group.name == 'Cached Group'
assert second.sender.group.name == 'Cached Group'
assert bot.calls == 1
@pytest.mark.asyncio
async def test_group_message_group_name_cache_expires(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
},
}
)
now = 1000.0
class Bot:
calls = 0
async def get_group_info(self, group_id):
self.calls += 1
return {'group_id': group_id, 'group_name': f'Group Name {self.calls}'}
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: now)
bot = Bot()
converter = AiocqhttpEventConverter()
first = await converter.target2yiri(event, bot)
now = 4601.0
second = await converter.target2yiri(event, bot)
assert first.sender.group.name == 'Group Name 1'
assert second.sender.group.name == 'Group Name 2'
assert bot.calls == 2
@pytest.mark.asyncio
async def test_group_message_group_name_uses_placeholder_when_lookup_fails(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
},
}
)
now = 1000.0
class Bot:
calls = 0
async def get_group_info(self, group_id):
self.calls += 1
raise RuntimeError('api unavailable')
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: now)
bot = Bot()
converter = AiocqhttpEventConverter()
converted = await converter.target2yiri(event, bot)
cached_failure = await converter.target2yiri(event, bot)
assert converted.sender.group.name == 'Group 2000'
assert cached_failure.sender.group.name == 'Group 2000'
assert bot.calls == 1
@pytest.mark.asyncio
async def test_group_message_group_name_retries_after_negative_cache_expires(monkeypatch):
event = aiocqhttp.Event(
{
'post_type': 'message',
'message_type': 'group',
'message_id': 1000,
'message': '',
'time': 1776491725,
'group_id': 2000,
'sender': {
'user_id': 3000,
'nickname': 'QQ Nickname',
'card': 'Group Card',
'role': 'member',
},
}
)
now = 1000.0
class Bot:
calls = 0
async def get_group_info(self, group_id):
self.calls += 1
if self.calls == 1:
raise RuntimeError('api unavailable')
return {'group_id': group_id, 'group_name': 'Recovered Group'}
monkeypatch.setattr('langbot.pkg.platform.sources.aiocqhttp.time.monotonic', lambda: now)
bot = Bot()
converter = AiocqhttpEventConverter()
failed = await converter.target2yiri(event, bot)
now = 1061.0
recovered = await converter.target2yiri(event, bot)
assert failed.sender.group.name == 'Group 2000'
assert recovered.sender.group.name == 'Recovered Group'
assert bot.calls == 2
@@ -0,0 +1,254 @@
"""Tests for DingTalk adapter helper behavior."""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from langbot.pkg.platform.sources.dingtalk import (
DingTalkAdapter,
_dingtalk_card_markdown,
_dingtalk_clean_form_content,
_dingtalk_completed_input_lines,
_dingtalk_extract_component_inputs,
_dingtalk_form_component_params,
_dingtalk_missing_completed_input_lines,
_dingtalk_pending_input_defs,
)
def test_dingtalk_select_component_params_expose_options():
params = _dingtalk_form_component_params(
{
'_current_input_field': 'choice',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
'inputs': {},
}
)
assert params['select_visible'] == 'true'
assert params['select_placeholder'] == 'choice'
assert params['select_options'] == ['A', 'B']
assert [option['value'] for option in params['index_o']] == ['A', 'B']
assert [option['value'] for option in params['test_index']] == ['A', 'B']
assert params['index_o'][0]['text']['zh_CN'] == 'A'
assert params['index_o'][0]['text']['en_US'] == 'A'
assert params['select_index'] == -1
def test_dingtalk_extract_select_from_builtin_result_dict():
inputs = _dingtalk_extract_component_inputs({'selectResult': {'index': 1, 'value': 'B'}})
assert inputs == {'select': 'B'}
def test_dingtalk_extract_select_from_template_param_string():
inputs = _dingtalk_extract_component_inputs({'select': '{"index": 1, "value": "B"}'})
assert inputs == {'select': '{"index": 1, "value": "B"}'}
def test_dingtalk_extract_input_and_select_together():
inputs = _dingtalk_extract_component_inputs(
{
'inputResult': {'value': 'looks good'},
'__built_in_selectResult__': {'index': 0, 'value': 'A'},
}
)
assert inputs == {'input': 'looks good', 'select': 'A'}
def test_dingtalk_extract_component_inputs_strips_card_line_endings():
inputs = _dingtalk_extract_component_inputs(
{
'inputResult': {'value': '回复我测试\r\n'},
'selectResult': {'value': '1\r'},
}
)
assert inputs == {'input': '回复我测试', 'select': '1'}
def test_dingtalk_pending_input_defs_includes_file_fields():
pending = _dingtalk_pending_input_defs(
{
'input_defs': [
{'output_variable_name': 'comment', 'type': 'paragraph'},
{'output_variable_name': 'files', 'type': 'file-list'},
],
'inputs': {'comment': 'ready'},
}
)
assert [field['output_variable_name'] for field in pending] == ['files']
def test_dingtalk_completed_input_lines_include_text_and_select_values():
lines = _dingtalk_completed_input_lines(
{
'all_input_defs': [
{'output_variable_name': 'comment', 'type': 'paragraph'},
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
},
],
'inputs': {'comment': 'looks good', 'choice': 'B'},
}
)
assert lines == ['✅ commentlooks good', '✅ choiceB']
def test_dingtalk_completed_inputs_are_not_repeated_when_already_interleaved():
form_data = {
'all_input_defs': [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{'output_variable_name': 'xiala', 'type': 'select'},
],
'inputs': {'us_input': '回复我测试\r', 'xiala': '1'},
}
form_content = '你好\n请输入你的问题\n✅ us_input:回复我测试\n请选择你的答案\n✅ xiala1'
assert _dingtalk_missing_completed_input_lines(form_data, form_content) == []
def test_dingtalk_completed_inputs_are_appended_when_template_does_not_render_them():
form_data = {
'all_input_defs': [{'output_variable_name': 'comment', 'type': 'paragraph'}],
'inputs': {'comment': 'ready'},
}
assert _dingtalk_missing_completed_input_lines(form_data, 'Please review') == ['✅ commentready']
def test_dingtalk_clean_form_content_uses_all_input_defs():
content = _dingtalk_clean_form_content(
{
'raw_form_content': 'Hello\n\n{{#$output.comment#}}\n\n{{#$output.choice#}}\n',
'input_defs': [],
'all_input_defs': [
{'output_variable_name': 'comment', 'type': 'paragraph'},
{'output_variable_name': 'choice', 'type': 'select'},
],
}
)
assert content == 'Hello'
def test_dingtalk_field_stage_keeps_prior_prompts_and_completed_values():
content = _dingtalk_clean_form_content(
{
'_current_input_field': 'choice',
'raw_form_content': ('Question\n{{#$output.comment#}}\nChoose an answer\n{{#$output.choice#}}'),
'form_content': 'Choose an answer',
'input_defs': [
{'output_variable_name': 'comment', 'type': 'paragraph'},
{'output_variable_name': 'choice', 'type': 'select'},
],
'inputs': {'comment': 'hello'},
}
)
assert '{{#$output.' not in content
assert content.index('Question') < content.index('comment')
assert content.index('comment') < content.index('Choose an answer')
def test_dingtalk_final_action_stage_interleaves_prompts_and_completed_values():
content = _dingtalk_clean_form_content(
{
'_action_select_only': True,
'raw_form_content': ('11\nQuestion\n{{#$output.comment#}}\nChoose an answer\n{{#$output.choice#}}'),
'all_input_defs': [
{'output_variable_name': 'comment', 'type': 'paragraph'},
{'output_variable_name': 'choice', 'type': 'select'},
],
'inputs': {'comment': 'hello', 'choice': 'B'},
}
)
assert '{{#$output.' not in content
assert content.startswith('11\nQuestion')
assert content.index('Question') < content.index('comment')
assert content.index('comment') < content.index('Choose an answer')
assert content.index('Choose an answer') < content.index('choice')
def test_dingtalk_card_markdown_preserves_internal_line_breaks():
assert _dingtalk_card_markdown('11\nQuestion\nCompleted') == '11<br>Question<br>Completed'
def _build_card_action_adapter() -> DingTalkAdapter:
adapter = DingTalkAdapter.model_construct(
card_state={
'card-1': {
'session_key': 'group_group-1',
'launcher_type': 'group',
'launcher_id': 'group-1',
'sender_user_id': 'initiator-1',
'form_token': 'token-1',
'workflow_run_id': 'run-1',
'actions': [{'id': 'approve', 'title': 'Approve'}],
'node_title': 'Review',
'form_content': 'Please review',
'input_defs': [],
'inputs': {},
}
},
active_turn_card={},
active_turn_text={},
)
adapter.logger = AsyncMock()
adapter.ap = MagicMock()
adapter.ap.platform_mgr.bots = []
adapter.ap.query_pool.add_query = AsyncMock()
return adapter
@pytest.mark.asyncio
async def test_dingtalk_group_card_action_uses_clicker_as_sender():
adapter = _build_card_action_adapter()
with patch.object(DingTalkAdapter, '_mark_card_resolved', new=AsyncMock()) as mark_resolved:
await adapter._on_card_action(
{
'out_track_id': 'card-1',
'user_id': 'reviewer-2',
'action_id': 'approve',
'params': {},
}
)
await asyncio.sleep(0)
call = adapter.ap.query_pool.add_query.await_args
assert call.kwargs['launcher_id'] == 'group-1'
assert call.kwargs['sender_id'] == 'reviewer-2'
assert call.kwargs['message_event'].sender.id == 'reviewer-2'
mark_resolved.assert_awaited_once()
@pytest.mark.asyncio
async def test_dingtalk_unknown_card_action_is_rejected():
adapter = _build_card_action_adapter()
await adapter._on_card_action(
{
'out_track_id': 'card-1',
'user_id': 'reviewer-2',
'action_id': 'not-on-card',
'params': {},
}
)
adapter.ap.query_pool.add_query.assert_not_awaited()
adapter.logger.warning.assert_awaited_once()
@@ -0,0 +1,42 @@
"""Tests for DingTalk API payload helpers."""
import json
from langbot.libs.dingtalk_api.api import _stringify_card_param_map
def test_dingtalk_card_param_map_stringifies_select_component_arrays():
params = _stringify_card_param_map(
{
'content': 'Pick one',
'btns': json.dumps([{'text': 'OK'}], ensure_ascii=False),
'select_options': ['A', 'B'],
'index_o': [
{
'value': 'A',
'text': {'zh_CN': 'A', 'en_US': 'A'},
}
],
'test_index': [
{
'value': 'A',
'text': {'zh_CN': 'A', 'en_US': 'A'},
}
],
'select_index': -1,
}
)
assert params['content'] == 'Pick one'
assert params['btns'] == '[{"text": "OK"}]'
assert params['select_options'] == '["A", "B"]'
assert json.loads(params['index_o'])[0]['value'] == 'A'
assert json.loads(params['test_index'])[0]['value'] == 'A'
assert params['select_index'] == '-1'
def test_dingtalk_card_param_map_stringifies_unregistered_structures():
params = _stringify_card_param_map({'other': ['A'], 'empty': None})
assert params['other'] == '["A"]'
assert params['empty'] == ''
@@ -0,0 +1,196 @@
"""Tests for Lark adapter helper behavior."""
from langbot.pkg.platform.sources.lark import (
LarkAdapter,
_lark_clean_form_content,
_lark_completed_input_lines,
_lark_current_input_defs,
_lark_extract_action_form_inputs,
_lark_should_update_stream_element,
_lark_visible_form_content,
)
def test_lark_current_input_defs_only_returns_active_stage():
input_defs = [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{'output_variable_name': 'xiala', 'type': 'select'},
]
assert _lark_current_input_defs(
{
'_current_input_field': 'xiala',
'input_defs': input_defs,
}
) == [input_defs[1]]
assert (
_lark_current_input_defs(
{
'_action_select_only': True,
'input_defs': input_defs,
}
)
== []
)
def test_lark_form_field_elements_only_render_active_stage():
adapter = LarkAdapter.model_construct()
form_data = {
'_current_input_field': 'xiala',
'input_defs': [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{
'output_variable_name': 'xiala',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['1', '2']},
},
],
}
elements, input_name_map, file_help_lines = adapter._build_lark_form_field_elements(form_data)
assert len(elements) == 1
assert elements[0]['tag'] == 'select_static'
assert elements[0]['label']['content'] == 'xiala'
assert list(input_name_map.values()) == ['xiala']
assert file_help_lines == []
def test_lark_form_stage_skips_closed_streaming_element_update():
assert not _lark_should_update_stream_element(
resume_from=False,
form_data={'_current_input_field': 'xiala'},
msg_seq=1,
is_final=True,
)
assert _lark_should_update_stream_element(
resume_from=False,
form_data=None,
msg_seq=1,
is_final=True,
)
def test_lark_final_action_stage_interleaves_prompts_and_completed_values():
form_content = _lark_visible_form_content(
{
'_action_select_only': True,
'raw_form_content': ('11\nQuestion\n{{#$output.us_input#}}\nChoose an answer\n{{#$output.xiala#}}\n'),
'all_input_defs': [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{'output_variable_name': 'xiala', 'type': 'select'},
],
'inputs': {'us_input': 'hello', 'xiala': '2'},
}
)
assert '{{#$output.' not in form_content
assert form_content.startswith('11\nQuestion')
assert form_content.index('Question') < form_content.index('us_input')
assert form_content.index('us_input') < form_content.index('Choose an answer')
assert form_content.index('Choose an answer') < form_content.index('xiala')
def test_lark_completed_input_lines_include_text_select_and_files():
lines = _lark_completed_input_lines(
{
'all_input_defs': [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{'output_variable_name': 'xiala', 'type': 'select'},
{'output_variable_name': 'files', 'type': 'file-list'},
],
'inputs': {
'us_input': '你好',
'xiala': 'or',
'files': [{'upload_file_id': 'file-1'}, {'upload_file_id': 'file-2'}],
},
}
)
assert lines == [
'✅ us_input:你好',
'✅ xialaor',
'✅ files2 file(s)',
]
def test_lark_clean_form_content_removes_all_input_placeholders():
content = _lark_clean_form_content(
'人工介入\n\n{{#$output.us_input#}}\n\n{{#$output.xiala#}}\n',
[
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{'output_variable_name': 'xiala', 'type': 'select'},
],
)
assert content == '人工介入'
def test_lark_extract_action_form_inputs_from_json_form_value():
class Action:
form_value = '{"Input_1_us_input_abcd12": "hello", "Select_2_xiala_abcd12": "B"}'
input_value = None
option = None
name = None
inputs = _lark_extract_action_form_inputs(
Action(),
{
'input_name_map': {
'Input_1_us_input_abcd12': 'us_input',
'Select_2_xiala_abcd12': 'xiala',
}
},
)
assert inputs == {'us_input': 'hello', 'xiala': 'B'}
def test_lark_extract_action_form_inputs_from_webhook_dict_action():
inputs = _lark_extract_action_form_inputs(
{
'form_value': {
'Input_1_us_input_abcd12': 'hello',
'Select_2_xiala_abcd12': {'value': 'B', 'text': {'content': 'Option B'}},
}
},
{
'input_name_map': {
'Input_1_us_input_abcd12': 'us_input',
'Select_2_xiala_abcd12': 'xiala',
}
},
)
assert inputs == {'us_input': 'hello', 'xiala': {'value': 'B', 'text': {'content': 'Option B'}}}
def test_lark_extract_action_form_inputs_maps_dotted_component_names():
inputs = _lark_extract_action_form_inputs(
{
'form_value': {
'Form_1_token_abcd12.Input_1_us_input_abcd12': 'hello',
}
},
{
'input_name_map': {
'Input_1_us_input_abcd12': 'us_input',
}
},
)
assert inputs == {'us_input': 'hello'}
def test_lark_completed_input_lines_display_select_value_from_object():
lines = _lark_completed_input_lines(
{
'all_input_defs': [
{'output_variable_name': 'xiala', 'type': 'select'},
],
'inputs': {'xiala': {'value': 'B', 'text': {'content': 'Option B'}}},
}
)
assert lines == ['✅ xialaB']
@@ -0,0 +1,224 @@
"""Tests for QQ Official keyboard payload helpers."""
import asyncio
import time
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.libs.qq_official_api.api import (
QQ_SELECT_ACTION_PREFIX,
build_keyboard_from_select_field,
get_select_field_options,
resolve_select_button_action,
)
def _select_form_data() -> dict:
return {
'_current_input_field': 'choice',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B', 'C']},
}
],
}
def test_qq_select_field_builds_callback_buttons():
keyboard = build_keyboard_from_select_field(_select_form_data(), buttons_per_row=2)
rows = keyboard['content']['rows']
assert [[button['render_data']['label'] for button in row['buttons']] for row in rows] == [
['A', 'B'],
['C'],
]
assert rows[0]['buttons'][0]['action']['data'] == f'{QQ_SELECT_ACTION_PREFIX}0'
assert rows[0]['buttons'][1]['action']['data'] == f'{QQ_SELECT_ACTION_PREFIX}1'
def test_qq_select_button_resolves_field_and_value():
form_data = _select_form_data()
assert get_select_field_options(form_data) == ('choice', ['A', 'B', 'C'])
assert resolve_select_button_action(form_data, f'{QQ_SELECT_ACTION_PREFIX}1') == ('choice', 'B')
assert resolve_select_button_action(form_data, f'{QQ_SELECT_ACTION_PREFIX}99') is None
def test_qq_select_keyboard_fits_twenty_five_options():
form_data = _select_form_data()
form_data['input_defs'][0]['option_source']['value'] = [f'Option {idx}' for idx in range(25)]
rows = build_keyboard_from_select_field(form_data)['content']['rows']
assert len(rows) == 5
assert all(len(row['buttons']) == 5 for row in rows)
def test_qq_non_select_field_does_not_build_keyboard():
form_data = {
'_current_input_field': 'comment',
'input_defs': [{'output_variable_name': 'comment', 'type': 'paragraph'}],
}
assert build_keyboard_from_select_field(form_data)['content']['rows'] == []
def _stream_test_adapter():
from langbot.pkg.platform.sources.qqofficial import QQOfficialAdapter
adapter = QQOfficialAdapter.model_construct()
adapter.logger = AsyncMock()
adapter.bot = MagicMock()
adapter.bot.send_stream_msg = AsyncMock(return_value={'id': 'stream-1'})
adapter.bot.send_markdown_keyboard = AsyncMock(return_value={'id': 'message-1'})
adapter.ap = None
adapter._stream_ctx = {}
adapter._stream_ctx_ts = {}
adapter._fallback_text = {}
adapter._fallback_text_ts = {}
return adapter
@pytest.mark.asyncio
async def test_qq_stream_uses_cumulative_chunks_as_snapshots():
adapter = _stream_test_adapter()
adapter._stream_ctx['message-1'] = {
'user_openid': 'user-1',
'msg_id': 'source-1',
'stream_msg_id': None,
'msg_seq': 1,
'index': 0,
'last_update_ts': 0,
'accumulated_text': '',
'sent_length': 0,
'session_started': False,
}
adapter._stream_ctx_ts['message-1'] = time.time()
source = MagicMock()
await adapter.reply_message_chunk(
source,
{'resp_message_id': 'message-1'},
platform_message.MessageChain([platform_message.Plain(text='<think>one')]),
)
await adapter.reply_message_chunk(
source,
{'resp_message_id': 'message-1'},
platform_message.MessageChain([platform_message.Plain(text='<think>one two')]),
is_final=True,
)
assert [call.kwargs['content'] for call in adapter.bot.send_stream_msg.await_args_list] == [
'<think>one',
' two',
]
@pytest.mark.asyncio
async def test_qq_non_streaming_fallback_keeps_latest_snapshot_only():
from langbot.pkg.platform.sources.qqofficial import QQOfficialAdapter
adapter = _stream_test_adapter()
source = MagicMock()
with patch.object(QQOfficialAdapter, 'reply_message', new=AsyncMock()) as reply_message:
await adapter.reply_message_chunk(
source,
{'resp_message_id': 'message-1'},
platform_message.MessageChain([platform_message.Plain(text='Hel')]),
)
await adapter.reply_message_chunk(
source,
{'resp_message_id': 'message-1'},
platform_message.MessageChain([platform_message.Plain(text='Hello')]),
is_final=True,
)
sent_chain = reply_message.await_args.args[1]
assert str(sent_chain) == 'Hello'
@pytest.mark.asyncio
async def test_qq_text_field_prompt_keeps_form_content():
from langbot.pkg.platform.sources.qqofficial import QQOfficialAdapter
adapter = _stream_test_adapter()
adapter._pending_forms = {}
adapter._session_event_ids = {}
adapter._anchor_msg_seq = {}
source = MagicMock()
source.d_id = 'source-1'
source.t = 'C2C_MESSAGE_CREATE'
event = MagicMock()
event.source_platform_object = source
event.sender.id = 'user-1'
form_data = {
'_current_input_field': 'us_input',
'node_title': 'Manual input',
'form_content': '1234\nEnter your question',
'input_defs': [{'output_variable_name': 'us_input', 'type': 'paragraph'}],
'actions': [{'id': 'yes', 'title': 'yes'}],
}
with patch.object(QQOfficialAdapter, '_resolve_target_from_event', return_value=('c2c', 'user-1')):
await adapter._handle_form_chunk(event, platform_message.MessageChain([]), form_data)
send_call = adapter.bot.send_markdown_keyboard.await_args.kwargs
assert send_call['markdown_content'] == '### Manual input\n\n1234\nEnter your question'
assert send_call['keyboard'] is None
@pytest.mark.asyncio
async def test_qq_select_click_enqueues_input_progress_query():
import langbot.pkg.core.app # noqa: F401
from langbot.pkg.platform.sources.qqofficial import QQOfficialAdapter
adapter = QQOfficialAdapter.model_construct()
adapter.logger = AsyncMock()
adapter.bot = MagicMock()
adapter.bot.ack_interaction = AsyncMock()
adapter.ap = MagicMock()
adapter.ap.platform_mgr.bots = []
adapter.ap.query_pool.add_query = AsyncMock()
adapter._pending_forms = {
'group_group-1': {
'form_data': {
**_select_form_data(),
'form_token': 'token-1',
'workflow_run_id': 'run-1',
'node_title': 'Review',
'actions': [{'id': 'approve', 'title': 'Approve'}],
},
'sender_id': 'initiator-1',
'posted_at': time.time(),
}
}
adapter._session_event_ids = {}
adapter._anchor_msg_seq = {}
await adapter._handle_interaction_create(
{
'id': 'interaction-1',
'chat_type': 1,
'group_openid': 'group-1',
'member_openid': 'reviewer-2',
'data': {'resolved': {'button_data': f'{QQ_SELECT_ACTION_PREFIX}1'}},
},
ws_event_id='event-1',
)
await asyncio.sleep(0)
call = adapter.ap.query_pool.add_query.await_args
form_action = call.kwargs['variables']['_dify_form_action']
assert call.kwargs['launcher_id'] == 'group-1'
assert call.kwargs['sender_id'] == 'reviewer-2'
assert form_action['action_id'] == ''
assert form_action['inputs'] == {'select': 'B'}
assert form_action['_current_input_field'] == 'choice'
assert form_action['_input_progress'] is True
adapter.bot.ack_interaction.assert_awaited_once_with('interaction-1', code=0)
@@ -0,0 +1,280 @@
"""
RuntimeBot.resolve_pipeline_uuid and _match_operator unit tests
"""
from unittest.mock import Mock
class TestMatchOperator:
"""Test the _match_operator static method."""
@staticmethod
def _get_class():
from langbot.pkg.platform.botmgr import RuntimeBot
return RuntimeBot
def test_eq(self):
cls = self._get_class()
assert cls._match_operator('hello', 'eq', 'hello') is True
assert cls._match_operator('hello', 'eq', 'world') is False
def test_neq(self):
cls = self._get_class()
assert cls._match_operator('hello', 'neq', 'world') is True
assert cls._match_operator('hello', 'neq', 'hello') is False
def test_contains(self):
cls = self._get_class()
assert cls._match_operator('hello world', 'contains', 'world') is True
assert cls._match_operator('hello world', 'contains', 'xyz') is False
def test_not_contains(self):
cls = self._get_class()
assert cls._match_operator('hello world', 'not_contains', 'xyz') is True
assert cls._match_operator('hello world', 'not_contains', 'world') is False
def test_starts_with(self):
cls = self._get_class()
assert cls._match_operator('hello world', 'starts_with', 'hello') is True
assert cls._match_operator('hello world', 'starts_with', 'world') is False
def test_regex(self):
cls = self._get_class()
assert cls._match_operator('hello123', 'regex', r'\d+') is True
assert cls._match_operator('hello', 'regex', r'\d+') is False
def test_regex_invalid_pattern(self):
cls = self._get_class()
assert cls._match_operator('hello', 'regex', r'[invalid') is False
def test_unknown_operator(self):
cls = self._get_class()
assert cls._match_operator('hello', 'unknown_op', 'hello') is False
class TestResolvePipelineUuid:
"""Test the resolve_pipeline_uuid method."""
@staticmethod
def _make_bot(default_pipeline: str, rules: list):
from langbot.pkg.platform.botmgr import RuntimeBot
bot_entity = Mock()
bot_entity.use_pipeline_uuid = default_pipeline
bot_entity.pipeline_routing_rules = rules
bot = object.__new__(RuntimeBot)
bot.bot_entity = bot_entity
return bot
def test_no_rules_returns_default(self):
bot = self._make_bot('default-uuid', [])
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi')
assert uuid == 'default-uuid'
assert routed is False
def test_none_rules_returns_default(self):
bot = self._make_bot('default-uuid', None)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi')
assert uuid == 'default-uuid'
assert routed is False
def test_launcher_type_match(self):
rules = [
{
'type': 'launcher_type',
'operator': 'eq',
'value': 'group',
'pipeline_uuid': 'group-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('group', '123', 'hi')
assert uuid == 'group-pipeline'
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi')
assert uuid == 'default-uuid'
assert routed is False
def test_launcher_id_match(self):
rules = [
{
'type': 'launcher_id',
'operator': 'eq',
'value': '12345',
'pipeline_uuid': 'vip-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '12345', 'hi')
assert uuid == 'vip-pipeline'
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '99999', 'hi')
assert uuid == 'default-uuid'
assert routed is False
def test_message_content_contains(self):
rules = [
{
'type': 'message_content',
'operator': 'contains',
'value': '紧急',
'pipeline_uuid': 'urgent-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', '这是紧急消息')
assert uuid == 'urgent-pipeline'
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '123', '普通消息')
assert uuid == 'default-uuid'
assert routed is False
def test_message_content_regex(self):
rules = [
{
'type': 'message_content',
'operator': 'regex',
'value': r'^/admin\b',
'pipeline_uuid': 'admin-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', '/admin help')
assert uuid == 'admin-pipeline'
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hello /admin')
assert uuid == 'default-uuid'
assert routed is False
def test_message_has_element_eq(self):
rules = [
{
'type': 'message_has_element',
'operator': 'eq',
'value': 'Image',
'pipeline_uuid': 'image-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi', ['Plain', 'Image'])
assert uuid == 'image-pipeline'
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi', ['Plain'])
assert uuid == 'default-uuid'
assert routed is False
def test_message_has_element_neq(self):
rules = [
{
'type': 'message_has_element',
'operator': 'neq',
'value': 'Image',
'pipeline_uuid': 'text-only-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi', ['Plain'])
assert uuid == 'text-only-pipeline'
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi', ['Plain', 'Image'])
assert uuid == 'default-uuid'
assert routed is False
def test_message_has_element_no_types_provided(self):
"""When element types are not provided, should not match."""
rules = [
{
'type': 'message_has_element',
'operator': 'eq',
'value': 'Image',
'pipeline_uuid': 'image-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi')
assert uuid == 'default-uuid'
assert routed is False
def test_first_match_wins(self):
rules = [
{
'type': 'launcher_type',
'operator': 'eq',
'value': 'group',
'pipeline_uuid': 'first-pipeline',
},
{
'type': 'launcher_type',
'operator': 'eq',
'value': 'group',
'pipeline_uuid': 'second-pipeline',
},
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('group', '123', 'hi')
assert uuid == 'first-pipeline'
assert routed is True
def test_skip_invalid_rules(self):
rules = [
{'type': '', 'operator': 'eq', 'value': 'x', 'pipeline_uuid': 'p1'},
{'type': 'launcher_type', 'operator': 'eq', 'value': 'person', 'pipeline_uuid': ''},
{'type': 'launcher_type', 'operator': 'eq', 'value': 'person', 'pipeline_uuid': 'valid'},
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi')
assert uuid == 'valid'
assert routed is True
def test_default_operator_is_eq(self):
rules = [
{
'type': 'launcher_type',
'value': 'person',
'pipeline_uuid': 'person-pipeline',
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'hi')
assert uuid == 'person-pipeline'
assert routed is True
def test_discard_pipeline(self):
"""When pipeline_uuid is __discard__, the message should be discarded."""
from langbot.pkg.platform.botmgr import RuntimeBot
rules = [
{
'type': 'message_content',
'operator': 'contains',
'value': 'spam',
'pipeline_uuid': RuntimeBot.PIPELINE_DISCARD,
}
]
bot = self._make_bot('default-uuid', rules)
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'this is spam')
assert uuid == RuntimeBot.PIPELINE_DISCARD
assert routed is True
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'normal message')
assert uuid == 'default-uuid'
assert routed is False
@@ -0,0 +1,158 @@
"""Tests for Telegram Dify form callback helpers."""
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
from telegram import ForceReply
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.platform.sources.telegram import (
TelegramAdapter,
_telegram_form_action_from_callback,
_telegram_select_field_options,
)
def _select_form_data() -> dict:
return {
'_current_input_field': 'choice',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B', 'C']},
}
],
}
def test_telegram_select_field_options_are_extracted():
assert _telegram_select_field_options(_select_form_data()) == ('choice', ['A', 'B', 'C'])
def test_telegram_select_callback_becomes_input_progress():
assert _telegram_form_action_from_callback({'f': 1, 'x': 1}) == {
'action_id': '',
'inputs': {'select': {'index': 1}},
'_input_progress': True,
}
def test_telegram_action_callback_remains_final_action():
assert _telegram_form_action_from_callback({'f': 1, 'a': 'approve'}) == {
'action_id': 'approve',
'inputs': {},
}
def test_telegram_invalid_select_callback_is_rejected():
assert _telegram_form_action_from_callback({'f': 1, 'x': -1}) is None
assert _telegram_form_action_from_callback({'f': 1, 'x': 'invalid'}) is None
def test_telegram_form_callback_cache_consumes_the_whole_form_group():
adapter = TelegramAdapter.model_construct()
adapter._form_action_titles = {}
adapter._cache_form_action_titles({'callback-a': 'A', 'callback-b': 'B'}, now=100.0)
assert adapter._take_form_action_title('callback-a', now=101.0) == 'A'
assert adapter._take_form_action_title('callback-a', now=101.0) is None
assert adapter._take_form_action_title('callback-b', now=101.0) is None
assert adapter._form_action_titles == {}
def test_telegram_form_callback_cache_prunes_expired_entries():
adapter = TelegramAdapter.model_construct()
adapter._form_action_titles = {}
adapter._cache_form_action_titles({'callback-a': 'A'}, now=100.0)
assert adapter._take_form_action_title('callback-a', now=100.0 + adapter._FORM_ACTION_CACHE_TTL) is None
assert adapter._form_action_titles == {}
def test_telegram_form_callback_cache_preserves_pipeline_uuid():
adapter = TelegramAdapter.model_construct()
adapter._form_action_titles = {}
adapter._cache_form_action_titles(
{'callback-a': 'Approve'},
pipeline_uuid='pipeline-routed',
now=100.0,
)
assert adapter._take_form_action_context('callback-a', now=101.0) == (
'Approve',
'pipeline-routed',
)
@pytest.mark.asyncio
async def test_telegram_select_field_sends_two_column_inline_keyboard():
bot = MagicMock()
bot.send_message = AsyncMock()
adapter = TelegramAdapter.model_construct(bot=bot, config={}, msg_stream_id={}, seq=1, listeners={})
adapter._form_action_titles = {}
update = MagicMock()
update.effective_chat.id = 123
update.effective_message.message_thread_id = None
event = platform_events.FriendMessage(
sender=platform_entities.Friend(id='user-1', nickname='', remark=''),
message_chain=platform_message.MessageChain([]),
source_platform_object=update,
)
form_data = {
**_select_form_data(),
'node_title': 'Review',
'form_content': 'Choose one',
'workflow_run_id': 'workflow-run-12345678',
'actions': [{'id': 'approve', 'title': 'Approve'}],
}
await adapter._send_form_action_buttons(event, form_data)
args = bot.send_message.await_args.kwargs
rows = args['reply_markup'].inline_keyboard
assert [[button.text for button in row] for row in rows] == [['A', 'B'], ['C']]
callback_data = rows[0][1].callback_data
assert len(callback_data.encode('utf-8')) <= 64
assert json.loads(callback_data)['x'] == 1
assert callback_data in adapter._form_action_titles
@pytest.mark.asyncio
async def test_telegram_text_field_does_not_show_action_buttons():
bot = MagicMock()
bot.send_message = AsyncMock()
adapter = TelegramAdapter.model_construct(bot=bot, config={}, msg_stream_id={}, seq=1, listeners={})
adapter._form_action_titles = {}
update = MagicMock()
update.effective_chat.id = 123
update.effective_message.message_thread_id = None
event = platform_events.FriendMessage(
sender=platform_entities.Friend(id='user-1', nickname='', remark=''),
message_chain=platform_message.MessageChain([]),
source_platform_object=update,
)
form_data = {
'_current_input_field': 'us_input',
'input_defs': [{'output_variable_name': 'us_input', 'type': 'paragraph'}],
'node_title': '人工介入',
'form_content': 'us_input (paragraph): reply "us_input: <value>"',
'workflow_run_id': 'workflow-run-12345678',
'actions': [{'id': 'yes', 'title': 'yes'}, {'id': 'no', 'title': 'no'}],
}
await adapter._send_form_action_buttons(event, form_data)
args = bot.send_message.await_args.kwargs
assert isinstance(args['reply_markup'], ForceReply)
assert args['reply_markup'].selective is False
assert args['reply_markup'].input_field_placeholder == 'us_input'
assert 'Please reply' not in args['text']
assert args['text'].startswith('[人工介入]')
assert 'us_input (paragraph)' in args['text']
assert adapter._form_action_titles == {}
@@ -0,0 +1,92 @@
"""Unit tests for WebSocketAdapter._process_image_components.
The web debug client uploads Image / Voice / File components carrying a storage
key in ``path``. This helper resolves each to a base64 data URI (so multimodal
LLM input and the Box sandbox inbox have usable bytes), then deletes the
consumed storage object and clears ``path``. Covers mimetype selection per
type and graceful error handling.
"""
from __future__ import annotations
import base64
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.platform.sources.websocket_adapter import WebSocketAdapter
def _make_adapter(load_return=b'hello', load_side_effect=None):
provider = Mock()
provider.load = AsyncMock(return_value=load_return, side_effect=load_side_effect)
provider.delete = AsyncMock()
ap = Mock()
ap.storage_mgr.storage_provider = provider
logger = Mock()
logger.error = AsyncMock()
# WebSocketAdapter is a pydantic model; bypass full __init__/validation.
adapter = WebSocketAdapter.model_construct(ap=ap, logger=logger)
return adapter, provider
@pytest.mark.asyncio
async def test_image_jpeg_mimetype_and_cleanup():
adapter, provider = _make_adapter(load_return=b'\xff\xd8\xff')
chain = [{'type': 'Image', 'path': 'storage://abc/photo.jpg'}]
await adapter._process_image_components(chain)
expected_b64 = base64.b64encode(b'\xff\xd8\xff').decode('utf-8')
assert chain[0]['base64'] == f'data:image/jpeg;base64,{expected_b64}'
assert chain[0]['path'] == '' # consumed
provider.delete.assert_awaited_once_with('storage://abc/photo.jpg')
@pytest.mark.asyncio
async def test_image_defaults_to_png():
adapter, _ = _make_adapter()
chain = [{'type': 'Image', 'path': 'storage://abc/blob'}]
await adapter._process_image_components(chain)
assert chain[0]['base64'].startswith('data:image/png;base64,')
@pytest.mark.asyncio
async def test_voice_uses_guessed_or_wav_mimetype():
adapter, _ = _make_adapter()
chain = [{'type': 'Voice', 'path': 'storage://abc/clip.wav'}]
await adapter._process_image_components(chain)
assert chain[0]['base64'].startswith('data:audio/')
@pytest.mark.asyncio
async def test_file_uses_octet_stream_fallback():
adapter, _ = _make_adapter()
chain = [{'type': 'File', 'path': 'storage://abc/unknownblob'}]
await adapter._process_image_components(chain)
assert chain[0]['base64'].startswith('data:application/octet-stream;base64,')
@pytest.mark.asyncio
async def test_skips_components_without_path_or_unknown_type():
adapter, provider = _make_adapter()
chain = [
{'type': 'Image', 'path': ''}, # no path
{'type': 'Plain', 'path': 'storage://abc/x'}, # not a file component
{'type': 'At', 'target': '123'}, # no path key at all
]
await adapter._process_image_components(chain)
provider.load.assert_not_awaited()
assert 'base64' not in chain[0]
assert 'base64' not in chain[1]
@pytest.mark.asyncio
async def test_load_failure_is_logged_not_raised():
adapter, _ = _make_adapter(load_side_effect=RuntimeError('storage down'))
chain = [{'type': 'File', 'path': 'storage://abc/doc.pdf'}]
# must not raise
await adapter._process_image_components(chain)
assert 'base64' not in chain[0]
adapter.logger.error.assert_awaited_once()
@@ -0,0 +1,475 @@
import sys
import types
import pytest
logger_module = types.ModuleType('langbot.pkg.platform.logger')
logger_module.EventLogger = object
sys.modules.setdefault('langbot.pkg.platform.logger', logger_module)
from langbot.libs.wecom_ai_bot_api.api import ( # noqa: E402
WecomBotClient,
build_button_interaction_payload,
build_human_input_template_card_payload,
build_button_interaction_update_card,
build_multiple_interaction_update_card,
extract_template_card_event_payload,
extract_template_card_selections,
extract_wecom_event_type,
extract_template_card_action,
build_human_input_text_prompt,
parse_select_button_action,
)
from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient # noqa: E402
def test_extract_template_card_action_supports_nested_button_key():
task_id, event_key, card_type = extract_template_card_action(
{
'taskId': 'task-1',
'cardType': 'button_interaction',
'button': {'key': 'approve'},
}
)
assert task_id == 'task-1'
assert event_key == 'approve'
assert card_type == 'button_interaction'
def test_extract_wecom_event_type_supports_top_level_template_card_event():
payload = {
'eventtype': 'template_card_event',
'template_card_event': {
'TaskId': 'task-1',
'CardType': 'multiple_interaction',
'ResponseData': '{"select_list":[{"question_key":"choice","option_id":"opt_2"}]}',
},
}
assert extract_wecom_event_type(payload) == 'template_card_event'
assert extract_template_card_event_payload(payload)['TaskId'] == 'task-1'
def test_extract_wecom_event_type_infers_template_card_event_from_top_level_card_fields():
payload = {
'TaskId': 'task-1',
'CardType': 'button_interaction',
'EventKey': 'approve',
}
assert extract_wecom_event_type(payload) == 'template_card_event'
assert extract_template_card_event_payload(payload)['EventKey'] == 'approve'
def test_build_button_interaction_update_card_marks_clicked_button():
card = build_button_interaction_update_card(
{
'node_title': 'Manual Review',
'form_content': 'Please choose one action.',
'actions': [
{'id': 'approve', 'title': 'Approve', 'button_style': 'primary'},
{'id': 'reject', 'title': 'Reject', 'button_style': 'danger'},
],
},
task_id='task-1',
action_id='reject',
source={'desc': 'LangBot'},
)
assert card['main_title'] == {'title': 'Manual Review'}
assert card['sub_title_text'] == 'Please choose one action.'
assert card['button_list'][0] == {'text': 'Approve', 'style': 2, 'key': 'approve'}
assert card['button_list'][1] == {
'text': '✅ Reject',
'style': 1,
'key': 'reject',
'replace_text': '✅ Reject',
}
assert card['source'] == {'desc': 'LangBot'}
def test_build_button_interaction_payload_uses_preselected_button_styles_before_click():
payload = build_button_interaction_payload(
{
'node_title': 'Manual Review',
'actions': [
{'id': 'approve', 'title': 'Approve', 'button_style': 'primary'},
{'id': 'reject', 'title': 'Reject', 'button_style': 'danger'},
],
},
task_id='task-1',
)
assert payload['template_card']['button_list'] == [
{'text': 'Approve', 'style': 2, 'key': 'approve'},
{'text': 'Reject', 'style': 2, 'key': 'reject'},
]
def test_build_payload_uses_multiple_interaction_for_pending_select_field():
payload = build_human_input_template_card_payload(
{
'node_title': 'Manual Review',
'form_content': 'Choose a label\n\n{{#$output.choice#}}',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
'inputs': {},
'actions': [{'id': 'yes', 'title': 'Yes'}],
},
task_id='task-1',
source={'desc': 'LangBot'},
)
card = payload['template_card']
assert card['card_type'] == 'multiple_interaction'
assert card['source'] == {'desc': 'LangBot'}
assert card['select_list'] == [
{
'question_key': 'choice',
'title': 'choice',
'selected_id': 'opt_1',
'option_list': [
{'id': 'opt_1', 'text': 'A'},
{'id': 'opt_2', 'text': 'B'},
],
}
]
assert card['submit_button'] == {'text': 'Submit', 'key': 'submit_human_input'}
def test_build_payload_can_emulate_select_as_buttons_for_wecombot_ws():
form_data = {
'node_title': 'Manual Review',
'form_content': 'Choose a label\n\n{{#$output.choice#}}',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
'inputs': {},
'actions': [{'id': 'yes', 'title': 'Yes'}],
}
payload = build_human_input_template_card_payload(
form_data,
task_id='task-1',
select_as_buttons=True,
)
card = payload['template_card']
assert card['card_type'] == 'button_interaction'
assert card['button_list'][0]['text'] == 'A'
assert parse_select_button_action(card['button_list'][1]['key'], form_data) == {'choice': 'B'}
def test_text_input_card_uses_current_stage_content_without_direct_reply_prompt():
payload = build_human_input_template_card_payload(
{
'node_title': '人工介入',
'form_content': '11\n请输入你的问题\n\n{{#$output.us_input#}}',
'raw_form_content': ('11\n请输入你的问题\n{{#$output.us_input#}}\n请选择你的答案\n{{#$output.xiala#}}'),
'input_defs': [
{
'output_variable_name': 'us_input',
'type': 'paragraph',
'label': '请输入你的问题',
}
],
'inputs': {},
'actions': [{'id': 'yes', 'title': 'yes'}],
'_current_input_field': 'us_input',
},
task_id='task-1',
)
card = payload['template_card']
assert 'desc' not in card['main_title']
assert card['sub_title_text'] == '11\n请输入你的问题'
assert card['button_list'] == []
def test_build_human_input_text_prompt_for_current_text_field():
prompt = build_human_input_text_prompt(
{
'node_title': '人工介入',
'form_content': '11\n请输入你的问题\n{{#$output.us_input#}}',
'raw_form_content': ('11\n请输入你的问题\n{{#$output.us_input#}}\n请选择你的答案\n{{#$output.xiala#}}'),
'input_defs': [
{
'output_variable_name': 'us_input',
'type': 'paragraph',
'label': '请输入你的问题',
}
],
'_current_input_field': 'us_input',
}
)
assert prompt == '人工介入\n\n11\n请输入你的问题'
@pytest.mark.asyncio
async def test_ws_push_form_pause_sends_text_prompt_without_empty_card():
client = WecomBotWsClient('bot-id', 'secret', object())
client._stream_ids['msg-1'] = 'req-1|stream-1'
client._stream_sessions['msg-1'] = {'user_id': 'user-1'}
sent = []
async def fake_reply_text(req_id, content):
sent.append((req_id, content))
return {}
client.reply_text = fake_reply_text
ok, stream_id, task_id = await client.push_form_pause(
'msg-1',
{
'node_title': '人工介入',
'form_content': '11\n请输入你的问题\n{{#$output.us_input#}}',
'raw_form_content': ('11\n请输入你的问题\n{{#$output.us_input#}}\n请选择你的答案\n{{#$output.xiala#}}'),
'input_defs': [
{
'output_variable_name': 'us_input',
'type': 'paragraph',
'label': '请输入你的问题',
}
],
'_current_input_field': 'us_input',
},
)
assert ok is True
assert stream_id == 'stream-1'
assert task_id is None
assert sent == [('req-1', '人工介入\n\n11\n请输入你的问题')]
assert client._pending_forms_by_task == {}
assert 'msg-1' not in client._stream_ids
@pytest.mark.asyncio
async def test_ws_stream_sends_cumulative_snapshots_to_wecom():
client = WecomBotWsClient('bot-id', 'secret', object())
client._stream_ids['msg-1'] = 'req-1|stream-1'
client._stream_sessions['msg-1'] = {}
sent = []
async def fake_reply_stream(req_id, stream_id, content, finish=False, feedback_id=''):
sent.append((req_id, stream_id, content, finish))
return {}
client.reply_stream = fake_reply_stream
assert await client.push_stream_chunk('msg-1', '', is_final=False)
assert await client.push_stream_chunk('msg-1', '你好', is_final=False)
assert await client.push_stream_chunk('msg-1', '你好', is_final=True)
assert sent == [
('req-1', 'stream-1', '', False),
('req-1', 'stream-1', '你好', False),
('req-1', 'stream-1', '你好', True),
]
@pytest.mark.asyncio
async def test_webhook_stream_queues_cumulative_snapshots_for_followups():
client = WecomBotClient('', '', '', object(), unified_mode=True)
session, _ = client.stream_sessions.create_or_get({'msgid': 'msg-1', 'chatid': '', 'from': {'userid': 'user-1'}})
assert await client.push_stream_chunk('msg-1', '', is_final=False)
assert await client.push_stream_chunk('msg-1', '你好', is_final=False)
assert await client.push_stream_chunk('msg-1', '你好', is_final=True)
chunks = [
await client.stream_sessions.consume(session.stream_id),
await client.stream_sessions.consume(session.stream_id),
await client.stream_sessions.consume(session.stream_id),
]
assert [(chunk.content, chunk.is_final) for chunk in chunks] == [
('', False),
('你好', False),
('你好', True),
]
def test_human_input_payload_keeps_action_select_stage_as_buttons():
payload = build_human_input_template_card_payload(
{
'node_title': 'Manual Review',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
'inputs': {'choice': 'B'},
'actions': [
{'id': 'approve', 'title': 'Approve'},
{'id': 'reject', 'title': 'Reject'},
],
'_action_select_only': True,
},
task_id='task-1',
)
card = payload['template_card']
assert card['card_type'] == 'button_interaction'
assert card['button_list'] == [
{'text': 'Approve', 'style': 2, 'key': 'approve'},
{'text': 'Reject', 'style': 2, 'key': 'reject'},
]
def test_extract_template_card_selections_maps_selected_id_to_option_text():
selections = extract_template_card_selections(
{
'SelectedItems': [
{'QuestionKey': 'choice', 'SelectedId': 'opt_2'},
],
},
{
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
},
)
assert selections == {'choice': 'B'}
def test_extract_template_card_selections_reads_nested_response_data_json():
selections = extract_template_card_selections(
{
'CardType': 'multiple_interaction',
'ResponseData': '{"select_list":[{"question_key":"choice","option_id":"opt_2"}]}',
},
{
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
},
)
assert selections == {'choice': 'B'}
def test_extract_template_card_selections_reads_response_data_direct_mapping():
selections = extract_template_card_selections(
{
'CardType': 'multiple_interaction',
'EventKey': 'submit_human_input',
'ResponseData': '{"choice":"opt_2"}',
},
{
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
},
)
assert selections == {'choice': 'B'}
def test_build_multiple_interaction_update_card_disables_selected_value_without_submitted_text():
card = build_multiple_interaction_update_card(
{
'node_title': 'Manual Review',
'form_content': 'Choose a label\n{{#$output.choice#}}',
'raw_form_content': 'Choose a label\n{{#$output.choice#}}',
'input_defs': [
{
'output_variable_name': 'choice',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['A', 'B']},
}
],
'_current_input_field': 'choice',
},
task_id='task-1',
selections={'choice': 'B'},
)
assert card['card_type'] == 'multiple_interaction'
assert card['main_title']['desc'] == 'Choose a label\n✅ choiceB'
assert card['submit_button']['text'] == ''
assert card['select_list'][0]['disable'] is True
assert card['select_list'][0]['selected_id'] == 'opt_2'
def test_select_stage_only_shows_current_prompt_in_a_separate_message():
raw_content = '11\n请输入你的问题\n{{#$output.us_input#}}\n请选择你的答案\n{{#$output.xiala#}}'
payload = build_human_input_template_card_payload(
{
'node_title': '人工介入',
'form_content': '请选择你的答案\n{{#$output.xiala#}}',
'raw_form_content': raw_content,
'input_defs': [
{
'output_variable_name': 'xiala',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['1', '2']},
}
],
'all_input_defs': [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{
'output_variable_name': 'xiala',
'type': 'select',
'option_source': {'type': 'constant', 'value': ['1', '2']},
},
],
'inputs': {'us_input': '你叫啥'},
'_current_input_field': 'xiala',
},
task_id='task-1',
)
assert payload['template_card']['main_title']['desc'] == '请选择你的答案'
def test_action_stage_only_shows_content_after_fields_without_placeholders():
raw_content = '11\n请输入你的问题\n{{#$output.us_input#}}\n请选择你的答案\n{{#$output.xiala#}}\n请选择操作'
payload = build_human_input_template_card_payload(
{
'node_title': '人工介入',
'form_content': raw_content,
'raw_form_content': raw_content,
'input_defs': [],
'all_input_defs': [
{'output_variable_name': 'us_input', 'type': 'paragraph'},
{'output_variable_name': 'xiala', 'type': 'select'},
],
'inputs': {'us_input': '你叫啥', 'xiala': '2'},
'actions': [
{'id': 'yes', 'title': 'yes'},
{'id': 'no', 'title': 'no'},
],
'_action_select_only': True,
},
task_id='task-1',
)
card = payload['template_card']
assert card['sub_title_text'] == '请选择操作'
assert '{{#$output.' not in card['sub_title_text']
assert [button['text'] for button in card['button_list']] == ['yes', 'no']
@@ -0,0 +1,91 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.platform.sources.wecomcs import WecomCSAdapter
class DummyLogger(abstract_platform_logger.AbstractEventLogger):
async def info(self, *args, **kwargs):
pass
async def debug(self, *args, **kwargs):
pass
async def warning(self, *args, **kwargs):
pass
async def error(self, *args, **kwargs):
pass
def make_adapter():
return WecomCSAdapter(
config={
'corpid': 'corp-id',
'secret': 'secret',
'token': 'token',
'EncodingAESKey': 'encoding-key',
},
logger=DummyLogger(),
)
@pytest.mark.asyncio
async def test_send_message_sends_text_to_customer_service_user():
adapter = make_adapter()
adapter.bot_account_id = 'kf-test'
adapter.bot = SimpleNamespace(send_text_msg=AsyncMock())
message = platform_message.MessageChain([platform_message.Plain(text='hello')])
await adapter.send_message('person', 'uexternal-user', message)
adapter.bot.send_text_msg.assert_awaited_once()
kwargs = adapter.bot.send_text_msg.await_args.kwargs
assert kwargs['open_kfid'] == 'kf-test'
assert kwargs['external_userid'] == 'external-user'
assert kwargs['content'] == 'hello'
assert kwargs['msgid'].startswith('langbot_')
@pytest.mark.asyncio
async def test_send_message_allows_explicit_open_kfid_in_target_id():
adapter = make_adapter()
adapter.bot = SimpleNamespace(send_text_msg=AsyncMock())
message = platform_message.MessageChain([platform_message.Plain(text='hello')])
await adapter.send_message('person', 'kf-explicit|uexternal-user', message)
kwargs = adapter.bot.send_text_msg.await_args.kwargs
assert kwargs['open_kfid'] == 'kf-explicit'
assert kwargs['external_userid'] == 'external-user'
@pytest.mark.asyncio
async def test_send_message_requires_open_kfid():
adapter = make_adapter()
adapter.bot = SimpleNamespace(send_text_msg=AsyncMock())
message = platform_message.MessageChain([platform_message.Plain(text='hello')])
with pytest.raises(ValueError, match='open_kfid is required'):
await adapter.send_message('person', 'uexternal-user', message)
adapter.bot.send_text_msg.assert_not_called()
@pytest.mark.asyncio
async def test_send_message_rejects_group_targets():
adapter = make_adapter()
adapter.bot_account_id = 'kf-test'
adapter.bot = SimpleNamespace(send_text_msg=AsyncMock())
message = platform_message.MessageChain([platform_message.Plain(text='hello')])
with pytest.raises(ValueError, match='only supports sending messages to person'):
await adapter.send_message('group', 'group-id', message)
adapter.bot.send_text_msg.assert_not_called()
+1
View File
@@ -0,0 +1 @@
# Plugin connector unit tests
@@ -0,0 +1,616 @@
"""Unit tests for plugin connector methods.
Tests cover:
- list_plugins() with filtering and sorting
- list_knowledge_engines() and list_parsers()
- RAG methods (ingest, retrieve, schema)
- Disabled plugin early returns
"""
from __future__ import annotations
import pytest
from unittest.mock import Mock, AsyncMock
from importlib import import_module
from tests.factories import text_query
def get_connector_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.plugin.connector')
def create_mock_app():
"""Create mock Application for testing."""
mock_app = Mock()
mock_app.logger = Mock()
mock_app.instance_config = Mock()
mock_app.instance_config.data = {'plugin': {'enable': True}}
mock_app.persistence_mgr = AsyncMock()
mock_app.persistence_mgr.execute_async = AsyncMock()
return mock_app
def create_mock_connector():
"""Create mock PluginRuntimeConnector instance for testing."""
connector = get_connector_module()
async def mock_disconnect_callback(conn):
pass
return connector.PluginRuntimeConnector(create_mock_app(), mock_disconnect_callback)
class TestListPlugins:
"""Tests for list_plugins method."""
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test returns empty list when plugin system disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.list_plugins()
assert result == []
@pytest.mark.asyncio
async def test_calls_handler_list_plugins(self):
"""Test that handler.list_plugins is called."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.list_plugins = AsyncMock(
return_value=[{'manifest': {'manifest': {'metadata': {'author': 'test', 'name': 'plugin'}}}}]
)
result = await connector.list_plugins()
connector.handler.list_plugins.assert_called_once()
assert result == [{'manifest': {'manifest': {'metadata': {'author': 'test', 'name': 'plugin'}}}}]
@pytest.mark.asyncio
async def test_filters_by_component_kinds(self):
"""Test that plugins are filtered by component kinds."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.list_plugins = AsyncMock(
return_value=[
{
'manifest': {'manifest': {'metadata': {'author': 'a', 'name': 'p1'}}},
'components': [{'manifest': {'manifest': {'kind': 'Command'}}}],
'debug': False,
},
{
'manifest': {'manifest': {'metadata': {'author': 'b', 'name': 'p2'}}},
'components': [{'manifest': {'manifest': {'kind': 'Tool'}}}],
'debug': False,
},
]
)
result = await connector.list_plugins(component_kinds=['Command'])
assert len(result) == 1
assert result[0]['manifest']['manifest']['metadata']['name'] == 'p1'
@pytest.mark.asyncio
async def test_sorts_debug_plugins_first(self):
"""Test that debug plugins are sorted first."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.list_plugins = AsyncMock(
return_value=[
{
'manifest': {'manifest': {'metadata': {'author': 'a', 'name': 'normal'}}},
'components': [],
'debug': False,
},
{
'manifest': {'manifest': {'metadata': {'author': 'b', 'name': 'debug'}}},
'components': [],
'debug': True,
},
]
)
connector.ap.persistence_mgr.execute_async = AsyncMock(return_value=Mock(__iter__=lambda self: iter([])))
result = await connector.list_plugins()
# Debug plugin should be first
assert result[0]['debug'] is True
class TestPluginDiagnostics:
@pytest.mark.asyncio
async def test_emit_event_preserves_response_sources(self):
connector = create_mock_connector()
query = text_query('hello')
event = query.message_event
object.__setattr__(event, 'query', query)
connector_module = get_connector_module()
original_from_event = connector_module.context.EventContext.from_event
original_model_validate = connector_module.context.EventContext.model_validate
response_sources = [
{
'kind': 'reply_message_chain',
'plugin': {'author': 'tester', 'name': 'demo'},
}
]
async def emit_event_response(event_context, include_plugins=None):
return {
'event_context': event_context,
'emitted_plugins': [],
'response_sources': response_sources,
}
connector.handler = AsyncMock()
connector.handler.emit_event = AsyncMock(side_effect=emit_event_response)
fake_event_ctx = Mock()
event_dump = event.model_dump()
event_dump['event_name'] = 'FriendMessage'
fake_event_ctx.model_dump.return_value = {
'query_id': query.query_id,
'eid': 0,
'event_name': 'FriendMessage',
'event': event_dump,
'is_prevent_default': False,
'is_prevent_postorder': False,
}
connector_module.context.EventContext.from_event = Mock(return_value=fake_event_ctx)
parsed_event_ctx = Mock()
connector_module.context.EventContext.model_validate = Mock(return_value=parsed_event_ctx)
try:
event_ctx = await connector.emit_event(event)
finally:
connector_module.context.EventContext.from_event = original_from_event
connector_module.context.EventContext.model_validate = original_model_validate
assert event_ctx is parsed_event_ctx
assert event_ctx._response_sources == response_sources
@pytest.mark.asyncio
async def test_emit_event_leaves_response_sources_absent_for_old_runtime(self):
connector = create_mock_connector()
query = text_query('hello')
event = query.message_event
object.__setattr__(event, 'query', query)
connector_module = get_connector_module()
original_from_event = connector_module.context.EventContext.from_event
original_model_validate = connector_module.context.EventContext.model_validate
async def emit_event_response(event_context, include_plugins=None):
return {
'event_context': event_context,
'emitted_plugins': [
{'manifest': {'metadata': {'author': 'tester', 'name': 'demo'}}},
],
}
connector.handler = AsyncMock()
connector.handler.emit_event = AsyncMock(side_effect=emit_event_response)
fake_event_ctx = Mock()
event_dump = event.model_dump()
event_dump['event_name'] = 'FriendMessage'
fake_event_ctx.model_dump.return_value = {
'query_id': query.query_id,
'eid': 0,
'event_name': 'FriendMessage',
'event': event_dump,
'is_prevent_default': False,
'is_prevent_postorder': False,
}
connector_module.context.EventContext.from_event = Mock(return_value=fake_event_ctx)
parsed_event_ctx = Mock()
connector_module.context.EventContext.model_validate = Mock(return_value=parsed_event_ctx)
try:
event_ctx = await connector.emit_event(event)
finally:
connector_module.context.EventContext.from_event = original_from_event
connector_module.context.EventContext.model_validate = original_model_validate
assert '_response_sources' not in vars(event_ctx)
assert event_ctx._emitted_plugins == [
{'manifest': {'metadata': {'author': 'tester', 'name': 'demo'}}},
]
@pytest.mark.asyncio
async def test_notify_plugin_diagnostic_skips_when_disabled(self):
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
connector.handler = AsyncMock()
await connector.notify_plugin_diagnostic({'code': 'response_delivery_failed'})
connector.handler.notify_plugin_diagnostic.assert_not_called()
@pytest.mark.asyncio
async def test_notify_plugin_diagnostic_is_best_effort(self):
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.notify_plugin_diagnostic = AsyncMock(side_effect=RuntimeError('action not found'))
await connector.notify_plugin_diagnostic({'code': 'response_delivery_failed'})
connector.handler.notify_plugin_diagnostic.assert_awaited_once()
connector.ap.logger.debug.assert_called_once()
class TestListKnowledgeEngines:
"""Tests for list_knowledge_engines method."""
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test returns empty list when plugin system disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.list_knowledge_engines()
assert result == []
@pytest.mark.asyncio
async def test_calls_handler_list_knowledge_engines(self):
"""Test that handler method is called."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.list_knowledge_engines = AsyncMock(
return_value=[{'plugin_id': 'author/engine', 'name': 'Engine'}]
)
result = await connector.list_knowledge_engines()
connector.handler.list_knowledge_engines.assert_called_once()
assert result == [{'plugin_id': 'author/engine', 'name': 'Engine'}]
class TestListParsers:
"""Tests for list_parsers method."""
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test returns empty list when plugin system disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.list_parsers()
assert result == []
@pytest.mark.asyncio
async def test_calls_handler_list_parsers(self):
"""Test that handler method is called."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.list_parsers = AsyncMock(
return_value=[{'plugin_id': 'author/parser', 'supported_mime_types': ['text/plain']}]
)
result = await connector.list_parsers()
connector.handler.list_parsers.assert_called_once()
assert result == [{'plugin_id': 'author/parser', 'supported_mime_types': ['text/plain']}]
class TestCallParser:
"""Tests for call_parser method."""
@pytest.mark.asyncio
async def test_calls_handler_parse_document(self):
"""Test that handler.parse_document is called with correct args."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.parse_document = AsyncMock(return_value={'content': 'parsed'})
result = await connector.call_parser(
'author/parser',
{'mime_type': 'text/plain', 'filename': 'test.txt'},
b'file content',
)
connector.handler.parse_document.assert_called_once_with(
'author',
'parser',
{'mime_type': 'text/plain', 'filename': 'test.txt'},
b'file content',
)
assert result['content'] == 'parsed'
class TestRAGMethods:
"""Tests for RAG-related methods."""
@pytest.mark.asyncio
async def test_call_rag_ingest(self):
"""Test call_rag_ingest calls handler with parsed plugin ID."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.rag_ingest_document = AsyncMock(return_value={'status': 'success'})
result = await connector.call_rag_ingest('author/engine', {'file': 'test.pdf'})
connector.handler.rag_ingest_document.assert_called_once_with('author', 'engine', {'file': 'test.pdf'})
assert result['status'] == 'success'
@pytest.mark.asyncio
async def test_call_rag_retrieve(self):
"""Test call_rag_retrieve calls handler."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.retrieve_knowledge = AsyncMock(
return_value={
'results': [
{'id': 'doc1', 'content': [{'type': 'text', 'text': 'test'}], 'metadata': {}, 'distance': 0.1}
]
}
)
result = await connector.call_rag_retrieve('author/engine', {'query': 'test'})
connector.handler.retrieve_knowledge.assert_called_once_with('author', 'engine', '', {'query': 'test'})
assert result == {
'results': [
{
'id': 'doc1',
'content': [{'type': 'text', 'text': 'test'}],
'metadata': {},
'distance': 0.1,
}
]
}
@pytest.mark.asyncio
async def test_get_rag_creation_schema(self):
"""Test get_rag_creation_schema calls handler."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.get_rag_creation_schema = AsyncMock(return_value={'properties': {'name': {'type': 'string'}}})
result = await connector.get_rag_creation_schema('author/engine')
connector.handler.get_rag_creation_schema.assert_called_once_with('author', 'engine')
assert result == {'properties': {'name': {'type': 'string'}}}
@pytest.mark.asyncio
async def test_get_rag_retrieval_schema(self):
"""Test get_rag_retrieval_schema calls handler."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.get_rag_retrieval_schema = AsyncMock(
return_value={'properties': {'top_k': {'type': 'integer'}}}
)
result = await connector.get_rag_retrieval_schema('author/engine')
connector.handler.get_rag_retrieval_schema.assert_called_once_with('author', 'engine')
assert result == {'properties': {'top_k': {'type': 'integer'}}}
@pytest.mark.asyncio
async def test_rag_on_kb_create(self):
"""Test rag_on_kb_create calls handler."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.rag_on_kb_create = AsyncMock(return_value={'status': 'ok'})
await connector.rag_on_kb_create('author/engine', 'kb-uuid', {'model': 'test'})
connector.handler.rag_on_kb_create.assert_called_once_with('author', 'engine', 'kb-uuid', {'model': 'test'})
@pytest.mark.asyncio
async def test_rag_on_kb_delete(self):
"""Test rag_on_kb_delete calls handler."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.rag_on_kb_delete = AsyncMock(return_value={'status': 'ok'})
await connector.rag_on_kb_delete('author/engine', 'kb-uuid')
connector.handler.rag_on_kb_delete.assert_called_once_with('author', 'engine', 'kb-uuid')
@pytest.mark.asyncio
async def test_call_rag_delete_document(self):
"""Test call_rag_delete_document calls handler."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.rag_delete_document = AsyncMock(return_value=True)
result = await connector.call_rag_delete_document('author/engine', 'doc-uuid', 'kb-uuid')
connector.handler.rag_delete_document.assert_called_once_with('author', 'engine', 'doc-uuid', 'kb-uuid')
assert result is True
class TestRetrieveKnowledge:
"""Tests for retrieve_knowledge method."""
@pytest.mark.asyncio
async def test_returns_empty_results_when_plugin_disabled(self):
"""Test returns empty when plugin disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.retrieve_knowledge('author', 'engine', 'retriever', {})
assert result == {'results': []}
class TestDisabledPluginEarlyReturns:
"""Tests for early returns when plugin system is disabled."""
@pytest.mark.asyncio
async def test_list_tools_returns_empty(self):
"""Test list_tools returns empty when disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.list_tools()
assert result == []
@pytest.mark.asyncio
async def test_list_commands_returns_empty(self):
"""Test list_commands returns empty when disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.list_commands()
assert result == []
@pytest.mark.asyncio
async def test_get_debug_info_returns_empty(self):
"""Test get_debug_info returns empty dict when disabled."""
connector_module = get_connector_module()
async def mock_disconnect(conn):
pass
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
result = await connector.get_debug_info()
assert result == {}
class TestGetPluginInfo:
"""Tests for get_plugin_info method."""
@pytest.mark.asyncio
async def test_calls_handler_get_plugin_info(self):
"""Test that handler.get_plugin_info is called."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.get_plugin_info = AsyncMock(return_value={'manifest': {'metadata': {'name': 'plugin'}}})
result = await connector.get_plugin_info('author', 'plugin')
connector.handler.get_plugin_info.assert_called_once_with('author', 'plugin')
assert result == {'manifest': {'metadata': {'name': 'plugin'}}}
class TestSetPluginConfig:
"""Tests for set_plugin_config method."""
@pytest.mark.asyncio
async def test_calls_handler_set_plugin_config(self):
"""Test that handler.set_plugin_config is called."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.set_plugin_config = AsyncMock(return_value={'status': 'ok'})
await connector.set_plugin_config('author', 'plugin', {'setting': 'value'})
connector.handler.set_plugin_config.assert_called_once_with('author', 'plugin', {'setting': 'value'})
class TestPingPluginRuntime:
"""Tests for ping_plugin_runtime method."""
@pytest.mark.asyncio
async def test_raises_when_handler_not_set(self):
"""Test that exception is raised when handler not initialized."""
get_connector_module()
connector = create_mock_connector()
# handler is not set
with pytest.raises(Exception, match='Plugin runtime is not connected') as exc_info:
await connector.ping_plugin_runtime()
assert 'not connected' in str(exc_info.value)
@pytest.mark.asyncio
async def test_calls_handler_ping(self):
"""Test that handler.ping is called."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.ping = AsyncMock(return_value={'status': 'ok'})
await connector.ping_plugin_runtime()
connector.handler.ping.assert_called_once()
@@ -0,0 +1,32 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError
def make_connector() -> PluginRuntimeConnector:
app = SimpleNamespace(instance_config=SimpleNamespace(data={'plugin': {'enable': True}}))
return PluginRuntimeConnector(app, AsyncMock())
@pytest.mark.asyncio
async def test_ping_plugin_runtime_raises_specific_error_when_not_connected():
connector = make_connector()
with pytest.raises(PluginRuntimeNotConnectedError, match='Plugin runtime is not connected'):
await connector.ping_plugin_runtime()
@pytest.mark.asyncio
async def test_ping_plugin_runtime_delegates_to_connected_handler():
connector = make_connector()
connector.handler = SimpleNamespace(ping=AsyncMock(return_value='pong'))
result = await connector.ping_plugin_runtime()
assert result == 'pong'
connector.handler.ping.assert_awaited_once()
@@ -0,0 +1,143 @@
"""Tests for PluginRuntimeConnector pure logic methods.
Tests methods that don't require real plugin runtime processes:
- _inspect_plugin_package: identity and deps extraction from zip files
- _parse_plugin_id: plugin ID string parsing
"""
from __future__ import annotations
import io
import zipfile
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
class TestExtractDepsMetadata:
"""Tests for dependency metadata extraction from plugin packages."""
def _create_connector(self):
"""Create a connector instance for testing."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
connector = PluginRuntimeConnector(mock_app, MagicMock())
return connector
def test_extract_deps_with_requirements_txt(self):
"""Extract dependency count from requirements.txt in plugin zip."""
connector = self._create_connector()
# Create a mock zip file with requirements.txt
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w') as zf:
zf.writestr('requirements.txt', 'requests>=2.0\nflask\n# comment\n\nnumpy')
zip_bytes = zip_buffer.getvalue()
task_context = SimpleNamespace(metadata={})
connector._inspect_plugin_package(zip_bytes, task_context)
assert task_context.metadata['deps_total'] == 3 # requests>=2.0, flask, numpy
# deps_list contains full requirement lines including version specifiers
assert 'requests>=2.0' in task_context.metadata['deps_list']
assert 'flask' in task_context.metadata['deps_list']
assert 'numpy' in task_context.metadata['deps_list']
def test_extract_deps_empty_requirements(self):
"""Handle empty requirements.txt."""
connector = self._create_connector()
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w') as zf:
zf.writestr('requirements.txt', '# only comments\n\n')
zip_bytes = zip_buffer.getvalue()
task_context = SimpleNamespace(metadata={})
connector._inspect_plugin_package(zip_bytes, task_context)
assert task_context.metadata['deps_total'] == 0
assert task_context.metadata['deps_list'] == []
def test_extract_deps_no_requirements_txt(self):
"""Handle zip without requirements.txt."""
connector = self._create_connector()
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w') as zf:
zf.writestr('plugin.py', 'print("hello")')
zip_bytes = zip_buffer.getvalue()
task_context = SimpleNamespace(metadata={})
connector._inspect_plugin_package(zip_bytes, task_context)
# No requirements.txt found, metadata unchanged
assert 'deps_total' not in task_context.metadata
def test_extract_deps_none_task_context(self):
"""Handle None task_context gracefully."""
connector = self._create_connector()
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w') as zf:
zf.writestr('requirements.txt', 'requests')
zip_bytes = zip_buffer.getvalue()
# Should return early without error
connector._inspect_plugin_package(zip_bytes, None)
def test_extract_deps_invalid_zip(self):
"""Handle invalid zip file gracefully."""
connector = self._create_connector()
# Not a valid zip
invalid_bytes = b'not a zip file'
task_context = SimpleNamespace(metadata={})
connector._inspect_plugin_package(invalid_bytes, task_context)
# Should catch exception and pass silently
assert 'deps_total' not in task_context.metadata
def test_extract_deps_nested_requirements(self):
"""Handle requirements.txt in nested directory."""
connector = self._create_connector()
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w') as zf:
zf.writestr('subdir/requirements.txt', 'pytest\nblack')
zip_bytes = zip_buffer.getvalue()
task_context = SimpleNamespace(metadata={})
connector._inspect_plugin_package(zip_bytes, task_context)
# Should find requirements.txt in subdirectory
assert task_context.metadata['deps_total'] == 2
class TestParsePluginId:
"""Tests for _parse_plugin_id static method."""
def test_parse_valid_plugin_id(self):
"""Parse valid plugin ID format 'author/name'."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
author, name = PluginRuntimeConnector._parse_plugin_id('myauthor/myplugin')
assert author == 'myauthor'
assert name == 'myplugin'
def test_parse_plugin_id_empty(self):
"""Empty plugin ID is invalid."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
with pytest.raises(ValueError):
PluginRuntimeConnector._parse_plugin_id('')
@@ -0,0 +1,55 @@
"""Unit tests for plugin connector static methods.
Tests cover:
- _parse_plugin_id() parsing and validation
"""
from __future__ import annotations
import pytest
from importlib import import_module
def get_connector_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.plugin.connector')
class TestParsePluginId:
"""Tests for _parse_plugin_id static method."""
def test_valid_plugin_id_simple(self):
"""Test parsing valid plugin ID with simple format."""
connector = get_connector_module()
author, name = connector.PluginRuntimeConnector._parse_plugin_id('langbot/rag-engine')
assert author == 'langbot'
assert name == 'rag-engine'
def test_invalid_plugin_id_no_slash(self):
"""Test that ValueError is raised when no slash present."""
connector = get_connector_module()
with pytest.raises(ValueError) as exc_info:
connector.PluginRuntimeConnector._parse_plugin_id('invalid-plugin-id')
assert 'Invalid plugin_id format' in str(exc_info.value)
assert 'invalid-plugin-id' in str(exc_info.value)
def test_invalid_plugin_id_empty_string(self):
"""Test that ValueError is raised for empty string."""
connector = get_connector_module()
with pytest.raises(ValueError) as exc_info:
connector.PluginRuntimeConnector._parse_plugin_id('')
assert 'Invalid plugin_id format' in str(exc_info.value)
def test_valid_plugin_id_single_character_parts(self):
"""Test parsing plugin ID with single character author and name."""
connector = get_connector_module()
author, name = connector.PluginRuntimeConnector._parse_plugin_id('a/b')
assert author == 'a'
assert name == 'b'
def test_valid_plugin_id_with_hyphens_and_underscores(self):
"""Test parsing plugin ID with hyphens and underscores."""
connector = get_connector_module()
author, name = connector.PluginRuntimeConnector._parse_plugin_id('lang-bot/my_rag_engine')
assert author == 'lang-bot'
assert name == 'my_rag_engine'
@@ -0,0 +1,211 @@
"""Unit tests for plugin connector _inspect_plugin_package method.
Tests cover:
- Extracting requirements.txt from ZIP
- Parsing dependency lines
- Handling missing requirements.txt
- Handling empty/malformed requirements.txt
"""
from __future__ import annotations
import zipfile
import io
from unittest.mock import Mock
from importlib import import_module
def get_connector_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.plugin.connector')
def create_mock_connector():
"""Create a mock PluginRuntimeConnector instance for testing."""
connector = get_connector_module()
mock_app = Mock()
mock_app.logger = Mock()
mock_app.instance_config = Mock()
mock_app.instance_config.data = {'plugin': {'enable': True}}
# Mock disconnect callback
async def mock_disconnect_callback(connector):
pass
return connector.PluginRuntimeConnector(mock_app, mock_disconnect_callback)
def create_zip_with_requirements(requirements_content: str) -> bytes:
"""Create a ZIP file containing requirements.txt with given content."""
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr('requirements.txt', requirements_content)
return buf.getvalue()
def create_zip_with_nested_requirements(requirements_content: str) -> bytes:
"""Create a ZIP file with requirements.txt in nested directory."""
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr('plugin/requirements.txt', requirements_content)
return buf.getvalue()
def create_zip_without_requirements() -> bytes:
"""Create a ZIP file without requirements.txt."""
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr('main.py', 'print("hello")')
zf.writestr('manifest.yaml', 'name: test')
return buf.getvalue()
class TestExtractDepsMetadata:
"""Tests for dependency metadata extraction from plugin packages."""
def test_extract_simple_requirements(self):
"""Test extracting simple requirements.txt."""
connector_instance = create_mock_connector()
# Create test ZIP
zip_bytes = create_zip_with_requirements('requests>=2.0\nflask==1.0\nnumpy')
# Create task context
task_context = Mock()
task_context.metadata = {}
connector_instance._inspect_plugin_package(zip_bytes, task_context)
assert task_context.metadata.get('deps_total') == 3
assert task_context.metadata.get('deps_list') == ['requests>=2.0', 'flask==1.0', 'numpy']
def test_extract_requirements_with_comments_and_empty_lines(self):
"""Test that comments and empty lines are filtered."""
connector_instance = create_mock_connector()
requirements = """# This is a comment
requests>=2.0
# Another comment
flask==1.0
numpy"""
zip_bytes = create_zip_with_requirements(requirements)
task_context = Mock()
task_context.metadata = {}
connector_instance._inspect_plugin_package(zip_bytes, task_context)
assert task_context.metadata.get('deps_total') == 3
assert '# This is a comment' not in task_context.metadata.get('deps_list', [])
def test_extract_nested_requirements(self):
"""Test extracting requirements.txt from nested directory."""
connector_instance = create_mock_connector()
zip_bytes = create_zip_with_nested_requirements('requests\nflask')
task_context = Mock()
task_context.metadata = {}
connector_instance._inspect_plugin_package(zip_bytes, task_context)
# Should find nested requirements.txt (ends with 'requirements.txt')
assert task_context.metadata.get('deps_total') == 2
def test_no_requirements_in_zip(self):
"""Test handling ZIP without requirements.txt."""
connector_instance = create_mock_connector()
zip_bytes = create_zip_without_requirements()
task_context = Mock()
task_context.metadata = {}
connector_instance._inspect_plugin_package(zip_bytes, task_context)
# metadata should remain empty (no deps found)
assert task_context.metadata.get('deps_total') is None
assert task_context.metadata.get('deps_list') is None
def test_empty_requirements_file(self):
"""Test handling empty requirements.txt."""
connector_instance = create_mock_connector()
zip_bytes = create_zip_with_requirements('')
task_context = Mock()
task_context.metadata = {}
connector_instance._inspect_plugin_package(zip_bytes, task_context)
# deps_total should be 0 (empty list after filtering)
assert task_context.metadata.get('deps_total') == 0
assert task_context.metadata.get('deps_list') == []
def test_requirements_only_comments(self):
"""Test handling requirements.txt with only comments."""
connector_instance = create_mock_connector()
requirements = """# Comment 1
# Comment 2
# Comment 3"""
zip_bytes = create_zip_with_requirements(requirements)
task_context = Mock()
task_context.metadata = {}
connector_instance._inspect_plugin_package(zip_bytes, task_context)
assert task_context.metadata.get('deps_total') == 0
assert task_context.metadata.get('deps_list') == []
def test_task_context_none_returns_early(self):
"""Test that method returns early when task_context is None."""
connector_instance = create_mock_connector()
zip_bytes = create_zip_with_requirements('requests')
# Should return without error when task_context is None
connector_instance._inspect_plugin_package(zip_bytes, None)
# No exception should be raised
def test_malformed_zip_handling(self):
"""Test handling malformed ZIP bytes."""
connector_instance = create_mock_connector()
# Invalid ZIP bytes
invalid_bytes = b'not a valid zip file'
task_context = Mock()
task_context.metadata = {}
# Should silently handle exception (pass in try/except)
connector_instance._inspect_plugin_package(invalid_bytes, task_context)
# metadata should remain unchanged
assert task_context.metadata == {}
def test_requirements_with_unicode_decode_error(self):
"""Test handling requirements.txt with non-UTF8 content."""
connector_instance = create_mock_connector()
# Create ZIP with non-UTF8 content in requirements.txt
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
# Write bytes that will cause decode issues
# \x80 is invalid UTF-8, but errors='ignore' will skip it
zf.writestr('requirements.txt', b'requests\nflask\n\x80invalid')
zip_bytes = buf.getvalue()
task_context = Mock()
task_context.metadata = {}
# errors='ignore' will decode \x80invalid as 'invalid' (skipping \x80)
connector_instance._inspect_plugin_package(zip_bytes, task_context)
# All 3 lines will be parsed (requests, flask, invalid)
assert task_context.metadata.get('deps_total') == 3
assert 'invalid' in task_context.metadata.get('deps_list', [])
+214
View File
@@ -0,0 +1,214 @@
"""Tests for RuntimeConnectionHandler helper functions.
Tests handler helper methods that don't require full handler setup.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, Mock
import pytest
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
def make_handler(app):
"""Create a RuntimeConnectionHandler with mocked external connection."""
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
return RuntimeConnectionHandler(Mock(), AsyncMock(return_value=True), app)
class TestHandlerQueryVariables:
"""Tests for handler query variable logic."""
@pytest.fixture
def mock_app(self):
"""Create mock app with query pool."""
app = SimpleNamespace()
app.query_pool = SimpleNamespace()
app.query_pool.cached_queries = {}
app.logger = SimpleNamespace()
app.logger.debug = MagicMock()
return app
@pytest.mark.asyncio
async def test_set_query_var_query_not_found(self, mock_app):
"""Test set_query_var returns error when query not found."""
runtime_handler = make_handler(mock_app)
response = await runtime_handler.actions[PluginToRuntimeAction.SET_QUERY_VAR.value](
{
'query_id': 'nonexistent-query',
'key': 'test_var',
'value': 'test_value',
}
)
assert response.code != 0
assert 'nonexistent-query' in response.message
@pytest.mark.asyncio
async def test_set_query_var_success(self, mock_app):
"""Test set_query_var sets variable on existing query."""
runtime_handler = make_handler(mock_app)
mock_query = SimpleNamespace()
mock_query.variables = {}
mock_app.query_pool.cached_queries['test-query'] = mock_query
response = await runtime_handler.actions[PluginToRuntimeAction.SET_QUERY_VAR.value](
{
'query_id': 'test-query',
'key': 'test_var',
'value': 'test_value',
}
)
assert response.code == 0
assert mock_query.variables['test_var'] == 'test_value'
@pytest.mark.asyncio
async def test_get_query_var_success(self, mock_app):
"""Test get_query_var retrieves variable from query."""
runtime_handler = make_handler(mock_app)
mock_query = SimpleNamespace()
mock_query.variables = {'existing_var': 'existing_value'}
mock_app.query_pool.cached_queries['test-query'] = mock_query
response = await runtime_handler.actions[PluginToRuntimeAction.GET_QUERY_VAR.value](
{
'query_id': 'test-query',
'key': 'existing_var',
}
)
assert response.code == 0
assert response.data == {'value': 'existing_value'}
@pytest.mark.asyncio
async def test_get_query_vars_multiple(self, mock_app):
"""Test get_query_vars returns the query's variable mapping."""
runtime_handler = make_handler(mock_app)
mock_query = SimpleNamespace()
mock_query.variables = {'var1': 'val1', 'var2': 'val2', 'var3': 'val3'}
mock_app.query_pool.cached_queries['test-query'] = mock_query
response = await runtime_handler.actions[PluginToRuntimeAction.GET_QUERY_VARS.value](
{
'query_id': 'test-query',
}
)
assert response.code == 0
assert response.data == {'vars': mock_query.variables}
class TestHandlerRagErrorResponse:
"""Tests for _make_rag_error_response helper."""
def test_make_rag_error_response_basic(self):
"""Test basic error response creation."""
from langbot.pkg.plugin.handler import _make_rag_error_response
error = Exception('test error')
response = _make_rag_error_response(error, 'TestError')
# ActionResponse is a pydantic model, check message field
assert 'TestError' in response.message
assert 'test error' in response.message
assert 'Exception' in response.message
def test_make_rag_error_response_with_context(self):
"""Test error response with extra context."""
from langbot.pkg.plugin.handler import _make_rag_error_response
error = ValueError('invalid input')
response = _make_rag_error_response(error, 'ValidationError', field='name', value='test')
assert 'ValidationError' in response.message
assert 'field=name' in response.message
assert 'value=test' in response.message
assert 'ValueError' in response.message
def test_make_rag_error_response_exception_type(self):
"""Test error response includes exception type."""
from langbot.pkg.plugin.handler import _make_rag_error_response
error = RuntimeError('connection failed')
response = _make_rag_error_response(error, 'ConnectionError')
assert 'RuntimeError' in response.message
assert 'ConnectionError' in response.message
assert 'connection failed' in response.message
def test_make_rag_error_response_empty_context(self):
"""Test error response with no extra context."""
from langbot.pkg.plugin.handler import _make_rag_error_response
error = KeyError('missing_key')
response = _make_rag_error_response(error, 'LookupError')
# No context parts means no brackets
assert '[' in response.message # Still has error type bracket
assert 'KeyError' in response.message
class TestHandlerPluginDiagnostic:
@pytest.mark.asyncio
async def test_notify_plugin_diagnostic_falls_back_to_raw_protocol_action(self):
"""Diagnostic forwarding works before the SDK enum exists."""
app = SimpleNamespace()
app.logger = SimpleNamespace(debug=MagicMock())
runtime_handler = make_handler(app)
runtime_handler.call_action = AsyncMock(return_value={})
payload = {'code': 'response_delivery_failed'}
await runtime_handler.notify_plugin_diagnostic(payload)
action = runtime_handler.call_action.await_args.args[0]
assert action.value == 'plugin_diagnostic'
assert runtime_handler.call_action.await_args.args[1] is payload
assert runtime_handler.call_action.await_args.kwargs['timeout'] == 5
def test_langbot_to_runtime_action_uses_enum_when_available(self):
"""The compatibility helper should prefer SDK enums once available."""
from langbot.pkg.plugin import handler as plugin_handler
sentinel = object()
original = plugin_handler.LangBotToRuntimeAction
plugin_handler.LangBotToRuntimeAction = SimpleNamespace(PLUGIN_DIAGNOSTIC=sentinel)
try:
assert plugin_handler._langbot_to_runtime_action('PLUGIN_DIAGNOSTIC', 'plugin_diagnostic') is sentinel
finally:
plugin_handler.LangBotToRuntimeAction = original
class TestConstantsSemanticVersion:
"""Tests for version constant access."""
def test_semantic_version_exists(self):
"""Test semantic_version is defined."""
from langbot.pkg.utils import constants
assert hasattr(constants, 'semantic_version')
assert constants.semantic_version.startswith('v')
def test_edition_exists(self):
"""Test edition constant is defined."""
from langbot.pkg.utils import constants
assert hasattr(constants, 'edition')
assert constants.edition == 'community'
def test_required_database_version_exists(self):
"""Test database version constant."""
from langbot.pkg.utils import constants
assert hasattr(constants, 'required_database_version')
assert isinstance(constants.required_database_version, int)
@@ -0,0 +1,425 @@
"""Unit tests for RuntimeConnectionHandler action handlers."""
from __future__ import annotations
import base64
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction, RuntimeToLangBotAction
def make_handler(app):
"""Create a RuntimeConnectionHandler with mocked external connection."""
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
return RuntimeConnectionHandler(Mock(), AsyncMock(return_value=True), app)
def make_result(first_item=None):
result = Mock()
result.first = Mock(return_value=first_item)
return result
def compiled_params(statement):
return statement.compile().params
class TestRagRerankAction:
"""Tests for RAG rerank action handler."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.model_mgr = Mock()
mock_app.logger = Mock()
return mock_app
@pytest.mark.asyncio
async def test_invokes_rerank_model_and_sorts_scores(self, app):
"""Rerank action uses the selected model and returns top scores."""
provider = Mock()
provider.invoke_rerank = AsyncMock(
return_value=[
{'index': 0, 'relevance_score': 0.2},
{'index': 1, 'relevance_score': 0.9},
]
)
rerank_model = SimpleNamespace(provider=provider)
app.model_mgr.get_rerank_model_by_uuid = AsyncMock(return_value=rerank_model)
runtime_handler = make_handler(app)
response = await runtime_handler.actions[PluginToRuntimeAction.INVOKE_RERANK.value](
{
'rerank_model_uuid': 'rerank-1',
'query': 'hello',
'documents': ['a', 'b'],
'top_k': 1,
'extra_args': {'return_documents': False},
}
)
assert response.code == 0
assert response.data['results'] == [{'index': 1, 'relevance_score': 0.9}]
app.model_mgr.get_rerank_model_by_uuid.assert_awaited_once_with('rerank-1')
provider.invoke_rerank.assert_awaited_once_with(
model=rerank_model,
query='hello',
documents=['a', 'b'],
extra_args={'return_documents': False},
)
@pytest.mark.asyncio
async def test_returns_error_when_rerank_model_missing(self, app):
"""Missing rerank model returns an action error."""
app.model_mgr.get_rerank_model_by_uuid = AsyncMock(side_effect=ValueError('not found'))
runtime_handler = make_handler(app)
response = await runtime_handler.actions[PluginToRuntimeAction.INVOKE_RERANK.value](
{
'rerank_model_uuid': 'missing',
'query': 'hello',
'documents': ['a'],
}
)
assert response.code != 0
assert 'Rerank model with rerank_model_uuid missing not found' in response.message
class TestInitializePluginSettings:
"""Tests for initialize_plugin_settings action handler."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock()
mock_app.logger = Mock()
return mock_app
@pytest.mark.asyncio
async def test_creates_new_setting_when_not_exists(self, app):
"""New plugin settings use default enabled, priority and config values."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.side_effect = [
make_result(),
Mock(),
]
response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value](
{
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
'install_source': 'local',
'install_info': {'path': '/test'},
}
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
assert insert_params == {
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
'install_source': 'local',
'install_info': {'path': '/test'},
'enabled': True,
'priority': 0,
'config': {},
}
@pytest.mark.asyncio
async def test_inherits_values_from_existing_setting(self, app):
"""Existing settings are replaced while preserving user-controlled values."""
runtime_handler = make_handler(app)
existing_setting = SimpleNamespace(
enabled=False,
priority=5,
config={'key': 'value'},
)
app.persistence_mgr.execute_async.side_effect = [
make_result(existing_setting),
Mock(),
Mock(),
]
response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value](
{
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
'install_source': 'github',
'install_info': {'repo': 'author/name'},
}
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 3
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0])
assert insert_params['enabled'] is False
assert insert_params['priority'] == 5
assert insert_params['config'] == {'key': 'value'}
assert insert_params['install_source'] == 'github'
assert insert_params['install_info'] == {'repo': 'author/name'}
class TestSetBinaryStorage:
"""Tests for set_binary_storage action handler with size limit validation."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.instance_config = Mock()
mock_app.instance_config.data = {
'plugin': {
'binary_storage': {
'max_value_bytes': 1024,
},
},
}
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=make_result())
mock_app.logger = Mock()
return mock_app
@staticmethod
def payload(value: bytes):
return {
'key': 'test-key',
'owner_type': 'plugin',
'owner': 'test-owner',
'value_base64': base64.b64encode(value).decode('utf-8'),
}
@pytest.mark.asyncio
async def test_rejects_value_exceeding_limit(self, app):
"""Values larger than max_value_bytes are rejected before persistence writes."""
runtime_handler = make_handler(app)
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](
self.payload(b'x' * 2048)
)
assert response.code != 0
assert '2048 > 1024 bytes' in response.message
app.persistence_mgr.execute_async.assert_not_awaited()
@pytest.mark.asyncio
async def test_accepts_value_within_limit_and_inserts_storage(self, app):
"""A new small value is inserted into binary storage."""
runtime_handler = make_handler(app)
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](
self.payload(b'x' * 512)
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
assert insert_params['unique_key'] == 'plugin:test-owner:test-key'
assert insert_params['value'] == b'x' * 512
@pytest.mark.asyncio
async def test_updates_existing_storage(self, app):
"""An existing binary storage row is updated instead of inserted."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.return_value = make_result(SimpleNamespace(value=b'old'))
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new'))
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
update_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
assert update_params['value'] == b'new'
@pytest.mark.asyncio
async def test_invalid_max_value_bytes_falls_back_to_default_limit(self, app):
"""Invalid max_value_bytes uses the 10MB default limit."""
runtime_handler = make_handler(app)
app.instance_config.data['plugin']['binary_storage']['max_value_bytes'] = 'invalid'
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](
self.payload(b'x' * (10 * 1024 * 1024 + 1))
)
assert response.code != 0
assert '10485761 > 10485760 bytes' in response.message
app.persistence_mgr.execute_async.assert_not_awaited()
@pytest.mark.asyncio
async def test_negative_limit_disables_size_check(self, app):
"""Negative max_value_bytes allows values larger than the normal default."""
runtime_handler = make_handler(app)
app.instance_config.data['plugin']['binary_storage']['max_value_bytes'] = -1
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](
self.payload(b'x' * 2048)
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
@pytest.mark.asyncio
async def test_zero_limit_rejects_non_empty_values(self, app):
"""A zero byte limit rejects non-empty values."""
runtime_handler = make_handler(app)
app.instance_config.data['plugin']['binary_storage']['max_value_bytes'] = 0
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'x'))
assert response.code != 0
assert '1 > 0 bytes' in response.message
app.persistence_mgr.execute_async.assert_not_awaited()
class TestGetPluginSettings:
"""Tests for get_plugin_settings action handler with defaults."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock()
return mock_app
@pytest.mark.asyncio
async def test_returns_defaults_when_setting_not_found(self, app):
"""Default plugin settings are returned when no persisted row exists."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.return_value = make_result()
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value](
{
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
}
)
assert response.code == 0
assert response.data == {
'enabled': True,
'priority': 0,
'plugin_config': {},
'install_source': 'local',
'install_info': {},
}
@pytest.mark.asyncio
async def test_returns_actual_values_when_setting_exists(self, app):
"""Persisted plugin setting values override defaults."""
runtime_handler = make_handler(app)
setting = SimpleNamespace(
enabled=False,
priority=10,
config={'custom': 'config'},
install_source='github',
install_info={'repo': 'test/repo'},
)
app.persistence_mgr.execute_async.return_value = make_result(setting)
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value](
{
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
}
)
assert response.code == 0
assert response.data == {
'enabled': False,
'priority': 10,
'plugin_config': {'custom': 'config'},
'install_source': 'github',
'install_info': {'repo': 'test/repo'},
}
class TestGetBinaryStorage:
"""Tests for get_binary_storage action handler."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock()
return mock_app
@pytest.mark.asyncio
async def test_returns_base64_encoded_value(self, app):
"""Stored bytes are returned as base64."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.return_value = make_result(SimpleNamespace(value=b'test binary content'))
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE.value](
{
'key': 'test-key',
'owner_type': 'plugin',
'owner': 'test-owner',
}
)
assert response.code == 0
assert response.data == {
'value_base64': base64.b64encode(b'test binary content').decode('utf-8'),
}
@pytest.mark.asyncio
async def test_returns_error_when_not_found(self, app):
"""Missing binary storage rows return an error response."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.return_value = make_result()
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE.value](
{
'key': 'test-key',
'owner_type': 'plugin',
'owner': 'test-owner',
}
)
assert response.code != 0
assert 'Storage with key test-key not found' in response.message
class TestHandlerQueryLookup:
"""Tests for query lookup in cached_queries."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.query_pool = Mock()
mock_app.query_pool.cached_queries = {}
mock_app.logger = Mock()
return mock_app
@pytest.mark.asyncio
async def test_query_not_found_returns_error(self, app):
"""Query-bound actions return error when query_id is not cached."""
runtime_handler = make_handler(app)
response = await runtime_handler.actions[PluginToRuntimeAction.GET_BOT_UUID.value](
{
'query_id': 'nonexistent-query',
}
)
assert response.code != 0
assert 'nonexistent-query' in response.message
@pytest.mark.asyncio
async def test_query_found_returns_success(self, app):
"""Query-bound actions read data from the cached query object."""
runtime_handler = make_handler(app)
query = SimpleNamespace(variables={}, bot_uuid='test-bot-uuid')
app.query_pool.cached_queries['existing-query'] = query
response = await runtime_handler.actions[PluginToRuntimeAction.GET_BOT_UUID.value](
{
'query_id': 'existing-query',
}
)
assert response.code == 0
assert response.data == {'bot_uuid': 'test-bot-uuid'}
@@ -0,0 +1,126 @@
"""Unit tests for plugin handler helper functions and methods.
Tests cover:
- _make_rag_error_response() helper function
- RuntimeConnectionHandler cleanup_plugin_data method
"""
from __future__ import annotations
import pytest
from unittest.mock import Mock, AsyncMock
from importlib import import_module
def get_handler_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.plugin.handler')
class TestMakeRagErrorResponse:
"""Tests for _make_rag_error_response helper function."""
def test_creates_error_response_with_exception(self):
"""Test basic error response creation."""
handler = get_handler_module()
error = ValueError('test error message')
result = handler._make_rag_error_response(error, 'TestError')
# ActionResponse.error() returns code=1 (error status)
assert result.code == 1
assert 'TestError' in result.message
assert 'ValueError' in result.message
assert 'test error message' in result.message
def test_includes_error_type_in_message(self):
"""Test that error type is included in message."""
handler = get_handler_module()
error = RuntimeError('something went wrong')
result = handler._make_rag_error_response(error, 'VectorStoreError')
assert '[VectorStoreError/RuntimeError]' in result.message
def test_includes_extra_context_in_message(self):
"""Test that extra context fields are included."""
handler = get_handler_module()
error = Exception('embedding failed')
result = handler._make_rag_error_response(
error,
'EmbeddingError',
embedding_model_uuid='test-uuid-123',
collection_id='collection-456',
)
assert 'embedding_model_uuid=test-uuid-123' in result.message
assert 'collection_id=collection-456' in result.message
def test_handles_exception_with_no_message(self):
"""Test handling exception with empty message."""
handler = get_handler_module()
error = Exception()
result = handler._make_rag_error_response(error, 'GenericError')
# ActionResponse.error() returns code=1 (error status)
assert result.code == 1
assert '[GenericError/Exception]' in result.message
def test_formats_context_with_multiple_fields(self):
"""Test multiple context fields are comma separated."""
handler = get_handler_module()
error = IOError('file not found')
result = handler._make_rag_error_response(
error,
'FileServiceError',
storage_path='/data/file.pdf',
kb_id='kb-001',
)
assert '[storage_path=/data/file.pdf, kb_id=kb-001]' in result.message
class TestCleanupPluginData:
"""Tests for cleanup_plugin_data method."""
@pytest.mark.asyncio
async def test_deletes_plugin_settings(self):
"""Test that plugin settings are deleted."""
handler_module = get_handler_module()
mock_app = Mock()
mock_app.persistence_mgr = AsyncMock()
mock_app.persistence_mgr.execute_async = AsyncMock()
# Mock the handler without connection (we only need ap)
handler_instance = Mock(spec=handler_module.RuntimeConnectionHandler)
handler_instance.ap = mock_app
# Call cleanup_plugin_data
await handler_module.RuntimeConnectionHandler.cleanup_plugin_data(
handler_instance, 'test-author', 'test-plugin'
)
# Verify plugin settings delete was called
calls = mock_app.persistence_mgr.execute_async.call_args_list
assert len(calls) >= 1
@pytest.mark.asyncio
async def test_deletes_binary_storage(self):
"""Test that binary storage is deleted."""
handler_module = get_handler_module()
mock_app = Mock()
mock_app.persistence_mgr = AsyncMock()
mock_app.persistence_mgr.execute_async = AsyncMock()
handler_instance = Mock(spec=handler_module.RuntimeConnectionHandler)
handler_instance.ap = mock_app
await handler_module.RuntimeConnectionHandler.cleanup_plugin_data(handler_instance, 'author', 'plugin-name')
# Should have at least 2 calls: one for settings, one for binary storage
assert mock_app.persistence_mgr.execute_async.call_count >= 2
@@ -0,0 +1,276 @@
"""Test plugin list filtering by component kinds."""
from unittest.mock import AsyncMock, MagicMock
import pytest
@pytest.mark.asyncio
async def test_plugin_list_filter_by_component_kinds():
"""Test that plugins can be filtered by component kinds."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock plugin data with different component kinds
mock_plugins = [
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author1',
'name': 'plugin_with_tool',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'Tool', 'metadata': {'name': 'tool1'}}}}],
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author2',
'name': 'plugin_with_knowledge_engine_only',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'KnowledgeEngine', 'metadata': {'name': 'retriever1'}}}}],
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author3',
'name': 'plugin_with_command',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'Command', 'metadata': {'name': 'cmd1'}}}}],
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author4',
'name': 'plugin_with_event_listener',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'EventListener', 'metadata': {'name': 'listener1'}}}}],
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author5',
'name': 'plugin_with_mixed_components',
}
}
},
'components': [
{'manifest': {'manifest': {'kind': 'KnowledgeEngine', 'metadata': {'name': 'retriever2'}}}},
{'manifest': {'manifest': {'kind': 'Tool', 'metadata': {'name': 'tool2'}}}},
],
},
]
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
# Test filtering by pipeline component kinds (Command, EventListener, Tool)
pipeline_component_kinds = ['Command', 'EventListener', 'Tool']
result = await connector.list_plugins(component_kinds=pipeline_component_kinds)
# Verify that only plugins with pipeline-related components are returned
assert len(result) == 4
plugin_names = [p['manifest']['manifest']['metadata']['name'] for p in result]
assert 'plugin_with_tool' in plugin_names
assert 'plugin_with_command' in plugin_names
assert 'plugin_with_event_listener' in plugin_names
assert 'plugin_with_mixed_components' in plugin_names
# Plugin with only KnowledgeEngine should NOT be included
assert 'plugin_with_knowledge_engine_only' not in plugin_names
@pytest.mark.asyncio
async def test_plugin_list_filter_no_filter():
"""Test that all plugins are returned when no filter is specified."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock plugin data with different component kinds
mock_plugins = [
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author1',
'name': 'plugin1',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'Tool', 'metadata': {'name': 'tool1'}}}}],
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author2',
'name': 'plugin2',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'KnowledgeEngine', 'metadata': {'name': 'retriever1'}}}}],
},
]
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
# Test without filter - should return all plugins
result = await connector.list_plugins()
assert len(result) == 2
plugin_names = [p['manifest']['manifest']['metadata']['name'] for p in result]
assert 'plugin1' in plugin_names
assert 'plugin2' in plugin_names
@pytest.mark.asyncio
async def test_plugin_list_filter_empty_result():
"""Test that empty list is returned when no plugins match the filter."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock plugin data - only KnowledgeEngine plugins
mock_plugins = [
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author1',
'name': 'plugin1',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'KnowledgeEngine', 'metadata': {'name': 'retriever1'}}}}],
},
]
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
# Filter by Tool kind - should return empty list
result = await connector.list_plugins(component_kinds=['Tool'])
assert len(result) == 0
@pytest.mark.asyncio
async def test_plugin_list_filter_plugin_without_components():
"""Test that plugins without components are excluded when filtering."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock plugin data - one with components, one without
mock_plugins = [
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author1',
'name': 'plugin_with_tool',
}
}
},
'components': [{'manifest': {'manifest': {'kind': 'Tool', 'metadata': {'name': 'tool1'}}}}],
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author2',
'name': 'plugin_without_components',
}
}
},
'components': [],
},
]
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
# Filter by Tool kind - should return only plugin with Tool
result = await connector.list_plugins(component_kinds=['Tool'])
assert len(result) == 1
assert result[0]['manifest']['manifest']['metadata']['name'] == 'plugin_with_tool'
@@ -0,0 +1,25 @@
"""Test plugin ID parsing validation."""
import pytest
from langbot.pkg.plugin.connector import PluginRuntimeConnector
def test_parse_plugin_id_accepts_author_name():
assert PluginRuntimeConnector._parse_plugin_id('langbot/rag-engine') == ('langbot', 'rag-engine')
@pytest.mark.parametrize(
'plugin_id',
[
'',
'author',
'author/',
'/name',
'author/name/extra',
'/',
],
)
def test_parse_plugin_id_rejects_malformed_ids(plugin_id):
with pytest.raises(ValueError, match='Expected'):
PluginRuntimeConnector._parse_plugin_id(plugin_id)
@@ -0,0 +1,228 @@
"""Test plugin list sorting functionality."""
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock
import pytest
@pytest.mark.asyncio
async def test_plugin_list_sorting_debug_first():
"""Test that debug plugins appear before non-debug plugins."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock plugin data with different debug states and timestamps
now = datetime.now()
mock_plugins = [
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author1',
'name': 'plugin1',
}
}
},
},
{
'debug': True,
'manifest': {
'manifest': {
'metadata': {
'author': 'author2',
'name': 'plugin2',
}
}
},
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author3',
'name': 'plugin3',
}
}
},
},
]
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query to return all timestamps in a single batch
async def mock_execute_async(query):
mock_result = MagicMock()
# Create mock rows for all plugins with timestamps
mock_rows = []
# plugin1: oldest, plugin2: middle, plugin3: newest
mock_row1 = MagicMock()
mock_row1.plugin_author = 'author1'
mock_row1.plugin_name = 'plugin1'
mock_row1.created_at = now - timedelta(days=2)
mock_rows.append(mock_row1)
mock_row2 = MagicMock()
mock_row2.plugin_author = 'author2'
mock_row2.plugin_name = 'plugin2'
mock_row2.created_at = now - timedelta(days=1)
mock_rows.append(mock_row2)
mock_row3 = MagicMock()
mock_row3.plugin_author = 'author3'
mock_row3.plugin_name = 'plugin3'
mock_row3.created_at = now
mock_rows.append(mock_row3)
# Make the result iterable
mock_result.__iter__ = lambda self: iter(mock_rows)
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
# Call list_plugins
result = await connector.list_plugins()
# Verify sorting: debug plugin should be first
assert len(result) == 3
assert result[0]['debug'] is True # plugin2 (debug)
assert result[0]['manifest']['manifest']['metadata']['name'] == 'plugin2'
# Remaining should be sorted by created_at (newest first)
assert result[1]['debug'] is False
assert result[1]['manifest']['manifest']['metadata']['name'] == 'plugin3' # newest non-debug
assert result[2]['debug'] is False
assert result[2]['manifest']['manifest']['metadata']['name'] == 'plugin1' # oldest non-debug
@pytest.mark.asyncio
async def test_plugin_list_sorting_by_installation_time():
"""Test that non-debug plugins are sorted by installation time (newest first)."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock plugin data - all non-debug with different installation times
now = datetime.now()
mock_plugins = [
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author1',
'name': 'oldest_plugin',
}
}
},
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author2',
'name': 'middle_plugin',
}
}
},
},
{
'debug': False,
'manifest': {
'manifest': {
'metadata': {
'author': 'author3',
'name': 'newest_plugin',
}
}
},
},
]
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query to return all timestamps in a single batch
async def mock_execute_async(query):
mock_result = MagicMock()
# Create mock rows for all plugins with timestamps
mock_rows = []
# oldest_plugin: oldest, middle_plugin: middle, newest_plugin: newest
mock_row1 = MagicMock()
mock_row1.plugin_author = 'author1'
mock_row1.plugin_name = 'oldest_plugin'
mock_row1.created_at = now - timedelta(days=10)
mock_rows.append(mock_row1)
mock_row2 = MagicMock()
mock_row2.plugin_author = 'author2'
mock_row2.plugin_name = 'middle_plugin'
mock_row2.created_at = now - timedelta(days=5)
mock_rows.append(mock_row2)
mock_row3 = MagicMock()
mock_row3.plugin_author = 'author3'
mock_row3.plugin_name = 'newest_plugin'
mock_row3.created_at = now
mock_rows.append(mock_row3)
# Make the result iterable
mock_result.__iter__ = lambda self: iter(mock_rows)
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
# Call list_plugins
result = await connector.list_plugins()
# Verify sorting: newest first
assert len(result) == 3
assert result[0]['manifest']['manifest']['metadata']['name'] == 'newest_plugin'
assert result[1]['manifest']['manifest']['metadata']['name'] == 'middle_plugin'
assert result[2]['manifest']['manifest']['metadata']['name'] == 'oldest_plugin'
@pytest.mark.asyncio
async def test_plugin_list_empty():
"""Test that empty plugin list is handled correctly."""
from langbot.pkg.plugin.connector import PluginRuntimeConnector
# Mock the application
mock_app = MagicMock()
mock_app.instance_config.data.get.return_value = {'enable': True}
mock_app.logger = MagicMock()
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
# Mock empty plugin list
connector.handler.list_plugins = AsyncMock(return_value=[])
# Call list_plugins
result = await connector.list_plugins()
# Verify empty list
assert len(result) == 0
+1
View File
@@ -0,0 +1 @@
"""Provider requester tests"""
+298
View File
@@ -0,0 +1,298 @@
"""
Test fixtures for provider/modelmgr tests.
Provides fake persistence, mock requester registry, and test utilities
without calling real LLM APIs or network requests.
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.provider.modelmgr import token
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.discover import engine as discover_engine
class FakeProviderAPIRequester(requester.ProviderAPIRequester):
"""Fake requester for testing that does not make real API calls."""
name = 'fake-requester'
default_config = {'base_url': 'https://fake-api.example.com', 'timeout': 30}
def __init__(self, ap, config: dict):
super().__init__(ap, config)
self._invoke_count = 0
self._last_messages = None
self._last_model = None
async def invoke_llm(
self,
query,
model: requester.RuntimeLLMModel,
messages: list,
funcs=None,
extra_args={},
remove_think=False,
):
"""Return a fake message response."""
self._invoke_count += 1
self._last_messages = messages
self._last_model = model
# Import the message entity for response
import langbot_plugin.api.entities.builtin.provider.message as provider_message
return provider_message.Message(
role='assistant',
content=[provider_message.ContentElement(type='text', text='Fake LLM response')],
)
async def invoke_llm_stream(
self,
query,
model: requester.RuntimeLLMModel,
messages: list,
funcs=None,
extra_args={},
remove_think=False,
):
"""Yield fake message chunks."""
import langbot_plugin.api.entities.builtin.provider.message as provider_message
yield provider_message.MessageChunk(
role='assistant',
content=[provider_message.ContentElement(type='text', text='Fake stream chunk')],
)
async def invoke_embedding(self, model, input_text: list, extra_args={}):
"""Return fake embedding vectors."""
return [[0.1, 0.2, 0.3] for _ in input_text]
async def invoke_rerank(self, model, query: str, documents: list, extra_args={}):
"""Return fake rerank results."""
return [{'index': i, 'relevance_score': 0.9 - i * 0.1} for i in range(len(documents))]
class AnotherFakeRequester(requester.ProviderAPIRequester):
"""Another fake requester for multi-requester tests."""
name = 'another-fake-requester'
default_config = {'base_url': 'https://another-fake.example.com'}
async def invoke_llm(self, query, model, messages, funcs=None, extra_args={}, remove_think=False):
import langbot_plugin.api.entities.builtin.provider.message as provider_message
return provider_message.Message(
role='assistant', content=[provider_message.ContentElement(type='text', text='Another response')]
)
async def invoke_rerank(self, model, query: str, documents: list, extra_args={}):
"""Return fake rerank results."""
return [{'index': i, 'relevance_score': 0.9 - i * 0.1} for i in range(len(documents))]
def _create_fake_component(name: str, requester_class: type) -> Mock:
"""Create a fake Component mock for a requester."""
# Use Mock to allow overriding get_python_component_class
component = Mock(spec=discover_engine.Component)
component.metadata = Mock()
component.metadata.name = name
component.get_python_component_class = Mock(return_value=requester_class)
return component
def _make_mock_result(items: list = None, first_item=None):
"""Create a mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
return result
def _make_row_mock(entity):
"""Create a mock Row-like object that can be unpacked via _mapping.
Note: This function returns the actual entity directly since Mock objects
don't pass isinstance(provider_info, sqlalchemy.Row) checks. The code
in modelmgr.load_provider handles this via the else branch.
"""
return entity
@pytest.fixture
def mock_app_for_modelmgr():
"""Provides a mock Application for ModelManager tests."""
app = SimpleNamespace()
app.logger = Mock()
app.logger.debug = Mock()
app.logger.info = Mock()
app.logger.warning = Mock()
app.logger.error = Mock()
# Fake persistence manager - returns empty results by default
app.persistence_mgr = SimpleNamespace()
async def default_execute(query):
return _make_mock_result([])
app.persistence_mgr.execute_async = AsyncMock(side_effect=default_execute)
# Fake discover engine
app.discover = SimpleNamespace()
app.discover.get_components_by_kind = Mock(return_value=[])
# Fake instance config
app.instance_config = SimpleNamespace()
app.instance_config.data = {'space': {'disable_models_service': True}}
# Other services (not used in basic tests)
app.space_service = AsyncMock()
app.llm_model_service = AsyncMock()
app.embedding_models_service = AsyncMock()
app.monitoring_service = AsyncMock()
return app
@pytest.fixture
def fake_requester_registry(mock_app_for_modelmgr):
"""Provides a ModelManager with fake requester registry."""
app = mock_app_for_modelmgr
# Create fake components
fake_component = _create_fake_component('fake-requester', FakeProviderAPIRequester)
another_component = _create_fake_component('another-fake-requester', AnotherFakeRequester)
app.discover.get_components_by_kind = Mock(return_value=[fake_component, another_component])
model_mgr = ModelManager(app)
return model_mgr
@pytest.fixture
def fake_persistence_data():
"""Provides fake persistence data for models and providers."""
provider_uuid = 'test-provider-uuid'
provider_uuid2 = 'test-provider-uuid-2'
providers = [
persistence_model.ModelProvider(
uuid=provider_uuid,
name='Test Provider',
requester='fake-requester',
base_url='https://test.example.com',
api_keys=['test-api-key-1', 'test-api-key-2'],
),
persistence_model.ModelProvider(
uuid=provider_uuid2,
name='Test Provider 2',
requester='another-fake-requester',
base_url='https://test2.example.com',
api_keys=['key-3'],
),
]
llm_models = [
persistence_model.LLMModel(
uuid='test-llm-uuid-1',
name='TestLLM-1',
provider_uuid=provider_uuid,
abilities=['func_call'],
extra_args={'temperature': 0.7},
),
persistence_model.LLMModel(
uuid='test-llm-uuid-2',
name='TestLLM-2',
provider_uuid=provider_uuid,
abilities=['vision'],
extra_args={},
),
]
embedding_models = [
persistence_model.EmbeddingModel(
uuid='test-embedding-uuid-1',
name='TestEmbedding-1',
provider_uuid=provider_uuid,
extra_args={'dimensions': 768},
),
]
rerank_models = [
persistence_model.RerankModel(
uuid='test-rerank-uuid-1',
name='TestRerank-1',
provider_uuid=provider_uuid2,
extra_args={},
),
]
return {
'providers': providers,
'llm_models': llm_models,
'embedding_models': embedding_models,
'rerank_models': rerank_models,
'provider_uuid': provider_uuid,
'provider_uuid2': provider_uuid2,
}
@pytest.fixture
def runtime_provider(fake_persistence_data, mock_app_for_modelmgr):
"""Provides a RuntimeProvider instance for testing."""
provider_entity = fake_persistence_data['providers'][0]
token_mgr = token.TokenManager(name=provider_entity.uuid, tokens=provider_entity.api_keys or [])
requester_inst = FakeProviderAPIRequester(mock_app_for_modelmgr, {'base_url': provider_entity.base_url})
return requester.RuntimeProvider(
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
)
@pytest.fixture
def runtime_llm_model(fake_persistence_data, runtime_provider):
"""Provides a RuntimeLLMModel instance for testing."""
model_entity = fake_persistence_data['llm_models'][0]
return requester.RuntimeLLMModel(
model_entity=model_entity,
provider=runtime_provider,
)
@pytest.fixture
def runtime_embedding_model(fake_persistence_data, runtime_provider):
"""Provides a RuntimeEmbeddingModel instance for testing."""
model_entity = fake_persistence_data['embedding_models'][0]
return requester.RuntimeEmbeddingModel(
model_entity=model_entity,
provider=runtime_provider,
)
@pytest.fixture
def runtime_rerank_model(fake_persistence_data, mock_app_for_modelmgr):
"""Provides a RuntimeRerankModel instance for testing."""
provider_entity = fake_persistence_data['providers'][1]
token_mgr = token.TokenManager(name=provider_entity.uuid, tokens=provider_entity.api_keys or [])
requester_inst = AnotherFakeRequester(mock_app_for_modelmgr, {'base_url': provider_entity.base_url})
provider = requester.RuntimeProvider(
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
)
model_entity = fake_persistence_data['rerank_models'][0]
return requester.RuntimeRerankModel(
model_entity=model_entity,
provider=provider,
)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,93 @@
"""Unit tests for LiteLLMRequester._convert_messages.
Focus: the content-part normalization that (a) converts image_base64 parts to
the OpenAI image_url shape and (b) drops non-image file parts (file_base64 /
file_url) which OpenAI-compatible chat models reject. The latter is essential
for Voice/File attachments — including ones replayed from conversation history —
since the agent consumes their bytes via the sandbox, not the model payload.
"""
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
def _make_requester() -> LiteLLMRequester:
# _convert_messages does not touch instance config, so bypass __init__.
return LiteLLMRequester.__new__(LiteLLMRequester)
def test_convert_messages_drops_file_base64_part():
req = _make_requester()
msg = provider_message.Message(
role='user',
content=[
provider_message.ContentElement.from_text('analyze this audio'),
provider_message.ContentElement.from_file_base64('data:audio/wav;base64,AAAA', 'voice.wav'),
],
)
out = req._convert_messages([msg])
parts = out[0]['content']
types = [p.get('type') for p in parts]
assert 'file_base64' not in types
assert types == ['text']
assert parts[0]['text'] == 'analyze this audio'
def test_convert_messages_drops_file_url_part():
req = _make_requester()
msg = provider_message.Message(
role='user',
content=[
provider_message.ContentElement.from_text('here is a doc'),
provider_message.ContentElement.from_file_url('http://example.com/report.xlsx', 'report.xlsx'),
],
)
out = req._convert_messages([msg])
types = [p.get('type') for p in out[0]['content']]
assert types == ['text']
def test_convert_messages_keeps_image_and_converts_to_image_url():
req = _make_requester()
msg = provider_message.Message(
role='user',
content=[
provider_message.ContentElement.from_text('look'),
provider_message.ContentElement.from_image_base64('data:image/png;base64,AAAA'),
],
)
out = req._convert_messages([msg])
parts = out[0]['content']
types = [p.get('type') for p in parts]
# image is preserved and reshaped to the OpenAI image_url form
assert types == ['text', 'image_url']
img_part = parts[1]
assert img_part['image_url'] == {'url': 'data:image/png;base64,AAAA'}
assert 'image_base64' not in img_part
def test_convert_messages_mixed_history_strips_only_files():
req = _make_requester()
# Simulate replayed history: an old voice turn + a current text turn.
history_voice = provider_message.Message(
role='user',
content=[
provider_message.ContentElement.from_text('old audio turn'),
provider_message.ContentElement.from_file_base64('data:audio/wav;base64,BBBB', 'voice.wav'),
],
)
current = provider_message.Message(
role='user',
content=[provider_message.ContentElement.from_text('now do the csv')],
)
out = req._convert_messages([history_voice, current])
assert [p.get('type') for p in out[0]['content']] == ['text']
assert [p.get('type') for p in out[1]['content']] == ['text']
def test_convert_messages_plain_string_content_untouched():
req = _make_requester()
msg = provider_message.Message(role='user', content='just text')
out = req._convert_messages([msg])
assert out[0]['content'] == 'just text'
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,146 @@
"""Unit tests for LocalAgentRunner._inject_inbound_attachments.
Covers the user -> sandbox attachment path added for the Box attachment
round-trip:
* materialized descriptors are stashed on the query and described to the model
via an appended text note (in-sandbox paths + outbox convention);
* non-image file parts (file_base64 / file_url) are stripped from the user
message content because OpenAI-compatible chat models reject them, while
image and text parts are kept for vision models;
* the helper is a no-op when the box service is unavailable or yields nothing,
and never raises into the chat turn on materialization failure.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
def _make_runner(box_service) -> LocalAgentRunner:
runner = LocalAgentRunner.__new__(LocalAgentRunner)
runner.ap = SimpleNamespace(logger=Mock(), box_service=box_service)
return runner
def _make_query():
return SimpleNamespace(variables={}, query_id='q-123')
def _box_service(attachments):
svc = SimpleNamespace(
available=True,
OUTBOX_MOUNT_DIR='/outbox',
materialize_inbound_attachments=AsyncMock(return_value=attachments),
)
return svc
@pytest.mark.asyncio
async def test_inject_strips_file_parts_and_appends_note():
box = _box_service([{'type': 'Voice', 'path': '/inbox/q-123/voice.wav', 'size': 176000}])
runner = _make_runner(box)
query = _make_query()
user_message = provider_message.Message(
role='user',
content=[
provider_message.ContentElement.from_text('transcribe this'),
provider_message.ContentElement.from_file_base64('data:audio/wav;base64,AAAA', 'voice.wav'),
],
)
await runner._inject_inbound_attachments(query, user_message)
types = [getattr(ce, 'type', None) for ce in user_message.content]
# file_base64 dropped; text kept; sandbox-path note appended as text
assert 'file_base64' not in types
assert types.count('text') == 2
note = user_message.content[-1].text
assert '/inbox/q-123/voice.wav' in note
assert '/outbox/q-123' in note
# descriptors stashed for downstream stages
assert query.variables['_sandbox_inbound_attachments'] == box.materialize_inbound_attachments.return_value
@pytest.mark.asyncio
async def test_inject_keeps_image_parts():
box = _box_service([{'type': 'Image', 'path': '/inbox/q-123/pic.png', 'size': 1234}])
runner = _make_runner(box)
query = _make_query()
user_message = provider_message.Message(
role='user',
content=[
provider_message.ContentElement.from_text('what is this'),
provider_message.ContentElement.from_image_base64('data:image/png;base64,iVBORw0K'),
],
)
await runner._inject_inbound_attachments(query, user_message)
types = [getattr(ce, 'type', None) for ce in user_message.content]
assert 'image_base64' in types # vision part preserved
assert types[-1] == 'text' # note appended last
@pytest.mark.asyncio
async def test_inject_promotes_string_content_to_list_with_note():
box = _box_service([{'type': 'File', 'path': '/inbox/q-123/data.csv', 'size': 42}])
runner = _make_runner(box)
query = _make_query()
user_message = provider_message.Message(role='user', content='clean this csv')
await runner._inject_inbound_attachments(query, user_message)
assert isinstance(user_message.content, list)
assert [getattr(ce, 'type', None) for ce in user_message.content] == ['text', 'text']
assert user_message.content[0].text == 'clean this csv'
assert '/inbox/q-123/data.csv' in user_message.content[1].text
@pytest.mark.asyncio
async def test_inject_noop_without_box_service():
runner = _make_runner(box_service=None)
query = _make_query()
user_message = provider_message.Message(role='user', content='hello')
await runner._inject_inbound_attachments(query, user_message)
assert user_message.content == 'hello'
assert '_sandbox_inbound_attachments' not in query.variables
@pytest.mark.asyncio
async def test_inject_noop_when_no_attachments():
box = _box_service([])
runner = _make_runner(box)
query = _make_query()
user_message = provider_message.Message(role='user', content='hello')
await runner._inject_inbound_attachments(query, user_message)
assert user_message.content == 'hello'
assert '_sandbox_inbound_attachments' not in query.variables
@pytest.mark.asyncio
async def test_inject_swallows_materialization_error():
box = SimpleNamespace(
available=True,
OUTBOX_MOUNT_DIR='/outbox',
materialize_inbound_attachments=AsyncMock(side_effect=RuntimeError('disk full')),
)
runner = _make_runner(box)
query = _make_query()
user_message = provider_message.Message(role='user', content='hello')
# must not raise
await runner._inject_inbound_attachments(query, user_message)
assert user_message.content == 'hello'
runner.ap.logger.warning.assert_called_once()
@@ -0,0 +1,281 @@
from __future__ import annotations
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.provider.runners.localagent import LocalAgentRunner, _StreamAccumulator
class RecordingProvider:
def __init__(self):
self.requests: list[dict] = []
async def invoke_llm(self, query, model, messages, funcs, extra_args=None, remove_think=None):
self.requests.append(
{
'messages': list(messages),
'funcs': list(funcs),
'remove_think': remove_think,
}
)
if len(self.requests) == 1:
return provider_message.Message(
role='assistant',
content='Let me calculate that exactly.',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(
name='exec',
arguments=json.dumps(
{'command': ("python - <<'PY'\nnums = [1, 2, 3, 4]\nprint(sum(nums) / len(nums))\nPY")}
),
),
)
],
)
tool_result = json.loads(messages[-1].content)
return provider_message.Message(
role='assistant',
content=f'The average is {tool_result["stdout"]}.',
)
class RecordingStreamProvider:
def __init__(self):
self.stream_requests: list[dict] = []
def invoke_llm_stream(self, query, model, messages, funcs, extra_args=None, remove_think=None):
self.stream_requests.append(
{
'messages': list(messages),
'funcs': list(funcs),
'remove_think': remove_think,
}
)
async def _stream():
if len(self.stream_requests) == 1:
yield provider_message.MessageChunk(
role='assistant',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(
name='exec',
arguments=json.dumps({'command': "python -c 'print(1)'"}),
),
)
],
is_final=True,
)
return
yield provider_message.MessageChunk(
role='assistant',
content='Tool execution failed.',
is_final=True,
)
return _stream()
def make_query() -> pipeline_query.Query:
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
return pipeline_query.Query.model_construct(
query_id='avg-query',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=12345,
message_chain=[],
message_event=None,
adapter=adapter,
pipeline_uuid='pipeline-uuid',
bot_uuid='bot-uuid',
pipeline_config={
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {'model': {'primary': 'test-model-uuid', 'fallbacks': []}, 'prompt': 'test-prompt'},
},
'output': {'misc': {'remove-think': False}},
},
prompt=SimpleNamespace(messages=[]),
messages=[],
user_message=provider_message.Message(
role='user',
content='Please calculate the average of 1, 2, 3, and 4.',
),
use_funcs=[SimpleNamespace(name='exec')],
use_llm_model_uuid='test-model-uuid',
variables={},
)
def test_stream_accumulator_merges_fragmented_tool_call_arguments():
accumulator = _StreamAccumulator(msg_sequence=1)
assert (
accumulator.add(
provider_message.MessageChunk(
role='assistant',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(name='exec', arguments='{"command":'),
)
],
)
)
is None
)
emitted = accumulator.add(
provider_message.MessageChunk(
role='assistant',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(name='exec', arguments='"pwd"}'),
)
],
is_final=True,
)
)
assert emitted is not None
final_msg = accumulator.final_message()
assert final_msg.tool_calls[0].function.name == 'exec'
assert final_msg.tool_calls[0].function.arguments == '{"command":"pwd"}'
@pytest.mark.asyncio
async def test_localagent_uses_exec_for_exact_calculation():
provider = RecordingProvider()
model = SimpleNamespace(
provider=provider,
model_entity=SimpleNamespace(
uuid='test-model-uuid',
name='test-model',
abilities=['func_call'],
extra_args={},
),
)
tool_manager = SimpleNamespace(
execute_func_call=AsyncMock(
return_value={
'session_id': 'avg-query',
'backend': 'podman',
'status': 'completed',
'ok': True,
'exit_code': 0,
'stdout': '2.5',
'stderr': '',
'duration_ms': 18,
}
)
)
app = SimpleNamespace(
logger=Mock(),
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
tool_mgr=tool_manager,
rag_mgr=SimpleNamespace(),
box_service=SimpleNamespace(
get_system_guidance=Mock(
return_value=(
'When the exec tool is available, use it for exact calculations, statistics, '
'structured data parsing, and code execution instead of estimating mentally. '
'Unless the user explicitly asks for the script, code, or implementation details, '
'do not include the generated script in the final answer. '
'A default workspace is mounted at /workspace for file tasks.'
)
),
),
skill_mgr=SimpleNamespace(
get_skills_for_pipeline=AsyncMock(return_value=[]),
detect_skill_activation=AsyncMock(return_value=None),
build_activation_prompt=Mock(return_value=None),
),
)
runner = LocalAgentRunner(app, pipeline_config={})
query = make_query()
results = [message async for message in runner.run(query)]
assert [message.role for message in results] == ['assistant', 'tool', 'assistant']
assert results[-1].content == 'The average is 2.5.'
tool_manager.execute_func_call.assert_awaited_once()
tool_name, tool_parameters = tool_manager.execute_func_call.await_args.args[:2]
assert tool_name == 'exec'
assert 'print(sum(nums) / len(nums))' in tool_parameters['command']
first_request = provider.requests[0]
assert any(
message.role == 'system'
and 'exec' in str(message.content)
and 'exact calculations' in str(message.content)
and 'Unless the user explicitly asks for the script' in str(message.content)
and '/workspace' in str(message.content)
for message in first_request['messages']
)
assert [tool.name for tool in first_request['funcs']] == ['exec']
@pytest.mark.asyncio
async def test_localagent_streaming_tool_error_yields_message_chunks():
provider = RecordingStreamProvider()
model = SimpleNamespace(
provider=provider,
model_entity=SimpleNamespace(
uuid='test-model-uuid',
name='test-model',
abilities=['func_call'],
extra_args={},
),
)
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=True)
query = make_query()
query.adapter = adapter
app = SimpleNamespace(
logger=Mock(),
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
tool_mgr=SimpleNamespace(execute_func_call=AsyncMock(side_effect=RuntimeError('boom'))),
rag_mgr=SimpleNamespace(),
box_service=SimpleNamespace(
get_system_guidance=Mock(return_value='sandbox guidance'),
),
skill_mgr=SimpleNamespace(
get_skills_for_pipeline=AsyncMock(return_value=[]),
detect_skill_activation=AsyncMock(return_value=None),
build_activation_prompt=Mock(return_value=None),
),
)
runner = LocalAgentRunner(app, pipeline_config={})
results = [message async for message in runner.run(query)]
assert all(isinstance(message, provider_message.MessageChunk) for message in results)
assert any(message.role == 'tool' and message.content == 'err: boom' for message in results)
@@ -0,0 +1,957 @@
"""Tests for MCP Box integration: path rewriting, host_path inference, config model, payloads.
Uses importlib.util.spec_from_file_location to load mcp.py directly without
triggering the circular import chain through the app module.
"""
from __future__ import annotations
import importlib
import importlib.util
import os
import sys
import tempfile
import types
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
# ---------------------------------------------------------------------------
# Load mcp.py directly from file path, with stub dependencies
# ---------------------------------------------------------------------------
def _stub_module(fqn: str, attrs: dict | None = None, is_package: bool = False):
"""Create or return a stub module and register it in sys.modules."""
if fqn in sys.modules:
mod = sys.modules[fqn]
else:
mod = types.ModuleType(fqn)
mod.__spec__ = importlib.machinery.ModuleSpec(fqn, None, is_package=is_package)
if is_package:
mod.__path__ = []
sys.modules[fqn] = mod
parts = fqn.rsplit('.', 1)
if len(parts) == 2 and parts[0] in sys.modules:
setattr(sys.modules[parts[0]], parts[1], mod)
if attrs:
for k, v in attrs.items():
setattr(mod, k, v)
return mod
@pytest.fixture(scope='module', autouse=True)
def mcp_module():
"""Load mcp.py with minimal stubs to avoid circular imports."""
saved = {}
def _save_and_stub(name, attrs=None, is_package=False):
saved[name] = sys.modules.get(name)
# Don't overwrite modules that already exist (from other test modules)
if name in sys.modules:
return
_stub_module(name, attrs, is_package)
# Stub entire dependency chains as packages / modules
_save_and_stub('langbot_plugin', is_package=True)
_save_and_stub('langbot_plugin.api', is_package=True)
_save_and_stub('langbot_plugin.api.entities', is_package=True)
_save_and_stub('langbot_plugin.api.entities.events', is_package=True)
_save_and_stub('langbot_plugin.api.entities.events.pipeline_query', {})
_save_and_stub('langbot_plugin.api.entities.builtin', is_package=True)
_save_and_stub('langbot_plugin.api.entities.builtin.resource', is_package=True)
_save_and_stub(
'langbot_plugin.api.entities.builtin.resource.tool',
{
'LLMTool': type('LLMTool', (), {}),
},
)
_save_and_stub('langbot_plugin.api.entities.builtin.provider', is_package=True)
_save_and_stub('langbot_plugin.api.entities.builtin.provider.message', {})
_save_and_stub('sqlalchemy', {'select': Mock()})
_save_and_stub('httpx', {'AsyncClient': Mock()})
_save_and_stub('mcp', {'ClientSession': Mock, 'StdioServerParameters': Mock}, is_package=True)
_save_and_stub('mcp.client', is_package=True)
_save_and_stub('mcp.client.stdio', {'stdio_client': Mock()})
_save_and_stub('mcp.client.sse', {'sse_client': Mock()})
_save_and_stub('mcp.client.streamable_http', {'streamable_http_client': Mock()})
_save_and_stub('mcp.client.websocket', {'websocket_client': Mock()})
# Stub the provider.tools.loader (source of circular import)
_save_and_stub('langbot', is_package=True)
_save_and_stub('langbot.pkg', is_package=True)
_save_and_stub('langbot.pkg.provider', is_package=True)
_save_and_stub('langbot.pkg.provider.tools', is_package=True)
_save_and_stub(
'langbot.pkg.provider.tools.loader',
{
'ToolLoader': type('ToolLoader', (), {'__init__': lambda self, ap: None}),
},
)
_save_and_stub('langbot.pkg.provider.tools.loaders', is_package=True)
_save_and_stub('langbot.pkg.core', is_package=True)
_save_and_stub('langbot.pkg.core.app', {'Application': type('Application', (), {})})
_save_and_stub('langbot.pkg.entity', is_package=True)
_save_and_stub('langbot.pkg.entity.persistence', is_package=True)
_save_and_stub('langbot.pkg.entity.persistence.mcp', {})
# box models
import enum as _enum
class _BPS(str, _enum.Enum):
RUNNING = 'running'
EXITED = 'exited'
_save_and_stub('langbot_plugin.box', is_package=True)
_save_and_stub('langbot_plugin.box.models', {'BoxManagedProcessStatus': _BPS})
# Now load mcp.py via spec_from_file_location
mod_fqn = 'langbot.pkg.provider.tools.loaders.mcp'
sys.modules.pop(mod_fqn, None)
mcp_path = os.path.join(
os.path.dirname(__file__),
'..',
'..',
'..',
'src',
'langbot',
'pkg',
'provider',
'tools',
'loaders',
'mcp.py',
)
mcp_path = os.path.normpath(mcp_path)
pkg_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(mcp_path))))
sys.modules['langbot.pkg'].__path__ = [pkg_root]
sys.modules['langbot.pkg.provider.tools.loaders'].__path__ = [os.path.dirname(mcp_path)]
spec = importlib.util.spec_from_file_location(mod_fqn, mcp_path)
mod = importlib.util.module_from_spec(spec)
sys.modules[mod_fqn] = mod
spec.loader.exec_module(mod)
yield mod
# Cleanup
sys.modules.pop(mod_fqn, None)
sys.modules.pop('langbot.pkg.provider.tools.loaders.mcp_stdio', None)
sys.modules.pop('langbot.pkg.box.workspace', None)
for name in reversed(list(saved)):
if saved[name] is None:
sys.modules.pop(name, None)
else:
sys.modules[name] = saved[name]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_ap():
ap = Mock()
ap.logger = Mock()
ap.box_service = Mock()
return ap
def _make_session(mcp_module, server_config: dict, ap=None):
if ap is None:
ap = _make_ap()
return mcp_module.RuntimeMCPSession(
server_name=server_config.get('name', 'test-server'),
server_config=server_config,
enable=True,
ap=ap,
)
# ── MCPServerBoxConfig ──────────────────────────────────────────────
class TestMCPServerBoxConfig:
def test_default_values(self, mcp_module):
cfg = mcp_module.MCPServerBoxConfig.model_validate({})
assert cfg.image is None
assert cfg.network == 'on'
assert cfg.host_path is None
assert cfg.host_path_mode == 'ro'
assert cfg.env == {}
assert cfg.startup_timeout_sec == 300
assert cfg.cpus is None
assert cfg.memory_mb is None
assert cfg.pids_limit is None
assert cfg.read_only_rootfs is None
def test_custom_values(self, mcp_module):
cfg = mcp_module.MCPServerBoxConfig.model_validate(
{
'image': 'node:20',
'network': 'on',
'host_path': '/home/user/mcp',
'host_path_mode': 'rw',
'env': {'FOO': 'bar'},
'startup_timeout_sec': 60,
'cpus': 2.0,
'memory_mb': 1024,
'pids_limit': 256,
'read_only_rootfs': False,
}
)
assert cfg.image == 'node:20'
assert cfg.network == 'on'
assert cfg.cpus == 2.0
assert cfg.memory_mb == 1024
def test_extra_fields_ignored(self, mcp_module):
cfg = mcp_module.MCPServerBoxConfig.model_validate(
{
'image': 'node:20',
'unknown_field': 'whatever',
}
)
assert cfg.image == 'node:20'
assert not hasattr(cfg, 'unknown_field')
# ── Path Rewriting ──────────────────────────────────────────────────
class TestRewritePath:
def test_no_host_path_returns_unchanged(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
assert s._rewrite_path('/some/path', None) == '/some/path'
def test_empty_path_returns_empty(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
assert s._rewrite_path('', '/home/user/mcp') == ''
def test_prefix_match_rewrites(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
result = s._rewrite_path('/home/user/mcp/server.py', '/home/user/mcp')
assert result == '/workspace/server.py'
def test_exact_match_rewrites_to_workspace(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
result = s._rewrite_path('/home/user/mcp', '/home/user/mcp')
assert result == '/workspace'
def test_non_matching_path_unchanged(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
result = s._rewrite_path('/opt/other/server.py', '/home/user/mcp')
assert result == '/opt/other/server.py'
def test_similar_prefix_not_rewritten(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
result = s._rewrite_path('/home/user/mcp-other/file.py', '/home/user/mcp')
assert result == '/home/user/mcp-other/file.py'
def test_nested_subpath_rewrites(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
result = s._rewrite_path('/home/user/mcp/src/lib/main.py', '/home/user/mcp')
assert result == '/workspace/src/lib/main.py'
# ── host_path Inference ─────────────────────────────────────────────
class TestInferHostPath:
def test_no_absolute_paths_returns_none(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': ['server.py'],
},
)
assert s._infer_host_path() is None
def test_nonexistent_path_returns_none(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': '/nonexistent/path/to/python',
'args': [],
},
)
assert s._infer_host_path() is None
def test_existing_absolute_path_infers_directory(self, mcp_module):
with tempfile.NamedTemporaryFile(suffix='.py') as f:
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [f.name],
},
)
result = s._infer_host_path()
assert result is not None
assert result == os.path.dirname(os.path.realpath(f.name))
# ── Build Box Session Payload ───────────────────────────────────────
class TestBuildBoxSessionPayload:
def test_minimal_config(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
payload = s._build_box_session_payload('session-123')
assert payload['session_id'] == 'session-123'
assert payload['workdir'] == '/workspace'
assert payload['env'] == {}
assert 'host_path' not in payload
def test_with_host_path(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
'box': {'host_path': '/home/user/mcp', 'host_path_mode': 'ro'},
},
)
payload = s._build_box_session_payload('session-123')
assert payload['host_path'] == '/home/user/mcp'
assert payload['host_path_mode'] == 'ro'
def test_optional_fields_included_when_set(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
'box': {'image': 'node:20', 'cpus': 2.0, 'memory_mb': 1024, 'pids_limit': 256},
},
)
payload = s._build_box_session_payload('session-123')
assert payload['image'] == 'node:20'
assert payload['cpus'] == 2.0
assert payload["memory_mb"] == 1024
assert payload['pids_limit'] == 256
def test_none_fields_excluded(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
payload = s._build_box_session_payload('session-123')
assert 'image' not in payload
assert 'cpus' not in payload
# ── Build Box Process Payload ───────────────────────────────────────
class TestBuildBoxProcessPayload:
def test_basic_payload(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': ['server.py'],
'env': {'KEY': 'val'},
},
)
payload = s._build_box_process_payload()
assert payload['command'] == 'python'
assert payload['args'] == ['server.py']
assert payload['env'] == {'KEY': 'val'}
assert payload['cwd'] == '/workspace'
def test_path_rewriting_applied(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': '/home/user/mcp/venv/bin/python',
'args': ['/home/user/mcp/server.py', '--config', '/home/user/mcp/config.json'],
'env': {},
'box': {'host_path': '/home/user/mcp'},
},
)
payload = s._build_box_process_payload()
# venv python is replaced with plain 'python' (deps installed in-container)
assert payload['command'] == 'python'
assert payload['args'] == ['/workspace/server.py', '--config', '/workspace/config.json']
def test_non_matching_args_not_rewritten(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': ['/opt/other/server.py', '--flag'],
'env': {},
'box': {'host_path': '/home/user/mcp'},
},
)
payload = s._build_box_process_payload()
assert payload['command'] == 'python'
assert payload['args'] == ['/opt/other/server.py', '--flag']
# ── Python Workspace Preparation ────────────────────────────────────
class TestPythonWorkspacePreparation:
def test_requirements_workspace_uses_venv_bootstrap(self, mcp_module, tmp_path):
host_path = tmp_path / 'mcp-source'
host_path.mkdir()
(host_path / 'requirements.txt').write_text('mcp==1.26.0\n', encoding='utf-8')
command = mcp_module.BoxStdioSessionRuntime.detect_install_command(
str(host_path),
'/workspace/.mcp/u1/workspace',
)
assert command is not None
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
assert 'python -m pip install -r "/workspace/.mcp/u1/workspace/requirements.txt"' in command
assert 'pip install --no-cache-dir -r' not in command
def test_staging_refresh_removes_stale_source_files_but_preserves_runtime_dirs(self, mcp_module, tmp_path):
source = tmp_path / 'source'
source.mkdir()
(source / 'server.py').write_text('print("new")\n', encoding='utf-8')
(source / 'requirements.txt').write_text('mcp==1.26.0\n', encoding='utf-8')
(source / '.env').write_text('TOKEN=new\n', encoding='utf-8')
process_root = tmp_path / 'shared' / '.mcp' / 'u1'
workspace = process_root / 'workspace'
(workspace / '.venv' / 'bin').mkdir(parents=True)
(workspace / '.venv' / 'bin' / 'python').write_text('', encoding='utf-8')
(workspace / '.langbot').mkdir()
(workspace / '.langbot' / 'python-env.lock').mkdir()
(workspace / '.env').write_text('TOKEN=old\n', encoding='utf-8')
(workspace / 'server.py').write_text('print("old")\n', encoding='utf-8')
(workspace / 'removed.py').write_text('stale\n', encoding='utf-8')
(workspace / 'removed_dir').mkdir()
(workspace / 'removed_dir' / 'old.txt').write_text('stale\n', encoding='utf-8')
mcp_module.BoxStdioSessionRuntime._copy_workspace_tree(str(source), str(process_root), str(workspace))
assert (workspace / 'server.py').read_text(encoding='utf-8') == 'print("new")\n'
assert (workspace / 'requirements.txt').read_text(encoding='utf-8') == 'mcp==1.26.0\n'
assert (workspace / '.env').read_text(encoding='utf-8') == 'TOKEN=new\n'
assert not (workspace / 'removed.py').exists()
assert not (workspace / 'removed_dir').exists()
assert (workspace / '.venv' / 'bin' / 'python').exists()
assert (workspace / '.langbot' / 'python-env.lock').is_dir()
def test_staging_refresh_ignores_unlink_race(self, mcp_module, tmp_path, monkeypatch):
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
source = tmp_path / 'source'
source.mkdir()
(source / 'server.py').write_text('print("new")\n', encoding='utf-8')
process_root = tmp_path / 'shared' / '.mcp' / 'u1'
workspace = process_root / 'workspace'
workspace.mkdir(parents=True)
stale_file = workspace / 'removed.py'
stale_file.write_text('stale\n', encoding='utf-8')
real_unlink = os.unlink
def unlink_with_race(path):
if os.fspath(path) == str(stale_file):
real_unlink(path)
raise FileNotFoundError(path)
real_unlink(path)
monkeypatch.setattr(mcp_stdio_module.os, 'unlink', unlink_with_race)
mcp_module.BoxStdioSessionRuntime._copy_workspace_tree(str(source), str(process_root), str(workspace))
assert not stale_file.exists()
assert (workspace / 'server.py').read_text(encoding='utf-8') == 'print("new")\n'
# ── get_runtime_info_dict ───────────────────────────────────────────
class TestGetRuntimeInfoDict:
def test_non_stdio_session(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'test-uuid',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
info = s.get_runtime_info_dict()
assert info['status'] == 'connecting'
assert 'box_session_id' not in info
def test_runtime_tools_include_parameters(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'test-uuid',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
s.functions = [
SimpleNamespace(
name='create-service',
description='Create a service',
parameters={
'type': 'object',
'properties': {
'project_id': {'type': 'string'},
},
'required': ['project_id'],
},
)
]
info = s.get_runtime_info_dict()
assert info['tools'][0]['parameters']['properties']['project_id']['type'] == 'string'
assert info['tools'][0]['parameters']['required'] == ['project_id']
def test_stdio_session_includes_box_info(self, mcp_module):
ap = _make_ap()
ap.box_service.available = True
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'test-uuid',
'mode': 'stdio',
'command': 'python',
'args': [],
},
ap=ap,
)
info = s.get_runtime_info_dict()
assert info['box_session_id'] == 'mcp-shared'
assert info['box_enabled'] is True
def test_transient_test_shares_session_but_isolated_by_process(self, mcp_module):
"""A transient config-page "test" now shares the same 'mcp-shared' Box
session as live servers (so a test reuses the running container / live
process instead of a cold per-test session bootstrap). Isolation is at
the PROCESS level: the test runs under its own process_id and only ever
stops that process_id, so it cannot disturb another server's live
process or the shared session itself."""
ap = _make_ap()
ap.box_service.available = True
transient = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'gen-uuid-123',
'mode': 'stdio',
'command': 'uvx',
'args': ['mcp-server-time'],
'_transient': True,
},
ap=ap,
)
live = _make_session(
mcp_module,
{
'name': 'time',
'uuid': 'real-uuid',
'mode': 'stdio',
'command': 'uvx',
'args': ['mcp-server-time'],
},
ap=ap,
)
assert transient.is_transient is True
assert live.is_transient is False
# Both share ONE Box session ...
assert transient._build_box_session_id() == 'mcp-shared'
assert live._build_box_session_id() == 'mcp-shared'
assert transient._build_box_session_id() == live._build_box_session_id()
# ... but are isolated by distinct process_ids within that session.
assert transient._box_stdio_runtime.process_id != live._box_stdio_runtime.process_id
def test_stdio_session_refuses_when_box_unavailable(self, mcp_module):
"""Policy: when Box is configured but unavailable (disabled in config
OR connection failed), stdio MCP servers are NOT treated as box-stdio.
``_init_stdio_python_server`` will raise a clear refusal at start
time; until then, the runtime info simply omits box_session_id so the
UI can render the disabled state cleanly."""
ap = _make_ap()
ap.box_service.available = False
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'test-uuid',
'mode': 'stdio',
'command': 'python',
'args': [],
},
ap=ap,
)
info = s.get_runtime_info_dict()
assert 'box_session_id' not in info
assert 'box_enabled' not in info
def test_stdio_session_without_box_service_uses_local_stdio(self, mcp_module):
ap = _make_ap()
del ap.box_service
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'test-uuid',
'mode': 'stdio',
'command': 'python',
'args': [],
},
ap=ap,
)
info = s.get_runtime_info_dict()
assert 'box_session_id' not in info
# ── Box config parsing ──────────────────────────────────────────────
class TestBoxConfigParsing:
def test_box_config_parsed_from_server_config(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
'box': {'image': 'node:20', 'host_path': '/home/user/mcp'},
},
)
assert isinstance(s.box_config, mcp_module.MCPServerBoxConfig)
assert s.box_config.image == 'node:20'
assert s.box_config.host_path == '/home/user/mcp'
def test_missing_box_key_uses_defaults(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'sse',
'command': 'python',
'args': [],
},
)
assert isinstance(s.box_config, mcp_module.MCPServerBoxConfig)
assert s.box_config.image is None
assert s.box_config.host_path_mode == 'ro'
@pytest.mark.asyncio
async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_module, tmp_path):
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
class FakeClientSession:
def __init__(self, *_args):
pass
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def initialize(self):
return None
@asynccontextmanager
async def fake_websocket_client(_url: str):
yield ('read-stream', 'write-stream')
mcp_stdio_module.ClientSession = FakeClientSession
mcp_stdio_module.websocket_client = fake_websocket_client
ap = _make_ap()
ap.box_service.available = True
ap.box_service.default_workspace = str(tmp_path / 'shared-box-workspace')
ap.box_service.create_session = AsyncMock(return_value={})
ap.box_service.build_spec = Mock(return_value='validated-spec')
ap.box_service.client = SimpleNamespace(
execute=AsyncMock(return_value=SimpleNamespace(ok=True, stderr='', exit_code=0))
)
ap.box_service.start_managed_process = AsyncMock(return_value={})
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box.example/process')
host_path = tmp_path / 'mcp-source'
host_path.mkdir()
server_file = host_path / 'server.py'
server_file.write_text('print("hello")\n', encoding='utf-8')
session = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'u1',
'mode': 'stdio',
'command': str(host_path / '.venv' / 'bin' / 'python'),
'args': [str(server_file)],
'box': {'host_path': str(host_path)},
},
ap=ap,
)
await session._init_box_stdio_server()
await session.exit_stack.aclose()
assert ap.box_service.create_session.await_count == 1
session_payload = ap.box_service.create_session.await_args.args[0]
assert session_payload['session_id'] == 'mcp-shared'
assert 'host_path' not in session_payload
assert ap.box_service.build_spec.call_count == 1
assert ap.box_service.build_spec.call_args.kwargs.get('skip_host_mount_validation', False) is False
assert ap.box_service.build_spec.call_args.args[0]['host_path'] == str(host_path)
staged_file = tmp_path / 'shared-box-workspace' / '.mcp' / 'u1' / 'workspace' / 'server.py'
assert staged_file.read_text(encoding='utf-8') == 'print("hello")\n'
process_payload = ap.box_service.start_managed_process.await_args.args[1]
assert process_payload['process_id'] == 'u1'
assert process_payload['command'] == 'python'
assert process_payload['args'] == ['/workspace/.mcp/u1/workspace/server.py']
assert process_payload['cwd'] == '/workspace/.mcp/u1/workspace'
@pytest.mark.asyncio
async def test_stdio_handshake_raises_coldstart_retry_while_process_alive(mcp_module, tmp_path, monkeypatch):
"""During a slow (npx) cold start the handshake fails while the managed
process is still alive. initialize() must raise _ColdStartRetry (so the
outer lifecycle loop reuses the live process and retries without stopping it
or consuming the fatal budget), NOT a fatal error."""
from contextlib import asynccontextmanager
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
class ColdClientSession:
def __init__(self, *_args):
pass
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def initialize(self):
# Process still cold-starting: handshake fails.
raise Exception('Connection closed')
@asynccontextmanager
async def fake_websocket_client(_url: str):
yield ('read-stream', 'write-stream')
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', ColdClientSession)
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False)
ap = _make_ap()
ap.box_service.available = True
ap.box_service.create_session = AsyncMock(return_value={})
ap.box_service.start_managed_process = AsyncMock(return_value={})
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p')
session = _make_session(
mcp_module,
{
'name': 'slow',
'uuid': 'slow-uuid',
'mode': 'stdio',
'command': 'npx',
'args': ['-y', 'some-mcp'],
},
ap=ap,
)
# Process is NOT exited (still cold-starting) and not yet running for reuse.
async def _not_exited():
return False
session._box_stdio_runtime._managed_process_has_exited = _not_exited
async def _not_running():
return False
session._box_stdio_runtime._managed_process_is_running = _not_running
with pytest.raises(mcp_stdio_module._ColdStartRetry):
await session._init_box_stdio_server()
# Process was started exactly once (the retry will reuse it, not rebuild).
assert ap.box_service.start_managed_process.await_count == 1
await session.exit_stack.aclose()
@pytest.mark.asyncio
async def test_stdio_handshake_raises_fatal_when_process_exited(mcp_module, tmp_path, monkeypatch):
"""If the handshake fails AND the process has definitively exited, that is a
real failure — initialize() must NOT swallow it as a cold-start retry."""
from contextlib import asynccontextmanager
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
class DeadClientSession:
def __init__(self, *_args):
pass
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def initialize(self):
raise Exception('Connection closed')
@asynccontextmanager
async def fake_websocket_client(_url: str):
yield ('read-stream', 'write-stream')
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', DeadClientSession)
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False)
ap = _make_ap()
ap.box_service.available = True
ap.box_service.create_session = AsyncMock(return_value={})
ap.box_service.start_managed_process = AsyncMock(return_value={})
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p')
session = _make_session(
mcp_module,
{'name': 'dead', 'uuid': 'dead-uuid', 'mode': 'stdio', 'command': 'npx', 'args': ['-y', 'x']},
ap=ap,
)
async def _exited():
return True
session._box_stdio_runtime._managed_process_has_exited = _exited
async def _not_running():
return False
session._box_stdio_runtime._managed_process_is_running = _not_running
with pytest.raises(Exception) as ei:
await session._init_box_stdio_server()
assert not isinstance(ei.value, mcp_stdio_module._ColdStartRetry)
await session.exit_stack.aclose()
@@ -0,0 +1,244 @@
from __future__ import annotations
import asyncio
import json
from contextlib import asynccontextmanager
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock, Mock
import httpx
import pytest
from aiohttp import web
from mcp import types as mcp_types
from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
class _TransportProbe:
def __init__(self, streamable_status: int | None) -> None:
self.streamable_status = streamable_status
self.streamable_posts = 0
self.streamable_messages: list[str] = []
self.sse_gets = 0
self.sse_messages: list[str] = []
self.streamable_request_started = asyncio.Event()
self.release_streamable_request = asyncio.Event()
self._sse_response: web.StreamResponse | None = None
async def handle_mcp_endpoint(self, request: web.Request) -> web.StreamResponse:
if request.method == 'POST':
self.streamable_posts += 1
self.streamable_request_started.set()
if self.streamable_status is None:
await self.release_streamable_request.wait()
return web.Response(status=204)
if self.streamable_status == 200:
message = await request.json()
method = message.get('method', '')
self.streamable_messages.append(method)
if method == 'initialize':
return web.json_response(
{
'jsonrpc': '2.0',
'id': message['id'],
'result': {
'protocolVersion': mcp_types.LATEST_PROTOCOL_VERSION,
'capabilities': {'tools': {}},
'serverInfo': {'name': 'streamable-test', 'version': '1.0.0'},
},
}
)
if method == 'tools/list':
return web.json_response(
{
'jsonrpc': '2.0',
'id': message['id'],
'result': {
'tools': [
{
'name': 'echo',
'description': 'Echo test input',
'inputSchema': {'type': 'object'},
}
]
},
}
)
return web.Response(status=202)
return web.Response(status=self.streamable_status)
self.sse_gets += 1
response = web.StreamResponse(
status=200,
headers={
'Content-Type': 'text/event-stream',
'Cache-Control': 'no-cache',
},
)
await response.prepare(request)
self._sse_response = response
await response.write(b'event: endpoint\ndata: /messages?session_id=test-session\n\n')
try:
while request.transport is not None and not request.transport.is_closing():
await asyncio.sleep(0.05)
except asyncio.CancelledError:
raise
return response
async def handle_sse_message(self, request: web.Request) -> web.Response:
message = await request.json()
method = message.get('method', '')
self.sse_messages.append(method)
if method == 'initialize':
response_message = {
'jsonrpc': '2.0',
'id': message['id'],
'result': {
'protocolVersion': mcp_types.LATEST_PROTOCOL_VERSION,
'capabilities': {},
'serverInfo': {'name': 'legacy-sse-test', 'version': '1.0.0'},
},
}
assert self._sse_response is not None
payload = json.dumps(response_message, separators=(',', ':'))
await self._sse_response.write(f'event: message\ndata: {payload}\n\n'.encode())
return web.Response(status=202)
@asynccontextmanager
async def _transport_server(streamable_status: int | None):
probe = _TransportProbe(streamable_status)
application = web.Application()
application.router.add_route('*', '/mcp', probe.handle_mcp_endpoint)
application.router.add_post('/messages', probe.handle_sse_message)
runner = web.AppRunner(application, shutdown_timeout=0.1)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
server = cast(asyncio.Server, site._server)
port = server.sockets[0].getsockname()[1]
try:
yield probe, f'http://127.0.0.1:{port}/mcp'
finally:
await runner.cleanup()
def _session(url: str, *, timeout: float = 2) -> RuntimeMCPSession:
app = cast(Any, SimpleNamespace(logger=Mock()))
return RuntimeMCPSession(
'remote-transport-test',
{'uuid': 'srv-1', 'mode': 'remote', 'url': url, 'timeout': timeout},
True,
app,
)
def _contains_http_status(exc: BaseException, status_code: int) -> bool:
return any(
isinstance(leaf, httpx.HTTPStatusError) and leaf.response.status_code == status_code
for leaf in RuntimeMCPSession._iter_exception_leaves(exc)
)
async def _close_session(session: RuntimeMCPSession) -> None:
await session.exit_stack.aclose()
@pytest.mark.asyncio
async def test_remote_transport_real_streamable_http_success_keeps_session_usable():
async with _transport_server(200) as (probe, url):
session = _session(url)
try:
await session._init_remote_server()
assert session.session is not None
tools = await session.session.list_tools()
assert [tool.name for tool in tools.tools] == ['echo']
assert probe.streamable_posts >= 2
assert probe.streamable_messages[:2] == ['initialize', 'notifications/initialized']
assert 'tools/list' in probe.streamable_messages
assert probe.sse_gets == 0
finally:
await _close_session(session)
@pytest.mark.asyncio
@pytest.mark.parametrize('status_code', [400, 404, 405])
async def test_remote_transport_real_streamable_http_error_falls_back_to_legacy_sse(status_code: int):
async with _transport_server(status_code) as (probe, url):
session = _session(url)
try:
await session._init_remote_server()
assert session.session is not None
assert probe.streamable_posts == 1
assert probe.sse_gets == 1
assert 'initialize' in probe.sse_messages
finally:
await _close_session(session)
@pytest.mark.asyncio
@pytest.mark.parametrize('status_code', [401, 403, 406, 415, 429, 500])
async def test_remote_transport_real_non_compatibility_error_does_not_fallback(status_code: int):
async with _transport_server(status_code) as (probe, url):
session = _session(url)
try:
with pytest.raises(BaseException) as exc_info:
await session._init_remote_server()
assert _contains_http_status(exc_info.value, status_code)
assert probe.streamable_posts == 1
assert probe.sse_gets == 0
finally:
await _close_session(session)
@pytest.mark.asyncio
async def test_remote_transport_real_timeout_does_not_fallback():
async with _transport_server(None) as (probe, url):
session = _session(url, timeout=0.05)
try:
with pytest.raises(BaseException) as exc_info:
await session._init_remote_server()
assert any(
isinstance(leaf, httpx.TimeoutException)
for leaf in RuntimeMCPSession._iter_exception_leaves(exc_info.value)
)
assert probe.streamable_posts == 1
assert probe.sse_gets == 0
finally:
probe.release_streamable_request.set()
await _close_session(session)
@pytest.mark.asyncio
@pytest.mark.parametrize('error_type', [httpx.ConnectError, httpx.ConnectTimeout])
async def test_remote_transport_connection_errors_do_not_fallback(error_type: type[httpx.RequestError]):
request = httpx.Request('POST', 'https://unreachable.invalid/mcp')
error = error_type('connection failed', request=request)
session = _session(str(request.url))
session._init_streamable_http_server = AsyncMock(side_effect=error)
session._init_sse_server = AsyncMock()
with pytest.raises(type(error)) as exc_info:
await session._init_remote_server()
assert exc_info.value is error
session._init_sse_server.assert_not_awaited()
@pytest.mark.asyncio
async def test_remote_transport_external_cancellation_is_not_converted_to_sse_fallback():
async with _transport_server(None) as (probe, url):
session = _session(url)
task = asyncio.create_task(session._init_remote_server())
await asyncio.wait_for(probe.streamable_request_started.wait(), timeout=2)
task.cancel()
try:
with pytest.raises(asyncio.CancelledError):
await task
assert probe.sse_gets == 0
finally:
probe.release_streamable_request.set()
await _close_session(session)
@@ -0,0 +1,323 @@
from __future__ import annotations
import base64
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import httpx
import pytest
from mcp import types as mcp_types
from langbot.pkg.provider.tools.loaders.mcp import (
MCP_RESOURCE_CONTEXT_QUERY_KEY,
MCP_RESOURCE_TRACE_QUERY_KEY,
MCP_TOOL_LIST_RESOURCES,
MCP_TOOL_READ_RESOURCE,
MCPLoader,
MCPSessionStatus,
RuntimeMCPSession,
)
from langbot.pkg.telemetry import features as telemetry_features
def _app() -> SimpleNamespace:
return SimpleNamespace(logger=Mock())
def _connected_session(
*,
name: str = 'docs',
uuid: str = 'srv-1',
resources: list[dict] | None = None,
templates: list[dict] | None = None,
) -> RuntimeMCPSession:
session = RuntimeMCPSession(name, {'uuid': uuid, 'mode': 'remote'}, True, _app())
session.status = MCPSessionStatus.CONNECTED
session.session = SimpleNamespace(read_resource=AsyncMock())
session.resources = resources or [
{
'uri': 'file:///README.md',
'name': 'README.md',
'title': '',
'description': '',
'mime_type': 'text/markdown',
'size': None,
'icons': [],
'annotations': {},
'_meta': {},
}
]
session.resource_templates = templates or []
return session
def _query() -> SimpleNamespace:
return SimpleNamespace(variables={})
def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
request = httpx.Request('POST', 'https://example.com/mcp')
response = httpx.Response(status_code, request=request)
return httpx.HTTPStatusError(f'HTTP {status_code}', request=request, response=response)
@pytest.mark.asyncio
async def test_remote_transport_falls_back_to_sse_for_compatible_http_status_in_exception_group():
session = RuntimeMCPSession(
'remote',
{'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
True,
_app(),
)
session._init_streamable_http_server = AsyncMock(
side_effect=ExceptionGroup('transport failed', [_http_status_error(405)])
)
session._init_sse_server = AsyncMock()
await session._init_remote_server()
session._init_streamable_http_server.assert_awaited_once()
session._init_sse_server.assert_awaited_once()
@pytest.mark.asyncio
async def test_remote_transport_does_not_fallback_for_auth_http_status():
session = RuntimeMCPSession(
'remote',
{'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
True,
_app(),
)
error = _http_status_error(403)
session._init_streamable_http_server = AsyncMock(side_effect=error)
session._init_sse_server = AsyncMock()
with pytest.raises(httpx.HTTPStatusError):
await session._init_remote_server()
session._init_streamable_http_server.assert_awaited_once()
session._init_sse_server.assert_not_awaited()
@pytest.mark.asyncio
async def test_read_resource_envelope_truncates_caches_and_records_trace():
session = _connected_session()
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
contents=[
mcp_types.TextResourceContents(
uri='file:///README.md',
mimeType='text/markdown',
text='abcdef',
)
]
)
query = _query()
first = await session.read_resource_envelope(
'file:///README.md',
max_bytes=4,
source='ui_preview',
query=query,
)
second = await session.read_resource_envelope(
'file:///README.md',
max_bytes=4,
source='agent_tool',
query=query,
)
assert first['contents'][0]['text'] == 'abcd'
assert first['contents'][0]['bytes'] == 6
assert first['truncated'] is True
assert first['cache_hit'] is False
assert second['cache_hit'] is True
assert second['source'] == 'agent_tool'
assert session.session.read_resource.await_count == 1
traces = query.variables[MCP_RESOURCE_TRACE_QUERY_KEY]
assert [trace['source'] for trace in traces] == ['ui_preview', 'agent_tool']
assert traces[1]['cache_hit'] is True
assert query.variables[telemetry_features.FEATURES_KEY]['mcp_resource_reads'] == {
'ui_preview': 1,
'agent_tool': 1,
}
@pytest.mark.asyncio
async def test_read_resource_envelope_shares_byte_budget_across_text_contents():
session = _connected_session()
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
contents=[
mcp_types.TextResourceContents(
uri='file:///README.md#first',
mimeType='text/plain',
text='abc',
),
mcp_types.TextResourceContents(
uri='file:///README.md#second',
mimeType='text/plain',
text='def',
),
]
)
envelope = await session.read_resource_envelope('file:///README.md', max_bytes=4)
assert [item['text'] for item in envelope['contents']] == ['abc', 'd']
assert envelope['contents'][0]['truncated'] is False
assert envelope['contents'][1]['truncated'] is True
assert envelope['bytes'] == 6
assert envelope['truncated'] is True
@pytest.mark.asyncio
async def test_read_resource_envelope_omits_binary_by_default():
session = _connected_session(
resources=[
{
'uri': 'file:///image.png',
'name': 'image.png',
'title': '',
'description': '',
'mime_type': 'image/png',
'size': 4,
'icons': [],
'annotations': {},
'_meta': {},
}
]
)
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
contents=[
mcp_types.BlobResourceContents(
uri='file:///image.png',
mimeType='image/png',
blob=base64.b64encode(b'\x00\x01\x02\x03').decode(),
)
]
)
envelope = await session.read_resource_envelope('file:///image.png')
content = envelope['contents'][0]
assert content['type'] == 'blob'
assert content['blob'] is None
assert content['bytes'] == 4
assert content['binary_omitted'] is True
assert envelope['truncated'] is True
assert envelope['warnings'] == ['Binary resource content omitted from response.']
@pytest.mark.asyncio
async def test_read_resource_envelope_rejects_unlisted_uri():
session = _connected_session()
with pytest.raises(ValueError, match='Resource URI is not available'):
await session.read_resource_envelope('file:///secret.txt')
session.session.read_resource.assert_not_called()
def test_resource_uri_allowed_supports_listed_templates_conservatively():
session = _connected_session(
resources=[],
templates=[
{
'uri_template': 'repo://{owner}/{repo}/file/{path}',
'name': 'repository file',
'title': '',
'description': '',
'mime_type': 'text/plain',
'icons': [],
'annotations': {},
'_meta': {},
}
],
)
assert session.resource_uri_allowed('repo://langbot-app/LangBot/file/src/main.py') is True
assert session.resource_uri_allowed('repo://langbot-app/LangBot/issues/1') is False
assert session.resource_uri_allowed('https://example.com/secret') is False
@pytest.mark.asyncio
async def test_mcp_loader_can_hide_synthetic_resource_tools():
loader = MCPLoader(_app())
session = _connected_session()
loader.sessions = {'docs': session}
with_resource_tools = await loader.get_tools(['srv-1'], include_resource_tools=True)
without_resource_tools = await loader.get_tools(['srv-1'], include_resource_tools=False)
assert {tool.name for tool in with_resource_tools} == {
MCP_TOOL_LIST_RESOURCES,
MCP_TOOL_READ_RESOURCE,
}
assert without_resource_tools == []
@pytest.mark.asyncio
async def test_mcp_loader_refuses_resource_tool_calls_when_agent_read_disabled():
loader = MCPLoader(_app())
session = _connected_session()
loader.sessions = {'docs': session}
query = SimpleNamespace(
variables={
'_pipeline_bound_mcp_servers': ['srv-1'],
'_pipeline_mcp_resource_agent_read_enabled': False,
}
)
result = await loader.invoke_tool(
MCP_TOOL_READ_RESOURCE,
{'server_name': 'docs', 'uri': 'file:///README.md'},
query,
)
assert result[0].text == 'Error: MCP resource agent reads are disabled.'
session.session.read_resource.assert_not_called()
@pytest.mark.asyncio
async def test_build_resource_context_for_query_uses_only_bound_attached_text_resources():
loader = MCPLoader(_app())
docs = _connected_session(name='docs', uuid='srv-1')
docs.session.read_resource.return_value = mcp_types.ReadResourceResult(
contents=[
mcp_types.TextResourceContents(
uri='file:///README.md',
mimeType='text/markdown',
text='LangBot MCP resource context',
)
]
)
other = _connected_session(name='other', uuid='srv-2')
other.session.read_resource.return_value = mcp_types.ReadResourceResult(
contents=[
mcp_types.TextResourceContents(
uri='file:///README.md',
mimeType='text/markdown',
text='must not be injected',
)
]
)
loader.sessions = {'docs': docs, 'other': other}
query = SimpleNamespace(
variables={
'_pipeline_bound_mcp_servers': ['srv-1'],
'_pipeline_mcp_resource_attachments': [
{'server_uuid': 'srv-1', 'server_name': 'docs', 'uri': 'file:///README.md', 'mode': 'pinned'},
{'server_uuid': 'srv-2', 'server_name': 'other', 'uri': 'file:///README.md', 'mode': 'pinned'},
],
}
)
context = await loader.build_resource_context_for_query(query)
assert '<mcp_resource ' in context
assert 'server="docs"' in context
assert 'LangBot MCP resource context' in context
assert 'must not be injected' not in context
assert query.variables[MCP_RESOURCE_CONTEXT_QUERY_KEY]['resource_count'] == 1
docs.session.read_resource.assert_awaited_once()
other.session.read_resource.assert_not_called()
@@ -0,0 +1,798 @@
"""
Unit tests for ModelManager in provider/modelmgr.
Tests model configuration management, requester selection, provider loading,
and error handling without calling real LLM APIs.
"""
from __future__ import annotations
import pytest
from unittest.mock import Mock
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.entity.errors import provider as provider_errors
from langbot.pkg.provider.modelmgr import token
from tests.unit_tests.provider.conftest import _make_mock_result, _make_row_mock
# ============================================================================
# ModelManager Initialization Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_initialize_with_fake_requesters(fake_requester_registry):
"""Test ModelManager initializes with fake requester registry."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
assert 'fake-requester' in model_mgr.requester_dict
assert 'another-fake-requester' in model_mgr.requester_dict
assert model_mgr.requester_dict['fake-requester'] is not None
assert len(model_mgr.requester_components) == 2
@pytest.mark.asyncio
async def test_model_manager_initialize_empty_registry(mock_app_for_modelmgr):
"""Test ModelManager handles empty requester registry."""
app = mock_app_for_modelmgr
app.discover.get_components_by_kind = Mock(return_value=[])
model_mgr = ModelManager(app)
await model_mgr.initialize()
assert model_mgr.requester_dict == {}
assert len(model_mgr.requester_components) == 0
@pytest.mark.asyncio
async def test_model_manager_skips_space_sync_when_disabled(mock_app_for_modelmgr):
"""Test ModelManager skips space sync when disabled in config."""
app = mock_app_for_modelmgr
app.instance_config.data = {'space': {'disable_models_service': True}}
model_mgr = ModelManager(app)
await model_mgr.initialize()
# Should not call space_service if disabled
app.space_service.get_models.assert_not_called()
# ============================================================================
# Model Loading Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_load_models_from_db(fake_requester_registry, fake_persistence_data):
"""Test ModelManager loads models from database correctly."""
model_mgr = fake_requester_registry
# Setup fake persistence responses - return entities directly (code handles non-Row entities)
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'llm_models' in query_str:
return _make_mock_result(fake_persistence_data['llm_models'])
elif 'embedding_models' in query_str:
return _make_mock_result(fake_persistence_data['embedding_models'])
elif 'rerank_models' in query_str:
return _make_mock_result(fake_persistence_data['rerank_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
# Check providers loaded
assert len(model_mgr.provider_dict) == 2
assert fake_persistence_data['provider_uuid'] in model_mgr.provider_dict
assert fake_persistence_data['provider_uuid2'] in model_mgr.provider_dict
# Check models loaded
assert len(model_mgr.llm_models) == 2
assert len(model_mgr.embedding_models) == 1
assert len(model_mgr.rerank_models) == 1
@pytest.mark.asyncio
async def test_model_manager_load_provider_unknown_requester(mock_app_for_modelmgr):
"""Test ModelManager raises RequesterNotFoundError for unknown requester."""
app = mock_app_for_modelmgr
app.discover.get_components_by_kind = Mock(return_value=[])
model_mgr = ModelManager(app)
await model_mgr.initialize()
provider_info = {
'uuid': 'unknown-provider',
'name': 'Unknown Provider',
'requester': 'non-existent-requester',
'base_url': 'https://unknown.com',
'api_keys': [],
}
with pytest.raises(provider_errors.RequesterNotFoundError) as exc_info:
await model_mgr.load_provider(provider_info)
assert exc_info.value.requester_name == 'non-existent-requester'
@pytest.mark.asyncio
async def test_model_manager_load_provider_from_dict(fake_requester_registry):
"""Test ModelManager loads provider from dict correctly."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
provider_info = {
'uuid': 'dict-provider-uuid',
'name': 'Dict Provider',
'requester': 'fake-requester',
'base_url': 'https://dict.example.com',
'api_keys': ['dict-key'],
}
runtime_provider = await model_mgr.load_provider(provider_info)
assert runtime_provider.provider_entity.uuid == 'dict-provider-uuid'
assert runtime_provider.provider_entity.name == 'Dict Provider'
assert runtime_provider.token_mgr.name == 'dict-provider-uuid'
assert runtime_provider.token_mgr.tokens == ['dict-key']
assert isinstance(runtime_provider.requester, requester.ProviderAPIRequester)
@pytest.mark.asyncio
async def test_model_manager_load_provider_from_entity(fake_requester_registry, fake_persistence_data):
"""Test ModelManager loads provider from persistence entity."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
provider_entity = fake_persistence_data['providers'][0]
runtime_provider = await model_mgr.load_provider(provider_entity)
assert runtime_provider.provider_entity.uuid == provider_entity.uuid
assert runtime_provider.requester is not None
# ============================================================================
# Model Query Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_get_model_by_uuid(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.get_model_by_uuid returns correct model."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'llm_models' in query_str:
return _make_mock_result(fake_persistence_data['llm_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
model = await model_mgr.get_model_by_uuid('test-llm-uuid-1')
assert model.model_entity.uuid == 'test-llm-uuid-1'
assert model.model_entity.name == 'TestLLM-1'
@pytest.mark.asyncio
async def test_model_manager_get_model_by_uuid_not_found(fake_requester_registry):
"""Test ModelManager.get_model_by_uuid raises ValueError for unknown model."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
with pytest.raises(ValueError) as exc_info:
await model_mgr.get_model_by_uuid('unknown-model-uuid')
assert 'unknown-model-uuid' in str(exc_info.value)
@pytest.mark.asyncio
async def test_model_manager_get_embedding_model_by_uuid(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.get_embedding_model_by_uuid returns correct model."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'embedding_models' in query_str:
return _make_mock_result(fake_persistence_data['embedding_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
model = await model_mgr.get_embedding_model_by_uuid('test-embedding-uuid-1')
assert model.model_entity.uuid == 'test-embedding-uuid-1'
@pytest.mark.asyncio
async def test_model_manager_get_embedding_model_by_uuid_not_found(fake_requester_registry):
"""Test ModelManager.get_embedding_model_by_uuid raises ValueError."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
with pytest.raises(ValueError):
await model_mgr.get_embedding_model_by_uuid('unknown-embedding-uuid')
@pytest.mark.asyncio
async def test_model_manager_get_rerank_model_by_uuid(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.get_rerank_model_by_uuid returns correct model."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'rerank_models' in query_str:
return _make_mock_result(fake_persistence_data['rerank_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
model = await model_mgr.get_rerank_model_by_uuid('test-rerank-uuid-1')
assert model.model_entity.uuid == 'test-rerank-uuid-1'
@pytest.mark.asyncio
async def test_model_manager_get_rerank_model_by_uuid_not_found(fake_requester_registry):
"""Test ModelManager.get_rerank_model_by_uuid raises ValueError."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
with pytest.raises(ValueError):
await model_mgr.get_rerank_model_by_uuid('unknown-rerank-uuid')
# ============================================================================
# Model Removal Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_remove_llm_model(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.remove_llm_model removes model correctly."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'llm_models' in query_str:
return _make_mock_result(fake_persistence_data['llm_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert len(model_mgr.llm_models) == 2
await model_mgr.remove_llm_model('test-llm-uuid-1')
assert len(model_mgr.llm_models) == 1
assert model_mgr.llm_models[0].model_entity.uuid == 'test-llm-uuid-2'
@pytest.mark.asyncio
async def test_model_manager_remove_llm_model_not_found(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.remove_llm_model handles unknown model gracefully."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'llm_models' in query_str:
return _make_mock_result(fake_persistence_data['llm_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
original_count = len(model_mgr.llm_models)
# Removing unknown model should do nothing (no error)
await model_mgr.remove_llm_model('unknown-model-uuid')
assert len(model_mgr.llm_models) == original_count
@pytest.mark.asyncio
async def test_model_manager_remove_embedding_model(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.remove_embedding_model removes model correctly."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'embedding_models' in query_str:
return _make_mock_result(fake_persistence_data['embedding_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert len(model_mgr.embedding_models) == 1
await model_mgr.remove_embedding_model('test-embedding-uuid-1')
assert len(model_mgr.embedding_models) == 0
@pytest.mark.asyncio
async def test_model_manager_remove_rerank_model(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.remove_rerank_model removes model correctly."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'rerank_models' in query_str:
return _make_mock_result(fake_persistence_data['rerank_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert len(model_mgr.rerank_models) == 1
await model_mgr.remove_rerank_model('test-rerank-uuid-1')
assert len(model_mgr.rerank_models) == 0
@pytest.mark.asyncio
async def test_model_manager_remove_provider(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.remove_provider removes provider correctly."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
return _make_mock_result(fake_persistence_data['providers'])
elif 'llm_models' in query_str:
return _make_mock_result(fake_persistence_data['llm_models'])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert fake_persistence_data['provider_uuid'] in model_mgr.provider_dict
await model_mgr.remove_provider(fake_persistence_data['provider_uuid'])
assert fake_persistence_data['provider_uuid'] not in model_mgr.provider_dict
# ============================================================================
# Requester Info Tests
# ============================================================================
def test_model_manager_get_available_requesters_info(fake_requester_registry):
"""Test ModelManager.get_available_requesters_info returns correct info."""
model_mgr = fake_requester_registry
model_mgr.requester_components = []
info = model_mgr.get_available_requesters_info('')
assert info == []
def test_model_manager_get_available_requesters_info_with_type_filter(fake_requester_registry):
"""Test ModelManager.get_available_requesters_info filters by model type."""
model_mgr = fake_requester_registry
from langbot.pkg.discover import engine as discover_engine
manifest = {
'apiVersion': 'v1',
'kind': 'LLMAPIRequester',
'metadata': {'name': 'test-req', 'label': {'en_US': 'Test'}, 'description': {'en_US': 'Test'}},
'spec': {'support_type': ['chat', 'embedding']},
'execution': {'python': {'path': 'fake', 'attr': 'FakeClass'}},
}
component = discover_engine.Component(owner='test', manifest=manifest, rel_path='fake.yaml')
model_mgr.requester_components = [component]
# Filter by chat type
info = model_mgr.get_available_requesters_info('chat')
assert len(info) == 1
assert info[0]['name'] == 'test-req'
# Filter by unsupported type
info = model_mgr.get_available_requesters_info('rerank')
assert len(info) == 0
def test_model_manager_get_available_requester_info_by_name(fake_requester_registry):
"""Test ModelManager.get_available_requester_info_by_name returns correct info."""
model_mgr = fake_requester_registry
from langbot.pkg.discover import engine as discover_engine
manifest = {
'apiVersion': 'v1',
'kind': 'LLMAPIRequester',
'metadata': {'name': 'named-req', 'label': {'en_US': 'Named'}, 'description': {'en_US': 'Named'}},
'spec': {'support_type': ['chat']},
'execution': {'python': {'path': 'fake', 'attr': 'FakeClass'}},
}
component = discover_engine.Component(owner='test', manifest=manifest, rel_path='fake.yaml')
model_mgr.requester_components = [component]
info = model_mgr.get_available_requester_info_by_name('named-req')
assert info is not None
assert info['name'] == 'named-req'
info = model_mgr.get_available_requester_info_by_name('unknown-req')
assert info is None
def test_model_manager_get_available_requester_manifest_by_name(fake_requester_registry):
"""Test ModelManager.get_available_requester_manifest_by_name returns component."""
model_mgr = fake_requester_registry
from langbot.pkg.discover import engine as discover_engine
manifest = {
'apiVersion': 'v1',
'kind': 'LLMAPIRequester',
'metadata': {'name': 'manifest-req', 'label': {'en_US': 'Manifest'}, 'description': {'en_US': 'Manifest'}},
'spec': {'support_type': ['chat']},
'execution': {'python': {'path': 'fake', 'attr': 'FakeClass'}},
}
component = discover_engine.Component(owner='test', manifest=manifest, rel_path='fake.yaml')
model_mgr.requester_components = [component]
comp = model_mgr.get_available_requester_manifest_by_name('manifest-req')
assert comp is not None
assert comp.metadata.name == 'manifest-req'
comp = model_mgr.get_available_requester_manifest_by_name('unknown-req')
assert comp is None
# ============================================================================
# Temporary Runtime Model Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_init_temporary_runtime_llm_model(fake_requester_registry):
"""Test ModelManager.init_temporary_runtime_llm_model creates model correctly."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
model_info = {
'uuid': 'temp-model-uuid',
'name': 'TempModel',
'provider': {
'uuid': 'temp-provider-uuid',
'name': 'Temp Provider',
'requester': 'fake-requester',
'base_url': 'https://temp.example.com',
'api_keys': ['temp-key'],
},
'abilities': ['func_call'],
'context_length': 128000,
'extra_args': {'temperature': 0.5},
}
runtime_model = await model_mgr.init_temporary_runtime_llm_model(model_info)
assert runtime_model.model_entity.uuid == 'temp-model-uuid'
assert runtime_model.model_entity.name == 'TempModel'
assert runtime_model.model_entity.context_length == 128000
assert runtime_model.model_entity.extra_args == {'temperature': 0.5}
assert 'context_length' not in runtime_model.model_entity.extra_args
assert runtime_model.provider.provider_entity.uuid == 'temp-provider-uuid'
assert runtime_model.provider.token_mgr.tokens == ['temp-key']
@pytest.mark.asyncio
async def test_model_manager_init_temporary_runtime_embedding_model(fake_requester_registry):
"""Test ModelManager.init_temporary_runtime_embedding_model creates model correctly."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
model_info = {
'uuid': 'temp-embedding-uuid',
'name': 'TempEmbedding',
'provider': {
'uuid': 'temp-provider-uuid',
'name': 'Temp Provider',
'requester': 'fake-requester',
'base_url': 'https://temp.example.com',
'api_keys': [],
},
'extra_args': {'dimensions': 512},
}
runtime_model = await model_mgr.init_temporary_runtime_embedding_model(model_info)
assert runtime_model.model_entity.uuid == 'temp-embedding-uuid'
assert runtime_model.model_entity.name == 'TempEmbedding'
@pytest.mark.asyncio
async def test_model_manager_init_temporary_runtime_rerank_model(fake_requester_registry):
"""Test ModelManager.init_temporary_runtime_rerank_model creates model correctly."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
model_info = {
'uuid': 'temp-rerank-uuid',
'name': 'TempRerank',
'provider': {
'uuid': 'temp-provider-uuid',
'name': 'Temp Provider',
'requester': 'fake-requester',
'base_url': 'https://temp.example.com',
'api_keys': [],
},
'extra_args': {},
}
runtime_model = await model_mgr.init_temporary_runtime_rerank_model(model_info)
assert runtime_model.model_entity.uuid == 'temp-rerank-uuid'
assert runtime_model.model_entity.name == 'TempRerank'
# ============================================================================
# Provider Reload Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_reload_provider(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.reload_provider reloads provider and updates model refs."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
# For initial load - return all providers
rows = [_make_row_mock(p) for p in fake_persistence_data['providers']]
return _make_mock_result(rows)
elif 'llm_models' in query_str:
rows = [_make_row_mock(m) for m in fake_persistence_data['llm_models']]
return _make_mock_result(rows)
elif 'embedding_models' in query_str:
rows = [_make_row_mock(m) for m in fake_persistence_data['embedding_models']]
return _make_mock_result(rows)
elif 'rerank_models' in query_str:
rows = [_make_row_mock(m) for m in fake_persistence_data['rerank_models']]
return _make_mock_result(rows)
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
original_provider = model_mgr.provider_dict[fake_persistence_data['provider_uuid']]
original_base_url = original_provider.provider_entity.base_url
# Setup for reload - return updated provider
async def reload_execute(query):
updated_provider = persistence_model.ModelProvider(
uuid=fake_persistence_data['provider_uuid'],
name='Updated Provider',
requester='fake-requester',
base_url='https://updated.example.com',
api_keys=['updated-key'],
)
return _make_mock_result([_make_row_mock(updated_provider)], first_item=_make_row_mock(updated_provider))
model_mgr.ap.persistence_mgr.execute_async = reload_execute
await model_mgr.reload_provider(fake_persistence_data['provider_uuid'])
updated_provider = model_mgr.provider_dict[fake_persistence_data['provider_uuid']]
assert updated_provider.provider_entity.base_url == 'https://updated.example.com'
assert updated_provider.provider_entity.base_url != original_base_url
@pytest.mark.asyncio
async def test_model_manager_reload_provider_not_found(fake_requester_registry):
"""Test ModelManager.reload_provider raises ProviderNotFoundError."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
async def fake_execute(query):
return _make_mock_result([], first_item=None)
model_mgr.ap.persistence_mgr.execute_async = fake_execute
with pytest.raises(provider_errors.ProviderNotFoundError) as exc_info:
await model_mgr.reload_provider('unknown-provider-uuid')
assert exc_info.value.provider_name == 'unknown-provider-uuid'
# ============================================================================
# Model Load with Provider Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_load_llm_model_with_provider(
fake_requester_registry, fake_persistence_data, runtime_provider
):
"""Test ModelManager.load_llm_model_with_provider creates RuntimeLLMModel."""
model_mgr = fake_requester_registry
model_entity = fake_persistence_data['llm_models'][0]
runtime_model = await model_mgr.load_llm_model_with_provider(model_entity, runtime_provider)
assert runtime_model.model_entity.uuid == model_entity.uuid
assert runtime_model.provider is runtime_provider
@pytest.mark.asyncio
async def test_model_manager_load_llm_model_with_provider_from_row(
fake_requester_registry, fake_persistence_data, runtime_provider
):
"""Test ModelManager.load_llm_model_with_provider handles Row objects."""
model_mgr = fake_requester_registry
model_entity = fake_persistence_data['llm_models'][0]
row_mock = _make_row_mock(model_entity)
runtime_model = await model_mgr.load_llm_model_with_provider(row_mock, runtime_provider)
assert runtime_model.model_entity.uuid == model_entity.uuid
@pytest.mark.asyncio
async def test_model_manager_load_embedding_model_with_provider(
fake_requester_registry, fake_persistence_data, runtime_provider
):
"""Test ModelManager.load_embedding_model_with_provider creates RuntimeEmbeddingModel."""
model_mgr = fake_requester_registry
model_entity = fake_persistence_data['embedding_models'][0]
runtime_model = await model_mgr.load_embedding_model_with_provider(model_entity, runtime_provider)
assert runtime_model.model_entity.uuid == model_entity.uuid
assert runtime_model.provider is runtime_provider
@pytest.mark.asyncio
async def test_model_manager_load_rerank_model_with_provider(fake_requester_registry, fake_persistence_data):
"""Test ModelManager.load_rerank_model_with_provider creates RuntimeRerankModel."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
provider_entity = fake_persistence_data['providers'][1]
token_mgr = token.TokenManager(name=provider_entity.uuid, tokens=provider_entity.api_keys or [])
requester_inst = model_mgr.requester_dict['another-fake-requester'](
ap=model_mgr.ap, config={'base_url': provider_entity.base_url}
)
await requester_inst.initialize()
provider = requester.RuntimeProvider(
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
)
model_entity = fake_persistence_data['rerank_models'][0]
runtime_model = await model_mgr.load_rerank_model_with_provider(model_entity, provider)
assert runtime_model.model_entity.uuid == model_entity.uuid
assert runtime_model.provider is provider
# ============================================================================
# Missing Provider Warning Tests
# ============================================================================
@pytest.mark.asyncio
async def test_model_manager_logs_warning_for_missing_provider(fake_requester_registry):
"""Test ModelManager logs warning when model's provider is missing."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
# Return empty providers
return _make_mock_result([])
elif 'llm_models' in query_str:
# Return model with missing provider
fake_model = persistence_model.LLMModel(
uuid='model-with-missing-provider',
name='MissingProviderModel',
provider_uuid='missing-provider-uuid',
abilities=[],
extra_args={},
)
return _make_mock_result([_make_row_mock(fake_model)])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
# Should have logged warning and skipped the model
assert len(model_mgr.llm_models) == 0
model_mgr.ap.logger.warning.assert_called()
@pytest.mark.asyncio
async def test_model_manager_handles_requester_not_found_gracefully(fake_requester_registry):
"""Test ModelManager handles RequesterNotFoundError during provider load."""
model_mgr = fake_requester_registry
async def fake_execute(query):
query_str = str(query)
if 'model_providers' in query_str:
# Return provider with unknown requester
fake_provider = persistence_model.ModelProvider(
uuid='provider-with-unknown-requester',
name='Unknown Requester Provider',
requester='unknown-requester-name',
base_url='https://unknown.com',
api_keys=[],
)
return _make_mock_result([_make_row_mock(fake_provider)])
elif 'llm_models' in query_str:
fake_model = persistence_model.LLMModel(
uuid='model-uuid',
name='Model',
provider_uuid='provider-with-unknown-requester',
abilities=[],
extra_args={},
)
return _make_mock_result([_make_row_mock(fake_model)])
return _make_mock_result([])
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
# Provider should be skipped
assert len(model_mgr.provider_dict) == 0
assert len(model_mgr.llm_models) == 0
model_mgr.ap.logger.warning.assert_called()
# ============================================================================
# Error Classes Tests
# ============================================================================
def test_requester_not_found_error_str():
"""Test RequesterNotFoundError string representation."""
error = provider_errors.RequesterNotFoundError('test-requester')
assert str(error) == 'Requester test-requester not found'
assert error.requester_name == 'test-requester'
def test_provider_not_found_error_str():
"""Test ProviderNotFoundError string representation."""
error = provider_errors.ProviderNotFoundError('test-provider')
assert str(error) == 'Provider test-provider not found'
assert error.provider_name == 'test-provider'
@@ -0,0 +1,227 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.api.http.service.model import _runtime_model_data
from langbot.pkg.api.http.service.provider import ModelProviderService
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.pipeline.preproc.preproc import PreProcessor
from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.modelmgr.token import TokenManager
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
def test_runtime_llm_model_data_preserves_uuid_after_update_payload_uuid_removed():
update_payload = {
'name': 'Qwen3.5-27B',
'provider_uuid': 'provider-uuid',
'abilities': [],
'extra_args': {},
}
runtime_entity = persistence_model.LLMModel(**_runtime_model_data('model-uuid', update_payload))
assert runtime_entity.uuid == 'model-uuid'
assert runtime_entity.name == 'Qwen3.5-27B'
def test_runtime_embedding_model_data_preserves_uuid_after_update_payload_uuid_removed():
update_payload = {
'name': 'embedding-model',
'provider_uuid': 'provider-uuid',
'extra_args': {},
}
runtime_entity = persistence_model.EmbeddingModel(**_runtime_model_data('embedding-uuid', update_payload))
assert runtime_entity.uuid == 'embedding-uuid'
assert runtime_entity.name == 'embedding-model'
def test_runtime_rerank_model_data_preserves_uuid_after_update_payload_uuid_removed():
update_payload = {
'name': 'rerank-model',
'provider_uuid': 'provider-uuid',
'extra_args': {},
}
runtime_entity = persistence_model.RerankModel(**_runtime_model_data('rerank-uuid', update_payload))
assert runtime_entity.uuid == 'rerank-uuid'
assert runtime_entity.name == 'rerank-model'
def test_normalize_space_provider_api_keys_filters_blank_values():
assert ModelProviderService._normalize_api_keys('space-key') == ['space-key']
assert ModelProviderService._normalize_api_keys(' trimmed-key ') == ['trimmed-key']
assert ModelProviderService._normalize_api_keys('') == []
assert ModelProviderService._normalize_api_keys(' ') == []
assert ModelProviderService._normalize_api_keys(None) == []
assert ModelProviderService._normalize_api_keys([' first-key ', '', 'first-key', 'second-key']) == [
'first-key',
'second-key',
]
def test_token_manager_filters_blank_and_duplicate_tokens():
token_mgr = TokenManager('provider-uuid', [' first-key ', '', 'first-key', 'second-key', ' '])
assert token_mgr.tokens == ['first-key', 'second-key']
assert token_mgr.get_token() == 'first-key'
def test_token_manager_next_token_ignores_empty_token_list():
token_mgr = TokenManager('provider-uuid', [])
token_mgr.next_token()
assert token_mgr.get_token() == ''
assert token_mgr.using_token_index == 0
@pytest.mark.asyncio
async def test_model_manager_initialize_skips_space_sync_after_timeout():
ap = SimpleNamespace()
ap.discover = SimpleNamespace(get_components_by_kind=Mock(return_value=[]))
ap.instance_config = SimpleNamespace(data={'space': {'models_sync_timeout': 0.01}})
ap.logger = Mock()
mgr = ModelManager(ap)
mgr.load_models_from_db = AsyncMock()
async def slow_sync():
await asyncio.sleep(1)
mgr.sync_new_models_from_space = AsyncMock(side_effect=slow_sync)
await mgr.initialize()
mgr.load_models_from_db.assert_awaited_once()
mgr.sync_new_models_from_space.assert_awaited_once()
ap.logger.warning.assert_any_call('LangBot Space model sync timed out after 0.01s, skipping startup sync.')
@pytest.mark.asyncio
async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline():
from langbot.pkg.api.http.service.model import LLMModelsService
model_uuid = 'qwen-model-uuid'
provider_uuid = 'ollama-provider-uuid'
ap = SimpleNamespace()
ap.logger = Mock()
ap.persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
ap.tool_mgr = SimpleNamespace(get_all_tools=AsyncMock(return_value=[]))
ap.skill_mgr = None # PreProcessor only uses skill_mgr for the local-agent skill-binding branch
ap.plugin_connector = SimpleNamespace(
emit_event=AsyncMock(return_value=SimpleNamespace(event=SimpleNamespace(default_prompt=[], prompt=[])))
)
ap.model_mgr = ModelManager(ap)
runtime_provider = Mock()
ap.model_mgr.provider_dict = {provider_uuid: runtime_provider}
ap.model_mgr.llm_models = [
requester.RuntimeLLMModel(
model_entity=persistence_model.LLMModel(
uuid=model_uuid,
name='old-qwen-name',
provider_uuid=provider_uuid,
abilities=[],
extra_args={},
),
provider=runtime_provider,
)
]
await LLMModelsService(ap).update_llm_model(
model_uuid,
{
'name': 'Qwen3.5-27B',
'provider_uuid': provider_uuid,
'abilities': [],
'extra_args': {},
},
)
runtime_model = await ap.model_mgr.get_model_by_uuid(model_uuid)
assert runtime_model.model_entity.uuid == model_uuid
assert runtime_model.model_entity.name == 'Qwen3.5-27B'
session = SimpleNamespace(
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
)
conversation = SimpleNamespace(
uuid='conversation-uuid',
create_time=None,
update_time=None,
prompt=SimpleNamespace(messages=[], copy=Mock(return_value=SimpleNamespace(messages=[]))),
messages=[],
)
ap.sess_mgr = SimpleNamespace(
get_session=AsyncMock(return_value=session),
get_conversation=AsyncMock(return_value=conversation),
)
message_chain = platform_message.MessageChain([platform_message.Plain(text='hello')])
sender = platform_entities.Friend(id=12345, nickname='Tester', remark=None)
message_event = platform_events.FriendMessage(
type='FriendMessage',
sender=sender,
message_chain=message_chain,
time=1710000000,
)
pipeline_config = {
'ai': {
'runner': {'runner': 'local-agent'},
'local-agent': {
'model': {'primary': model_uuid, 'fallbacks': []},
'prompt': [],
'knowledge-bases': [],
},
},
'trigger': {'misc': {'combine-quote-message': False}},
'output': {'misc': {'remove-think': False}},
}
query = pipeline_query.Query.model_construct(
query_id='query-id',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=12345,
message_chain=message_chain,
message_event=message_event,
adapter=AsyncMock(),
pipeline_uuid='pipeline-uuid',
bot_uuid='bot-uuid',
pipeline_config=pipeline_config,
session=None,
prompt=None,
messages=[],
user_message=None,
use_funcs=[],
use_llm_model_uuid=None,
variables={},
resp_messages=[],
resp_message_chain=None,
current_stage_name=None,
)
result = await PreProcessor(ap).process(query, 'PreProcessor')
processed_query = result.new_query
assert processed_query.use_llm_model_uuid == model_uuid
runner = SimpleNamespace(ap=ap, pipeline_config=pipeline_config)
candidates = await LocalAgentRunner._get_model_candidates(runner, processed_query)
assert [model.model_entity.uuid for model in candidates] == [model_uuid]
@@ -0,0 +1,172 @@
"""Unit tests for provider_specific_fields round-trip in LiteLLMRequester.
This tests the fix for GitHub issue #1899: Gemini requires thought_signature
to be preserved across tool call rounds for function calls to work correctly.
"""
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
def _make_requester() -> LiteLLMRequester:
# _convert_messages and _normalize_stream_tool_calls do not touch instance config.
return LiteLLMRequester.__new__(LiteLLMRequester)
def test_convert_messages_preserves_tool_call_provider_specific_fields():
"""Tool calls should retain provider_specific_fields through _convert_messages."""
req = _make_requester()
msg = provider_message.Message(
role='assistant',
content=None,
tool_calls=[
provider_message.ToolCall(
id='call_123',
type='function',
function=provider_message.FunctionCall(
name='search',
arguments='{"query": "test"}',
),
provider_specific_fields={
'thought_signature': 'c2tpcF90aG91Z2h0X3NpZ25hdHVyZQ==',
},
),
],
)
out = req._convert_messages([msg])
assert len(out) == 1
assert out[0]['tool_calls'] is not None
assert len(out[0]['tool_calls']) == 1
tc = out[0]['tool_calls'][0]
assert tc['id'] == 'call_123'
assert tc['function']['name'] == 'search'
assert 'provider_specific_fields' in tc
assert tc['provider_specific_fields']['thought_signature'] == 'c2tpcF90aG91Z2h0X3NpZ25hdHVyZQ=='
def test_convert_messages_preserves_message_provider_specific_fields():
"""Messages should retain provider_specific_fields through _convert_messages."""
req = _make_requester()
msg = provider_message.Message(
role='assistant',
content='Hello',
provider_specific_fields={
'thought_signatures': ['sig1', 'sig2'],
},
)
out = req._convert_messages([msg])
assert len(out) == 1
assert 'provider_specific_fields' in out[0]
assert out[0]['provider_specific_fields']['thought_signatures'] == ['sig1', 'sig2']
def test_normalize_stream_tool_calls_preserves_provider_specific_fields():
"""Streaming tool calls should retain provider_specific_fields."""
req = _make_requester()
tool_call_state: dict[int, dict] = {}
# Simulate first chunk with id and type
raw_tool_calls_1 = [
{
'index': 0,
'id': 'call_abc',
'type': 'function',
'function': {
'name': 'get_weather',
'arguments': '',
},
'provider_specific_fields': {
'thought_signature': 'dGVzdF9zaWduYXR1cmU=',
},
},
]
result_1 = req._normalize_stream_tool_calls(raw_tool_calls_1, tool_call_state)
assert result_1 is not None
assert len(result_1) == 1
assert result_1[0]['provider_specific_fields']['thought_signature'] == 'dGVzdF9zaWduYXR1cmU='
# Simulate second chunk without provider_specific_fields (should be retained from state)
raw_tool_calls_2 = [
{
'index': 0,
'function': {
'arguments': '{"city": "Tokyo"}',
},
},
]
result_2 = req._normalize_stream_tool_calls(raw_tool_calls_2, tool_call_state)
assert result_2 is not None
assert len(result_2) == 1
# Should retain the provider_specific_fields from the first chunk
assert result_2[0]['provider_specific_fields']['thought_signature'] == 'dGVzdF9zaWduYXR1cmU='
assert result_2[0]['function']['arguments'] == '{"city": "Tokyo"}'
def test_normalize_stream_tool_calls_merges_function_level_psf():
"""Function-level provider_specific_fields should be merged into tool-level."""
req = _make_requester()
tool_call_state: dict[int, dict] = {}
raw_tool_calls = [
{
'index': 0,
'id': 'call_xyz',
'type': 'function',
'function': {
'name': 'search',
'arguments': '{}',
'provider_specific_fields': {
'thought_signature': 'ZnVuY19sZXZlbF9zaWc=',
},
},
},
]
result = req._normalize_stream_tool_calls(raw_tool_calls, tool_call_state)
assert result is not None
assert result[0]['provider_specific_fields']['thought_signature'] == 'ZnVuY19sZXZlbF9zaWc='
def test_tool_call_roundtrip_through_message_entity():
"""Full round-trip: LiteLLM response dict -> Message entity -> _convert_messages."""
# Simulate what LiteLLM returns for a Gemini tool call response
message_data = {
'role': 'assistant',
'content': None,
'tool_calls': [
{
'id': 'call_gemini_123',
'type': 'function',
'function': {
'name': 'web_search',
'arguments': '{"query": "test"}',
},
'provider_specific_fields': {
'thought_signature': 'Z2VtaW5pX3NpZ25hdHVyZQ==',
},
},
],
'provider_specific_fields': {
'thought_signatures': ['Z2VtaW5pX3NpZ25hdHVyZQ=='],
},
}
# Parse into Message entity (this is what invoke_llm does)
msg = provider_message.Message(**message_data)
# Verify the entity has the fields
assert msg.tool_calls is not None
assert len(msg.tool_calls) == 1
assert msg.tool_calls[0].provider_specific_fields is not None
assert msg.tool_calls[0].provider_specific_fields['thought_signature'] == 'Z2VtaW5pX3NpZ25hdHVyZQ=='
assert msg.provider_specific_fields is not None
assert msg.provider_specific_fields['thought_signatures'] == ['Z2VtaW5pX3NpZ25hdHVyZQ==']
# Convert back to dict for LiteLLM (this is what _convert_messages does)
req = _make_requester()
out = req._convert_messages([msg])
# Verify the fields are preserved in the output
assert out[0]['tool_calls'][0]['provider_specific_fields']['thought_signature'] == 'Z2VtaW5pX3NpZ25hdHVyZQ=='
assert out[0]['provider_specific_fields']['thought_signatures'] == ['Z2VtaW5pX3NpZ25hdHVyZQ==']
@@ -0,0 +1,640 @@
"""
Unit tests for ProviderAPIRequester base class and runtime entities in provider/modelmgr.
Tests requester initialization, configuration handling, token management,
and runtime model/provider behavior without calling real LLM APIs.
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.provider.modelmgr import token
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.provider.modelmgr.errors import RequesterError
# ============================================================================
# ProviderAPIRequester Base Class Tests
# ============================================================================
class TestableRequester(requester.ProviderAPIRequester):
"""Testable requester subclass for testing base class behavior."""
name = 'testable-requester'
default_config = {
'base_url': 'https://default.example.com',
'timeout': 60,
'max_retries': 3,
}
async def invoke_llm(
self,
query,
model: requester.RuntimeLLMModel,
messages: list,
funcs=None,
extra_args={},
remove_think=False,
):
import langbot_plugin.api.entities.builtin.provider.message as provider_message
return provider_message.Message(
role='assistant',
content=[provider_message.ContentElement(type='text', text='Testable response')],
)
def test_requester_base_class_is_abstract():
"""Test ProviderAPIRequester cannot be instantiated directly."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
# ProviderAPIRequester has abstract methods, but ABCMeta allows instantiation
# if you don't call the abstract methods. Test that it has abstract methods.
assert hasattr(requester.ProviderAPIRequester, 'invoke_llm')
# Check that invoke_llm is abstract
assert hasattr(requester.ProviderAPIRequester.invoke_llm, '__isabstractmethod__')
def test_requester_default_config_merged():
"""Test requester merges default config with provided config."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {'base_url': 'https://custom.example.com', 'custom_key': 'custom_value'})
assert inst.requester_cfg['base_url'] == 'https://custom.example.com'
assert inst.requester_cfg['timeout'] == 60 # from default
assert inst.requester_cfg['max_retries'] == 3 # from default
assert inst.requester_cfg['custom_key'] == 'custom_value' # custom added
def test_requester_default_config_not_modified():
"""Test that default_config dict is not modified when merging."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {'base_url': 'https://override.example.com'})
assert TestableRequester.default_config['base_url'] == 'https://default.example.com'
assert inst.requester_cfg['base_url'] == 'https://override.example.com'
def test_requester_empty_config_uses_defaults():
"""Test requester uses defaults when empty config provided."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {})
assert inst.requester_cfg == inst.default_config
@pytest.mark.asyncio
async def test_requester_initialize_is_callable():
"""Test requester initialize method is callable (default is pass)."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {})
await inst.initialize()
# No exception should occur
@pytest.mark.asyncio
async def test_requester_scan_models_not_implemented():
"""Test scan_models raises NotImplementedError by default."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {})
await inst.initialize()
with pytest.raises(NotImplementedError) as exc_info:
await inst.scan_models()
assert 'does not support model scanning' in str(exc_info.value)
@pytest.mark.asyncio
async def test_requester_invoke_rerank_not_implemented():
"""Test invoke_rerank raises NotImplementedError by default."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {})
await inst.initialize()
# Create fake model
fake_provider_entity = persistence_model.ModelProvider(
uuid='provider-uuid',
name='Provider',
requester='test',
base_url='https://test.com',
api_keys=[],
)
fake_token_mgr = token.TokenManager(name='test', tokens=[])
fake_requester = inst
fake_provider = requester.RuntimeProvider(
provider_entity=fake_provider_entity,
token_mgr=fake_token_mgr,
requester=fake_requester,
)
fake_model_entity = persistence_model.RerankModel(
uuid='model-uuid',
name='Model',
provider_uuid='provider-uuid',
extra_args={},
)
fake_model = requester.RuntimeRerankModel(
model_entity=fake_model_entity,
provider=fake_provider,
)
with pytest.raises(NotImplementedError) as exc_info:
await inst.invoke_rerank(fake_model, 'query', ['doc1', 'doc2'])
assert 'does not support rerank' in str(exc_info.value)
# ============================================================================
# TokenManager Tests
# ============================================================================
def test_token_manager_initial_state():
"""Test TokenManager initial state."""
mgr = token.TokenManager(name='test-manager', tokens=['key1', 'key2', 'key3'])
assert mgr.name == 'test-manager'
assert mgr.tokens == ['key1', 'key2', 'key3']
assert mgr.using_token_index == 0
def test_token_manager_get_token():
"""Test TokenManager.get_token returns current token."""
mgr = token.TokenManager(name='test', tokens=['key1', 'key2'])
assert mgr.get_token() == 'key1'
def test_token_manager_get_token_empty():
"""Test TokenManager.get_token returns empty string when no tokens."""
mgr = token.TokenManager(name='test', tokens=[])
assert mgr.get_token() == ''
def test_token_manager_next_token_cycles():
"""Test TokenManager.next_token cycles through tokens."""
mgr = token.TokenManager(name='test', tokens=['key1', 'key2', 'key3'])
assert mgr.get_token() == 'key1'
mgr.next_token()
assert mgr.get_token() == 'key2'
mgr.next_token()
assert mgr.get_token() == 'key3'
# Should cycle back to first
mgr.next_token()
assert mgr.get_token() == 'key1'
def test_token_manager_next_token_single():
"""Test TokenManager.next_token with single token."""
mgr = token.TokenManager(name='test', tokens=['single-key'])
mgr.next_token()
assert mgr.get_token() == 'single-key'
mgr.next_token()
assert mgr.get_token() == 'single-key'
def test_token_manager_next_token_empty():
"""Test TokenManager.next_token with empty tokens doesn't error."""
mgr = token.TokenManager(name='test', tokens=[])
assert mgr.next_token() is None
assert mgr.get_token() == ''
# ============================================================================
# RuntimeProvider Tests
# ============================================================================
def test_runtime_provider_initialization(runtime_provider, fake_persistence_data):
"""Test RuntimeProvider initialization."""
provider = runtime_provider
provider_entity = fake_persistence_data['providers'][0]
assert provider.provider_entity.uuid == provider_entity.uuid
assert provider.provider_entity.name == provider_entity.name
assert provider.token_mgr.name == provider_entity.uuid
assert provider.token_mgr.tokens == provider_entity.api_keys
assert isinstance(provider.requester, requester.ProviderAPIRequester)
def test_runtime_provider_has_invoke_methods(runtime_provider):
"""Test RuntimeProvider has invoke methods that delegate to requester."""
provider = runtime_provider
assert hasattr(provider, 'invoke_llm')
assert hasattr(provider, 'invoke_llm_stream')
assert hasattr(provider, 'invoke_embedding')
assert hasattr(provider, 'invoke_rerank')
@pytest.mark.asyncio
async def test_runtime_provider_invoke_llm_delegates(runtime_provider, runtime_llm_model):
"""Test RuntimeProvider.invoke_llm delegates to requester."""
provider = runtime_provider
# Track that requester was called
provider.requester._invoke_count = 0
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
# Create minimal query for testing (bypass validation)
query = pipeline_query.Query.model_construct(
query_id='test-query',
launcher_type='person',
launcher_id=12345,
sender_id=12345,
message_chain=None,
message_event=None,
adapter=None,
pipeline_uuid='pipeline-uuid',
bot_uuid='bot-uuid',
pipeline_config={'ai': {}, 'output': {}, 'trigger': {}},
session=None,
prompt=None,
messages=[],
user_message=None,
use_funcs=[],
use_llm_model_uuid=None,
variables={},
resp_messages=[],
resp_message_chain=None,
current_stage_name=None,
)
messages = [
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
]
result = await provider.invoke_llm(query, runtime_llm_model, messages)
assert provider.requester._invoke_count == 1
assert provider.requester._last_messages == messages
assert provider.requester._last_model == runtime_llm_model
assert result.role == 'assistant'
@pytest.mark.asyncio
async def test_runtime_provider_invoke_llm_stream_yields_chunks(runtime_provider, runtime_llm_model):
"""Test RuntimeProvider.invoke_llm_stream yields chunks from requester."""
provider = runtime_provider
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
query = pipeline_query.Query.model_construct(
query_id='test-stream',
launcher_type='person',
launcher_id=12345,
sender_id=12345,
message_chain=None,
message_event=None,
adapter=None,
pipeline_uuid='pipeline-uuid',
bot_uuid='bot-uuid',
pipeline_config={'ai': {}, 'output': {}, 'trigger': {}},
session=None,
prompt=None,
messages=[],
user_message=None,
use_funcs=[],
use_llm_model_uuid=None,
variables={},
resp_messages=[],
resp_message_chain=None,
current_stage_name=None,
)
messages = [
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
]
chunks = []
async for chunk in provider.invoke_llm_stream(query, runtime_llm_model, messages):
chunks.append(chunk)
assert len(chunks) == 1
assert chunks[0].role == 'assistant'
@pytest.mark.asyncio
async def test_runtime_provider_invoke_embedding_returns_vectors(runtime_provider, runtime_embedding_model):
"""Test RuntimeProvider.invoke_embedding returns embedding vectors."""
provider = runtime_provider
result = await provider.invoke_embedding(runtime_embedding_model, ['text1', 'text2'])
assert len(result) == 2
assert result[0] == [0.1, 0.2, 0.3]
@pytest.mark.asyncio
async def test_runtime_provider_invoke_rerank_returns_scores(runtime_provider, runtime_rerank_model):
"""Test RuntimeProvider.invoke_rerank returns relevance scores."""
# Need to use the correct provider for rerank model
provider = runtime_rerank_model.provider
result = await provider.invoke_rerank(runtime_rerank_model, 'query', ['doc1', 'doc2', 'doc3'])
assert len(result) == 3
assert result[0]['index'] == 0
assert result[0]['relevance_score'] == 0.9
# ============================================================================
# RuntimeLLMModel Tests
# ============================================================================
def test_runtime_llm_model_initialization(runtime_llm_model, fake_persistence_data):
"""Test RuntimeLLMModel initialization."""
model = runtime_llm_model
model_entity = fake_persistence_data['llm_models'][0]
assert model.model_entity.uuid == model_entity.uuid
assert model.model_entity.name == model_entity.name
assert model.model_entity.abilities == model_entity.abilities
assert model.model_entity.extra_args == model_entity.extra_args
assert model.provider is not None
def test_runtime_llm_model_provider_ref(runtime_llm_model):
"""Test RuntimeLLMModel has correct provider reference."""
model = runtime_llm_model
assert model.provider.provider_entity is not None
assert model.provider.token_mgr is not None
assert model.provider.requester is not None
# ============================================================================
# RuntimeEmbeddingModel Tests
# ============================================================================
def test_runtime_embedding_model_initialization(runtime_embedding_model, fake_persistence_data):
"""Test RuntimeEmbeddingModel initialization."""
model = runtime_embedding_model
model_entity = fake_persistence_data['embedding_models'][0]
assert model.model_entity.uuid == model_entity.uuid
assert model.model_entity.name == model_entity.name
assert model.model_entity.extra_args == model_entity.extra_args
assert model.provider is not None
# ============================================================================
# RuntimeRerankModel Tests
# ============================================================================
def test_runtime_rerank_model_initialization(runtime_rerank_model, fake_persistence_data):
"""Test RuntimeRerankModel initialization."""
model = runtime_rerank_model
model_entity = fake_persistence_data['rerank_models'][0]
assert model.model_entity.uuid == model_entity.uuid
assert model.model_entity.name == model_entity.name
assert model.model_entity.extra_args == model_entity.extra_args
assert model.provider is not None
# ============================================================================
# RequesterError Tests
# ============================================================================
def test_requester_error_message_format():
"""Test RequesterError message format."""
error = RequesterError('API returned 500')
assert '模型请求失败' in str(error)
assert 'API returned 500' in str(error)
def test_requester_error_is_exception():
"""Test RequesterError is Exception subclass."""
error = RequesterError('test')
assert isinstance(error, Exception)
# ============================================================================
# ProviderAPIRequester Config Validation Tests
# ============================================================================
def test_requester_with_missing_base_url():
"""Test requester handles missing base_url in config."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
# If base_url is in default_config, it will be used
inst = TestableRequester(mock_app, {'timeout': 30})
assert inst.requester_cfg['base_url'] == 'https://default.example.com'
def test_requester_with_none_values():
"""Test requester handles None values in config."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = TestableRequester(mock_app, {'timeout': None, 'base_url': 'https://test.com'})
# None values are kept in the merged config
assert inst.requester_cfg['timeout'] is None
class RequesterWithNoDefaults(requester.ProviderAPIRequester):
"""Requester with empty defaults for testing."""
name = 'no-defaults-requester'
default_config = {}
async def invoke_llm(self, query, model, messages, funcs=None, extra_args={}, remove_think=False):
pass
def test_requester_empty_defaults_with_empty_config():
"""Test requester with empty defaults and empty config."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = RequesterWithNoDefaults(mock_app, {})
assert inst.requester_cfg == {}
def test_requester_empty_defaults_with_values():
"""Test requester with empty defaults receives config values."""
mock_app = SimpleNamespace()
mock_app.logger = Mock()
inst = RequesterWithNoDefaults(mock_app, {'base_url': 'https://custom.com', 'api_key': 'key'})
assert inst.requester_cfg['base_url'] == 'https://custom.com'
assert inst.requester_cfg['api_key'] == 'key'
# ============================================================================
# RuntimeProvider Error Handling Tests
# ============================================================================
class ErrorThrowingRequester(requester.ProviderAPIRequester):
"""Requester that throws errors for testing."""
name = 'error-requester'
default_config = {}
async def invoke_llm(self, query, model, messages, funcs=None, extra_args={}, remove_think=False):
raise RequesterError('Simulated API error')
@pytest.mark.asyncio
async def test_runtime_provider_invoke_llm_propagates_error(mock_app_for_modelmgr):
"""Test RuntimeProvider.invoke_llm propagates requester errors."""
mock_app = mock_app_for_modelmgr
# Add monitoring_service for error handling path
mock_app.monitoring_service = AsyncMock()
requester_inst = ErrorThrowingRequester(mock_app, {})
await requester_inst.initialize()
provider_entity = persistence_model.ModelProvider(
uuid='error-provider',
name='Error Provider',
requester='error-requester',
base_url='https://error.com',
api_keys=['error-key'],
)
token_mgr = token.TokenManager(name='error-provider', tokens=['error-key'])
provider = requester.RuntimeProvider(
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
)
model_entity = persistence_model.LLMModel(
uuid='error-model',
name='Error Model',
provider_uuid='error-provider',
abilities=[],
extra_args={},
)
model = requester.RuntimeLLMModel(model_entity=model_entity, provider=provider)
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
query = pipeline_query.Query.model_construct(
query_id='error-query',
launcher_type='person',
launcher_id=12345,
sender_id=12345,
message_chain=None,
message_event=None,
adapter=None,
pipeline_uuid='pipeline-uuid',
bot_uuid='bot-uuid',
pipeline_config={'ai': {}, 'output': {}, 'trigger': {}},
session=None,
prompt=None,
messages=[],
user_message=None,
use_funcs=[],
use_llm_model_uuid=None,
variables={},
resp_messages=[],
resp_message_chain=None,
current_stage_name=None,
)
messages = [
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
]
with pytest.raises(RequesterError):
await provider.invoke_llm(query, model, messages)
# ============================================================================
# LLMModelInfo Tests (from entities.py)
# ============================================================================
def test_llm_model_info_basic():
"""Test LLMModelInfo basic structure."""
from langbot.pkg.provider.modelmgr.entities import LLMModelInfo
mock_app = SimpleNamespace()
mock_app.logger = Mock()
fake_requester = TestableRequester(mock_app, {})
fake_token_mgr = token.TokenManager(name='test', tokens=['key'])
info = LLMModelInfo(
name='test-model',
model_name='gpt-4',
token_mgr=fake_token_mgr,
requester=fake_requester,
tool_call_supported=True,
vision_supported=False,
)
assert info.name == 'test-model'
assert info.model_name == 'gpt-4'
assert info.tool_call_supported == True
assert info.vision_supported == False
def test_llm_model_info_optional_fields():
"""Test LLMModelInfo optional fields default values."""
from langbot.pkg.provider.modelmgr.entities import LLMModelInfo
mock_app = SimpleNamespace()
mock_app.logger = Mock()
fake_requester = TestableRequester(mock_app, {})
fake_token_mgr = token.TokenManager(name='test', tokens=['key'])
info = LLMModelInfo(
name='minimal-model',
token_mgr=fake_token_mgr,
requester=fake_requester,
)
assert info.model_name is None
assert info.tool_call_supported == False # default
assert info.vision_supported == False # default

Some files were not shown because too many files have changed in this diff Show More