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
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:
@@ -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 环境或集成测试,不属于纯单元测试范畴。
|
||||
@@ -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',
|
||||
}
|
||||
@@ -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}
|
||||
@@ -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
@@ -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,
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
# Unit tests for command module
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=[])
|
||||
@@ -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
|
||||
@@ -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'
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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']
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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')
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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}
|
||||
@@ -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')
|
||||
@@ -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
|
||||
@@ -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 == ['✅ comment:looks good', '✅ choice:B']
|
||||
|
||||
|
||||
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✅ xiala:1'
|
||||
|
||||
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') == ['✅ comment:ready']
|
||||
|
||||
|
||||
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:你好',
|
||||
'✅ xiala:or',
|
||||
'✅ files:2 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 == ['✅ xiala:B']
|
||||
@@ -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✅ choice:B'
|
||||
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()
|
||||
@@ -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', [])
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
"""Provider requester tests"""
|
||||
@@ -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
Reference in New Issue
Block a user