mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 20:06:06 +00:00
feat(tenancy): implement workspace isolation
This commit is contained in:
Junyan Qin
@@ -1,389 +1,565 @@
|
||||
"""Unit tests for API knowledge service.
|
||||
|
||||
Tests cover:
|
||||
- Knowledge base CRUD operations
|
||||
- Capability checking
|
||||
- Knowledge engine discovery
|
||||
- File operations
|
||||
"""
|
||||
"""Tests for the tenant-aware knowledge service facade."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from importlib import import_module
|
||||
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.api.http.service.knowledge import KnowledgeService
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
|
||||
|
||||
def get_knowledge_service_module():
|
||||
"""Lazy import to avoid circular import issues."""
|
||||
return import_module('langbot.pkg.api.http.service.knowledge')
|
||||
CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=2,
|
||||
)
|
||||
|
||||
|
||||
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 _Rows:
|
||||
def __init__(self, rows=()):
|
||||
self.rows = list(rows)
|
||||
|
||||
def all(self):
|
||||
return self.rows
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.rows)
|
||||
|
||||
|
||||
def _app():
|
||||
return SimpleNamespace(
|
||||
logger=Mock(),
|
||||
rag_mgr=SimpleNamespace(
|
||||
get_all_knowledge_base_details=AsyncMock(return_value=[]),
|
||||
get_knowledge_base_details=AsyncMock(return_value=None),
|
||||
create_knowledge_base=AsyncMock(),
|
||||
remove_knowledge_base_from_runtime=AsyncMock(),
|
||||
load_knowledge_base=AsyncMock(),
|
||||
get_knowledge_base_by_uuid=AsyncMock(return_value=None),
|
||||
delete_knowledge_base=AsyncMock(),
|
||||
),
|
||||
persistence_mgr=SimpleNamespace(
|
||||
execute_async=AsyncMock(return_value=_Rows()),
|
||||
serialize_model=Mock(return_value={}),
|
||||
),
|
||||
plugin_connector=SimpleNamespace(
|
||||
is_enable_plugin=True,
|
||||
require_workspace_context=AsyncMock(side_effect=lambda context: context),
|
||||
get_rag_creation_schema=AsyncMock(return_value={}),
|
||||
get_rag_retrieval_schema=AsyncMock(return_value={}),
|
||||
list_knowledge_engines=AsyncMock(return_value=[]),
|
||||
list_parsers=AsyncMock(return_value=[]),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_and_get_forward_explicit_context():
|
||||
app = _app()
|
||||
app.rag_mgr.get_all_knowledge_base_details.return_value = [{'uuid': 'kb-a'}]
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb-a'}
|
||||
service = KnowledgeService(app)
|
||||
|
||||
assert await service.get_knowledge_bases(CONTEXT) == [{'uuid': 'kb-a'}]
|
||||
assert await service.get_knowledge_base(CONTEXT, 'kb-a') == {'uuid': 'kb-a'}
|
||||
app.rag_mgr.get_all_knowledge_base_details.assert_awaited_once_with(CONTEXT)
|
||||
app.rag_mgr.get_knowledge_base_details.assert_awaited_once_with(CONTEXT, 'kb-a')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_context_fails_closed_before_plugin_or_manager_access():
|
||||
app = _app()
|
||||
service = KnowledgeService(app)
|
||||
|
||||
with pytest.raises(WorkspaceRequiredError):
|
||||
await service.get_knowledge_bases(None)
|
||||
with pytest.raises(WorkspaceRequiredError):
|
||||
await service.create_knowledge_base(None, {'knowledge_engine_plugin_id': 'author/engine'})
|
||||
app.plugin_connector.get_rag_creation_schema.assert_not_awaited()
|
||||
app.rag_mgr.create_knowledge_base.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_validates_schema_and_binds_context():
|
||||
app = _app()
|
||||
app.plugin_connector.get_rag_creation_schema.return_value = {
|
||||
'schema': [{'name': 'endpoint', 'label': {'en_US': 'Endpoint'}, 'required': True}]
|
||||
}
|
||||
app.rag_mgr.create_knowledge_base.return_value = SimpleNamespace(uuid='kb-created')
|
||||
service = KnowledgeService(app)
|
||||
|
||||
with pytest.raises(ValueError, match='Endpoint is required'):
|
||||
await service.create_knowledge_base(
|
||||
CONTEXT,
|
||||
{'knowledge_engine_plugin_id': 'author/engine'},
|
||||
)
|
||||
|
||||
result = await service.create_knowledge_base(
|
||||
CONTEXT,
|
||||
{
|
||||
'name': 'KB',
|
||||
'description': 'desc',
|
||||
'knowledge_engine_plugin_id': 'author/engine',
|
||||
'creation_settings': {'endpoint': 'https://example.invalid'},
|
||||
},
|
||||
)
|
||||
assert result == 'kb-created'
|
||||
app.rag_mgr.create_knowledge_base.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
name='KB',
|
||||
knowledge_engine_plugin_id='author/engine',
|
||||
creation_settings={'endpoint': 'https://example.invalid'},
|
||||
retrieval_settings={},
|
||||
description='desc',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_rejects_guessed_uuid_and_scopes_reload():
|
||||
app = _app()
|
||||
service = KnowledgeService(app)
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.update_knowledge_base(CONTEXT, 'kb-other', {'name': 'stolen'})
|
||||
app.persistence_mgr.execute_async.assert_not_awaited()
|
||||
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb-a', 'workspace_uuid': 'workspace-a'}
|
||||
await service.update_knowledge_base(CONTEXT, 'kb-a', {'name': 'updated', 'uuid': 'ignored'})
|
||||
app.rag_mgr.remove_knowledge_base_from_runtime.assert_awaited_once_with(CONTEXT, 'kb-a')
|
||||
app.rag_mgr.load_knowledge_base.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
{'uuid': 'kb-a', 'workspace_uuid': 'workspace-a'},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_retrieve_uses_execution_context():
|
||||
app = _app()
|
||||
entry = SimpleNamespace(model_dump=Mock(return_value={'id': 'entry-a'}))
|
||||
runtime_kb = SimpleNamespace(retrieve=AsyncMock(return_value=[entry]))
|
||||
app.rag_mgr.get_knowledge_base_by_uuid.return_value = runtime_kb
|
||||
service = KnowledgeService(app)
|
||||
|
||||
assert await service.retrieve_knowledge_base(CONTEXT, 'kb-a', 'query', {'top_k': 3}) == [{'id': 'entry-a'}]
|
||||
runtime_kb.retrieve.assert_awaited_once_with(CONTEXT, 'query', settings={'top_k': 3})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_retrieve_cross_workspace_uuid_is_not_found():
|
||||
app = _app()
|
||||
service = KnowledgeService(app)
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.retrieve_knowledge_base(CONTEXT, 'kb-other', 'query')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_listing_checks_parent_knowledge_base_first():
|
||||
app = _app()
|
||||
service = KnowledgeService(app)
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.get_files_by_knowledge_base(CONTEXT, 'kb-other')
|
||||
app.persistence_mgr.execute_async.assert_not_awaited()
|
||||
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb-a'}
|
||||
row = SimpleNamespace(uuid='file-a')
|
||||
app.persistence_mgr.execute_async.return_value = _Rows([row])
|
||||
app.persistence_mgr.serialize_model.return_value = {'uuid': 'file-a'}
|
||||
assert await service.get_files_by_knowledge_base(CONTEXT, 'kb-a') == [{'uuid': 'file-a'}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_and_delete_file_require_runtime_parent_and_capability():
|
||||
app = _app()
|
||||
runtime_kb = SimpleNamespace(
|
||||
store_file=AsyncMock(return_value='task-a'),
|
||||
delete_file=AsyncMock(),
|
||||
)
|
||||
app.rag_mgr.get_knowledge_base_by_uuid.return_value = runtime_kb
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'knowledge_engine': {'capabilities': ['doc_ingestion']}}
|
||||
service = KnowledgeService(app)
|
||||
|
||||
assert await service.store_file(CONTEXT, 'kb-a', 'upload.pdf', 'author/parser') == 'task-a'
|
||||
runtime_kb.store_file.assert_awaited_once_with(CONTEXT, 'upload.pdf', parser_plugin_id='author/parser')
|
||||
await service.delete_file(CONTEXT, 'kb-a', 'file-a')
|
||||
runtime_kb.delete_file.assert_awaited_once_with(CONTEXT, 'file-a')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_knowledge_base_rejects_cross_workspace_uuid():
|
||||
app = _app()
|
||||
service = KnowledgeService(app)
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.delete_knowledge_base(CONTEXT, 'kb-other')
|
||||
app.rag_mgr.delete_knowledge_base.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_and_parser_discovery_require_context_and_filter_results():
|
||||
app = _app()
|
||||
app.plugin_connector.list_knowledge_engines.return_value = [{'plugin_id': 'author/engine'}]
|
||||
app.plugin_connector.list_parsers.return_value = [
|
||||
{'id': 'text', 'supported_mime_types': ['text/plain']},
|
||||
{'id': 'pdf', 'supported_mime_types': ['application/pdf']},
|
||||
]
|
||||
service = KnowledgeService(app)
|
||||
|
||||
assert await service.list_knowledge_engines(CONTEXT) == [{'plugin_id': 'author/engine'}]
|
||||
assert await service.list_parsers(CONTEXT, 'application/pdf') == [
|
||||
{'id': 'pdf', 'supported_mime_types': ['application/pdf']}
|
||||
]
|
||||
with pytest.raises(WorkspaceRequiredError):
|
||||
await service.list_parsers(None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_discovery_rejects_connector_workspace_or_generation_mismatch():
|
||||
app = _app()
|
||||
app.plugin_connector.require_workspace_context.side_effect = WorkspaceNotFoundError('Plugin resource not found')
|
||||
service = KnowledgeService(app)
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
|
||||
await service.list_knowledge_engines(CONTEXT)
|
||||
|
||||
app.plugin_connector.list_knowledge_engines.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schema_validation_refences_before_second_runtime_call():
|
||||
app = _app()
|
||||
app.plugin_connector.require_workspace_context.side_effect = [
|
||||
CONTEXT,
|
||||
WorkspaceNotFoundError('Plugin resource not found'),
|
||||
]
|
||||
service = KnowledgeService(app)
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
|
||||
await service.create_knowledge_base(
|
||||
CONTEXT,
|
||||
{'knowledge_engine_plugin_id': 'author/engine'},
|
||||
)
|
||||
|
||||
app.plugin_connector.get_rag_creation_schema.assert_awaited_once()
|
||||
app.plugin_connector.get_rag_retrieval_schema.assert_not_awaited()
|
||||
app.rag_mgr.create_knowledge_base.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_schemas_are_context_gated_and_fail_soft_on_connector_error():
|
||||
app = _app()
|
||||
app.plugin_connector.get_rag_creation_schema.return_value = {'schema': ['creation']}
|
||||
app.plugin_connector.get_rag_retrieval_schema.side_effect = RuntimeError('offline')
|
||||
service = KnowledgeService(app)
|
||||
|
||||
assert await service.get_engine_creation_schema(CONTEXT, 'author/engine') == {'schema': ['creation']}
|
||||
assert await service.get_engine_retrieval_schema(CONTEXT, 'author/engine') == {}
|
||||
with pytest.raises(WorkspaceRequiredError):
|
||||
await service.get_engine_creation_schema(None, 'author/engine')
|
||||
|
||||
|
||||
# Preserve the original service regression matrix with the new explicit
|
||||
# Workspace context. These intentionally overlap a few isolation-focused
|
||||
# tests above so legacy business behavior cannot disappear behind new guards.
|
||||
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()
|
||||
app = _app()
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
service = KnowledgeService(app)
|
||||
|
||||
assert service.ap is mock_app
|
||||
assert service.ap is 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'}])
|
||||
app = _app()
|
||||
app.rag_mgr.get_all_knowledge_base_details.return_value = [{'uuid': 'kb1', 'name': 'KB1'}]
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_knowledge_bases()
|
||||
result = await KnowledgeService(app).get_knowledge_bases(CONTEXT)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]['uuid'] == 'kb1'
|
||||
assert result == [{'uuid': 'kb1', 'name': 'KB1'}]
|
||||
app.rag_mgr.get_all_knowledge_base_details.assert_awaited_once_with(CONTEXT)
|
||||
|
||||
@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=[])
|
||||
app = _app()
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_knowledge_bases()
|
||||
|
||||
assert result == []
|
||||
assert await KnowledgeService(app).get_knowledge_bases(CONTEXT) == []
|
||||
|
||||
|
||||
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'})
|
||||
app = _app()
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb1', 'name': 'KB1'}
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_knowledge_base('kb1')
|
||||
result = await KnowledgeService(app).get_knowledge_base(CONTEXT, 'kb1')
|
||||
|
||||
assert result['uuid'] == 'kb1'
|
||||
assert result == {'uuid': 'kb1', 'name': '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)
|
||||
app = _app()
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_knowledge_base('nonexistent')
|
||||
|
||||
assert result is None
|
||||
assert await KnowledgeService(app).get_knowledge_base(CONTEXT, 'nonexistent') 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)
|
||||
app = _app()
|
||||
app.rag_mgr.create_knowledge_base.return_value = SimpleNamespace(uuid='new_kb_uuid')
|
||||
service = KnowledgeService(app)
|
||||
kb_data = {
|
||||
'name': 'Test KB',
|
||||
'knowledge_engine_plugin_id': 'author/engine',
|
||||
'description': 'Test description',
|
||||
}
|
||||
|
||||
result = await service.create_knowledge_base(kb_data)
|
||||
result = await service.create_knowledge_base(CONTEXT, kb_data)
|
||||
|
||||
assert result == 'new_kb_uuid'
|
||||
mock_app.rag_mgr.create_knowledge_base.assert_called_once()
|
||||
app.rag_mgr.create_knowledge_base.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
name='Test KB',
|
||||
knowledge_engine_plugin_id='author/engine',
|
||||
creation_settings={},
|
||||
retrieval_settings={},
|
||||
description='Test description',
|
||||
)
|
||||
|
||||
@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()
|
||||
app = _app()
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
with pytest.raises(ValueError, match='knowledge_engine_plugin_id is required'):
|
||||
await KnowledgeService(app).create_knowledge_base(CONTEXT, {'name': 'Test'})
|
||||
|
||||
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)
|
||||
app.rag_mgr.create_knowledge_base.assert_not_awaited()
|
||||
|
||||
@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)
|
||||
app = _app()
|
||||
app.rag_mgr.create_knowledge_base.return_value = SimpleNamespace(uuid='new_kb_uuid')
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
await KnowledgeService(app).create_knowledge_base(
|
||||
CONTEXT,
|
||||
{'knowledge_engine_plugin_id': 'author/engine'},
|
||||
)
|
||||
|
||||
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'
|
||||
assert app.rag_mgr.create_knowledge_base.await_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()
|
||||
app = _app()
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb1', 'name': 'Updated'}
|
||||
service = KnowledgeService(app)
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
|
||||
# Pass both mutable and immutable fields
|
||||
await service.update_knowledge_base(
|
||||
CONTEXT,
|
||||
'kb1',
|
||||
{
|
||||
'name': 'New Name',
|
||||
'description': 'New desc',
|
||||
'uuid': 'should_be_filtered', # immutable
|
||||
'uuid': 'should_be_filtered',
|
||||
},
|
||||
)
|
||||
|
||||
# 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
|
||||
update_statement = app.persistence_mgr.execute_async.await_args_list[0].args[0]
|
||||
params = update_statement.compile().params
|
||||
assert params['name'] == 'New Name'
|
||||
assert params['description'] == 'New desc'
|
||||
assert 'uuid' not in params
|
||||
app.rag_mgr.remove_knowledge_base_from_runtime.assert_awaited_once_with(CONTEXT, 'kb1')
|
||||
|
||||
@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()
|
||||
app = _app()
|
||||
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb1'}
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
await KnowledgeService(app).update_knowledge_base(
|
||||
CONTEXT,
|
||||
'kb1',
|
||||
{'uuid': 'should_be_filtered'},
|
||||
)
|
||||
|
||||
# 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()
|
||||
app.persistence_mgr.execute_async.assert_not_awaited()
|
||||
app.rag_mgr.remove_knowledge_base_from_runtime.assert_not_awaited()
|
||||
|
||||
|
||||
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']}}
|
||||
)
|
||||
app = _app()
|
||||
app.rag_mgr.get_knowledge_base_details.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
|
||||
await KnowledgeService(app)._check_doc_capability(CONTEXT, 'kb1', 'document upload')
|
||||
|
||||
@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)
|
||||
app = _app()
|
||||
|
||||
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)
|
||||
with pytest.raises(WorkspaceNotFoundError, match='Knowledge base not found'):
|
||||
await KnowledgeService(app)._check_doc_capability(
|
||||
CONTEXT,
|
||||
'nonexistent',
|
||||
'test operation',
|
||||
)
|
||||
|
||||
@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']}}
|
||||
)
|
||||
app = _app()
|
||||
app.rag_mgr.get_knowledge_base_details.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)
|
||||
with pytest.raises(Exception, match='does not support document upload'):
|
||||
await KnowledgeService(app)._check_doc_capability(
|
||||
CONTEXT,
|
||||
'kb1',
|
||||
'document upload',
|
||||
)
|
||||
|
||||
|
||||
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'}]
|
||||
)
|
||||
app = _app()
|
||||
app.plugin_connector.list_knowledge_engines.return_value = [{'id': 'engine1', 'name': 'Engine 1'}]
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.list_knowledge_engines()
|
||||
result = await KnowledgeService(app).list_knowledge_engines(CONTEXT)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]['id'] == 'engine1'
|
||||
assert result == [{'id': 'engine1', 'name': 'Engine 1'}]
|
||||
|
||||
@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
|
||||
app = _app()
|
||||
app.plugin_connector.is_enable_plugin = False
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.list_knowledge_engines()
|
||||
|
||||
assert result == []
|
||||
assert await KnowledgeService(app).list_knowledge_engines(CONTEXT) == []
|
||||
app.plugin_connector.list_knowledge_engines.assert_not_awaited()
|
||||
|
||||
@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'))
|
||||
app = _app()
|
||||
app.plugin_connector.list_knowledge_engines.side_effect = RuntimeError('Connection error')
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.list_knowledge_engines()
|
||||
|
||||
assert result == []
|
||||
mock_app.logger.warning.assert_called_once()
|
||||
assert await KnowledgeService(app).list_knowledge_engines(CONTEXT) == []
|
||||
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']},
|
||||
]
|
||||
)
|
||||
app = _app()
|
||||
app.plugin_connector.list_parsers.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()
|
||||
result = await KnowledgeService(app).list_parsers(CONTEXT)
|
||||
|
||||
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']},
|
||||
]
|
||||
)
|
||||
app = _app()
|
||||
app.plugin_connector.list_parsers.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')
|
||||
result = await KnowledgeService(app).list_parsers(CONTEXT, 'application/pdf')
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]['id'] == 'parser2'
|
||||
assert result == [{'id': 'parser2', 'supported_mime_types': ['application/pdf']}]
|
||||
|
||||
@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
|
||||
app = _app()
|
||||
app.plugin_connector.is_enable_plugin = False
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.list_parsers()
|
||||
|
||||
assert result == []
|
||||
assert await KnowledgeService(app).list_parsers(CONTEXT) == []
|
||||
app.plugin_connector.list_parsers.assert_not_awaited()
|
||||
|
||||
|
||||
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'}}}
|
||||
)
|
||||
app = _app()
|
||||
app.plugin_connector.get_rag_creation_schema.return_value = {'properties': {'name': {'type': 'string'}}}
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_engine_creation_schema('author/engine')
|
||||
result = await KnowledgeService(app).get_engine_creation_schema(
|
||||
CONTEXT,
|
||||
'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'}}}
|
||||
)
|
||||
app = _app()
|
||||
app.plugin_connector.get_rag_retrieval_schema.return_value = {'properties': {'top_k': {'type': 'integer'}}}
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_engine_retrieval_schema('author/engine')
|
||||
result = await KnowledgeService(app).get_engine_retrieval_schema(
|
||||
CONTEXT,
|
||||
'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'))
|
||||
app = _app()
|
||||
app.plugin_connector.get_rag_creation_schema.side_effect = RuntimeError('Plugin error')
|
||||
|
||||
service = knowledge_module.KnowledgeService(mock_app)
|
||||
result = await service.get_engine_creation_schema('author/engine')
|
||||
result = await KnowledgeService(app).get_engine_creation_schema(
|
||||
CONTEXT,
|
||||
'author/engine',
|
||||
)
|
||||
|
||||
assert result == {}
|
||||
mock_app.logger.warning.assert_called_once()
|
||||
app.logger.warning.assert_called_once()
|
||||
|
||||
|
||||
class TestKnowledgeBaseSecretViews:
|
||||
@pytest.mark.asyncio
|
||||
async def test_creation_settings_are_redacted_for_resource_view_only(self):
|
||||
app = _app()
|
||||
raw = {
|
||||
'uuid': 'kb-secret',
|
||||
'creation_settings': {
|
||||
'dify_apikey': 'dify-secret',
|
||||
'headers': {'Authorization': 'Bearer secret'},
|
||||
},
|
||||
}
|
||||
app.rag_mgr.get_all_knowledge_base_details.return_value = [raw]
|
||||
service = KnowledgeService(app)
|
||||
|
||||
redacted = await service.get_knowledge_bases(CONTEXT)
|
||||
manager_view = await service.get_knowledge_bases(CONTEXT, include_secret=True)
|
||||
|
||||
assert redacted[0]['creation_settings']['dify_apikey'] == '***'
|
||||
assert redacted[0]['creation_settings']['headers']['Authorization'] == '***'
|
||||
assert manager_view[0]['creation_settings']['dify_apikey'] == 'dify-secret'
|
||||
assert raw['creation_settings']['dify_apikey'] == 'dify-secret'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_masked_creation_secret_is_rejected(self):
|
||||
app = _app()
|
||||
|
||||
with pytest.raises(ValueError, match='no existing value'):
|
||||
await KnowledgeService(app).create_knowledge_base(
|
||||
CONTEXT,
|
||||
{
|
||||
'knowledge_engine_plugin_id': 'author/engine',
|
||||
'creation_settings': {'dify_apikey': '***'},
|
||||
},
|
||||
)
|
||||
|
||||
app.rag_mgr.create_knowledge_base.assert_not_awaited()
|
||||
|
||||
Reference in New Issue
Block a user