test: 补充知识处理、检索与 Agent 回归用例
覆盖切块、向量检索、工具调用和历史预算；模型相关用例使用桩。
This commit is contained in:
WangShen007 committed
1 parent
362d79a5c2
commit
f4d8019829
123 files changed
+29005
No files matched your search
@@ -0,0 +1,340 @@
|
||||
"""
|
||||
Copyright 2024-2026 ChatterMate
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from app.knowledge.knowledge_base import KnowledgeManager
|
||||
from app.models.knowledge import Knowledge, SourceType
|
||||
from app.models.knowledge_to_agent import KnowledgeToAgent
|
||||
from uuid import UUID, uuid4
|
||||
import os
|
||||
from agno.vectordb.pgvector import PgVector, SearchType
|
||||
|
||||
@pytest.fixture
|
||||
def mock_db():
|
||||
"""Create a mock database session"""
|
||||
db = MagicMock()
|
||||
return db
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ai_config():
|
||||
"""Create a mock AI config"""
|
||||
config = MagicMock()
|
||||
config.encrypted_api_key = "encrypted_key"
|
||||
return config
|
||||
|
||||
@pytest.fixture
|
||||
def mock_knowledge():
|
||||
"""Create a mock knowledge entry"""
|
||||
knowledge_id = uuid4()
|
||||
org_id = uuid4()
|
||||
return Knowledge(
|
||||
id=knowledge_id,
|
||||
organization_id=org_id,
|
||||
source="test_source",
|
||||
source_type=SourceType.FILE,
|
||||
table_name="test_table",
|
||||
schema="test_schema"
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vector_db():
|
||||
"""Create a mock vector database"""
|
||||
vector_db = MagicMock(spec=PgVector)
|
||||
vector_db.table_name = "test_table"
|
||||
vector_db.schema = "test_schema"
|
||||
return vector_db
|
||||
|
||||
@pytest.fixture
|
||||
def knowledge_manager(mock_db, mock_ai_config, mock_vector_db):
|
||||
"""Create a KnowledgeManager instance with mocked dependencies"""
|
||||
org_id = uuid4()
|
||||
agent_id = uuid4()
|
||||
|
||||
with patch('app.knowledge.knowledge_base.SessionLocal') as mock_session_local, \
|
||||
patch('app.knowledge.knowledge_base.AIConfigRepository') as mock_ai_config_repo, \
|
||||
patch('app.knowledge.knowledge_base.OptimizedPgVector') as mock_pg_vector, \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo, \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeToAgentRepository') as mock_link_repo:
|
||||
|
||||
# Configure basic mocks
|
||||
mock_session_local.return_value.__enter__.return_value = mock_db
|
||||
mock_ai_config_repo.return_value.get_active_config.return_value = mock_ai_config
|
||||
mock_pg_vector.return_value = mock_vector_db
|
||||
|
||||
# Configure repository mocks
|
||||
mock_knowledge_repo_instance = MagicMock()
|
||||
mock_link_repo_instance = MagicMock()
|
||||
mock_knowledge_repo.return_value = mock_knowledge_repo_instance
|
||||
mock_link_repo.return_value = mock_link_repo_instance
|
||||
|
||||
manager = KnowledgeManager(org_id=org_id, agent_id=str(agent_id))
|
||||
|
||||
# Replace repository instances with mocks
|
||||
manager.knowledge_repo = mock_knowledge_repo_instance
|
||||
manager.link_repo = mock_link_repo_instance
|
||||
|
||||
return manager
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_pdf_urls_success(knowledge_manager):
|
||||
"""A PDF URL is downloaded through the guarded fetcher, then indexed."""
|
||||
urls = ["https://example.com/docs/test.pdf"]
|
||||
|
||||
with patch('app.knowledge.knowledge_base.download_pdf') as mock_download, \
|
||||
patch.object(knowledge_manager, 'add_pdf_files', AsyncMock(return_value=True)) as mock_add_files, \
|
||||
patch('app.knowledge.knowledge_base.delete_file_from_s3', AsyncMock()) as mock_delete:
|
||||
result = await knowledge_manager.add_pdf_urls(urls)
|
||||
|
||||
assert result is True
|
||||
url, temp_path = mock_download.call_args.args
|
||||
assert url == urls[0]
|
||||
assert os.path.basename(temp_path) == "test.pdf"
|
||||
mock_add_files.assert_awaited_once_with([temp_path], filename="test")
|
||||
mock_delete.assert_not_called() # not one of our S3 uploads
|
||||
assert not os.path.exists(os.path.dirname(temp_path)) # temp dir cleaned up
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_pdf_urls_failure(knowledge_manager):
|
||||
"""A refused or failed download is reported, and nothing is indexed."""
|
||||
from app.knowledge.url_safety import BlockedHostError
|
||||
|
||||
with patch('app.knowledge.knowledge_base.download_pdf',
|
||||
side_effect=BlockedHostError("internal")), \
|
||||
patch.object(knowledge_manager, 'add_pdf_files', AsyncMock()) as mock_add_files:
|
||||
result = await knowledge_manager.add_pdf_urls(["http://169.254.169.254/latest"])
|
||||
|
||||
assert result is False
|
||||
mock_add_files.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("url, name", [
|
||||
("https://example.com/a/guide.pdf", "guide"),
|
||||
("https://example.com/a/My%20Guide.pdf?x=1", "My Guide"),
|
||||
# An encoded traversal decodes to a separator; only the last part survives.
|
||||
("https://b.s3.amazonaws.com/knowledge/..%2F..%2Fapp%2Fmain.py", "main"),
|
||||
("https://example.com/", "example.com"),
|
||||
])
|
||||
def test_pdf_name_from_url_never_carries_a_path(url, name):
|
||||
from app.knowledge.knowledge_base import _pdf_name_from_url
|
||||
assert _pdf_name_from_url(url) == name
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_this_orgs_uploaded_pdf_is_deleted_from_s3(knowledge_manager):
|
||||
from app.knowledge.knowledge_base import settings
|
||||
own = f"https://{settings.S3_BUCKET}.s3.us-east-1.amazonaws.com/knowledge/{knowledge_manager.org_id}/f.pdf?X-Amz-Signature=x"
|
||||
other_org = f"https://{settings.S3_BUCKET}.s3.us-east-1.amazonaws.com/knowledge/{uuid4()}/f.pdf"
|
||||
# Another bucket, laid out so a path-style read lands on our org's key.
|
||||
foreign = f"https://evil.s3.amazonaws.com/x/knowledge/{knowledge_manager.org_id}/f.pdf"
|
||||
|
||||
with patch('app.knowledge.knowledge_base.delete_file_from_s3', AsyncMock(return_value=True)) as mock_delete:
|
||||
for url in (own, other_org, foreign):
|
||||
await knowledge_manager._delete_uploaded_pdf(url)
|
||||
|
||||
mock_delete.assert_awaited_once_with(own)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_websites_success(knowledge_manager, mock_knowledge):
|
||||
"""Test successful addition of websites"""
|
||||
# Setup
|
||||
urls = ["http://example.com"]
|
||||
max_links = 5
|
||||
|
||||
with patch('app.knowledge.knowledge_base.EnhancedWebsiteKnowledgeBase') as mock_web_kb, \
|
||||
patch.object(knowledge_manager, '_add_knowledge_source', return_value=mock_knowledge) as mock_add_source:
|
||||
# Configure mock
|
||||
mock_web_kb_instance = MagicMock()
|
||||
mock_web_kb.return_value = mock_web_kb_instance
|
||||
|
||||
# Execute
|
||||
result = await knowledge_manager.add_websites(urls, max_links)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
mock_web_kb.assert_called_once()
|
||||
mock_web_kb_instance.load.assert_called_once()
|
||||
mock_add_source.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_pdf_files_success(knowledge_manager, mock_knowledge):
|
||||
"""Test successful addition of PDF files"""
|
||||
# Setup
|
||||
files = ["/path/to/test.pdf"]
|
||||
|
||||
with patch('app.knowledge.knowledge_base.EnhancedPDFKnowledgeBase') as mock_pdf_kb, \
|
||||
patch('app.knowledge.knowledge_base.PDFImageReader') as mock_pdf_reader, \
|
||||
patch('os.path.exists') as mock_exists, \
|
||||
patch('os.path.getsize') as mock_getsize, \
|
||||
patch('os.path.basename') as mock_basename, \
|
||||
patch('os.path.splitext') as mock_splitext, \
|
||||
patch.object(knowledge_manager, '_add_knowledge_source', return_value=mock_knowledge) as mock_add_source:
|
||||
# Configure mocks
|
||||
mock_pdf_kb_instance = MagicMock()
|
||||
mock_pdf_kb.return_value = mock_pdf_kb_instance
|
||||
mock_exists.return_value = True
|
||||
mock_getsize.return_value = 1024
|
||||
mock_basename.return_value = "test.pdf"
|
||||
mock_splitext.return_value = ("test", ".pdf")
|
||||
|
||||
# Execute
|
||||
result = await knowledge_manager.add_pdf_files(files)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
mock_pdf_kb.assert_called_once()
|
||||
mock_pdf_kb_instance.load.assert_called_once()
|
||||
mock_add_source.assert_called_once()
|
||||
|
||||
def test_get_knowledge_base_for_agent(knowledge_manager, mock_knowledge):
|
||||
"""Test getting knowledge base for a specific agent"""
|
||||
# Setup
|
||||
mock_knowledge.agent_links = [
|
||||
MagicMock(agent_id=knowledge_manager.agent_id)
|
||||
]
|
||||
|
||||
with patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo_class:
|
||||
mock_knowledge_repo_instance = MagicMock()
|
||||
mock_knowledge_repo_class.return_value = mock_knowledge_repo_instance
|
||||
mock_knowledge_repo_instance.get_by_agent.return_value = [mock_knowledge]
|
||||
|
||||
# Execute
|
||||
result = knowledge_manager.get_knowledge_base()
|
||||
|
||||
# Assert
|
||||
assert len(result) == 1
|
||||
assert result[0]['id'] == mock_knowledge.id
|
||||
assert result[0]['source'] == mock_knowledge.source
|
||||
assert result[0]['source_type'] == mock_knowledge.source_type
|
||||
mock_knowledge_repo_instance.get_by_agent.assert_called_once_with(knowledge_manager.agent_id)
|
||||
|
||||
def test_get_knowledge_base_for_org(knowledge_manager, mock_knowledge):
|
||||
"""Test getting knowledge base for an organization"""
|
||||
# Setup
|
||||
knowledge_manager.agent_id = None
|
||||
mock_knowledge.agent_links = []
|
||||
|
||||
with patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo_class:
|
||||
mock_knowledge_repo_instance = MagicMock()
|
||||
mock_knowledge_repo_class.return_value = mock_knowledge_repo_instance
|
||||
mock_knowledge_repo_instance.get_by_org.return_value = [mock_knowledge]
|
||||
|
||||
# Execute
|
||||
result = knowledge_manager.get_knowledge_base()
|
||||
|
||||
# Assert
|
||||
assert len(result) == 1
|
||||
assert result[0]['id'] == mock_knowledge.id
|
||||
assert result[0]['source'] == mock_knowledge.source
|
||||
assert result[0]['source_type'] == mock_knowledge.source_type
|
||||
mock_knowledge_repo_instance.get_by_org.assert_called_once_with(knowledge_manager.org_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_knowledge_pdf_file(knowledge_manager):
|
||||
"""Test processing PDF file knowledge"""
|
||||
# Setup
|
||||
queue_item = MagicMock()
|
||||
queue_item.id = 123 # ensure a real integer ID to avoid DB adapter issues
|
||||
queue_item.source_type = "pdf_file"
|
||||
queue_item.source = "/path/to/test.pdf"
|
||||
|
||||
with patch.object(knowledge_manager, 'add_pdf_files') as mock_add_pdf_files, \
|
||||
patch('os.path.exists') as mock_exists, \
|
||||
patch('os.remove') as mock_remove, \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as mock_queue_repo_cls:
|
||||
# Configure mocks
|
||||
mock_add_pdf_files.return_value = True
|
||||
mock_exists.return_value = True
|
||||
mock_queue_repo = MagicMock()
|
||||
mock_queue_repo.update_progress.return_value = True
|
||||
mock_queue_repo.update_status.return_value = True
|
||||
mock_queue_repo_cls.return_value = mock_queue_repo
|
||||
|
||||
# Execute
|
||||
result = await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
mock_add_pdf_files.assert_called_once_with([queue_item.source])
|
||||
mock_exists.assert_called_once_with(queue_item.source)
|
||||
mock_remove.assert_called_once_with(queue_item.source)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_knowledge_unsupported_type(knowledge_manager):
|
||||
"""Test processing unsupported knowledge type"""
|
||||
# Setup
|
||||
queue_item = MagicMock()
|
||||
queue_item.source_type = "unsupported_type"
|
||||
|
||||
# Execute and Assert
|
||||
with pytest.raises(ValueError, match="Unsupported source type: unsupported_type"):
|
||||
await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
# ---- PDFs over a limit: refused with the reason, not "failed to process" ----
|
||||
|
||||
from app.knowledge.pdf_download import PdfTooLargeError # noqa: E402
|
||||
from app.knowledge.pdf_limits import PdfTooManyPages # noqa: E402
|
||||
from tests.knowledge.test_pdf_limits import pdf_with_pages # noqa: E402
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pdf_over_the_page_limit_is_refused_before_it_is_read(knowledge_manager, tmp_path):
|
||||
path = tmp_path / "catalogue.pdf"
|
||||
path.write_bytes(pdf_with_pages(4))
|
||||
|
||||
with patch('app.core.config.settings.KB_PDF_MAX_PAGES', 3), \
|
||||
patch('app.knowledge.knowledge_base.EnhancedPDFKnowledgeBase') as reader, \
|
||||
pytest.raises(PdfTooManyPages, match="4 pages; the limit is 3"):
|
||||
await knowledge_manager.add_pdf_files([str(path)])
|
||||
|
||||
reader.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pdf_link_over_the_size_limit_says_so_without_the_link(knowledge_manager):
|
||||
signed = "https://bucket.s3.amazonaws.com/knowledge/x/big.pdf?X-Amz-Signature=secret"
|
||||
with patch('app.knowledge.knowledge_base.download_pdf', side_effect=PdfTooLargeError(25 * 1024 * 1024)), \
|
||||
pytest.raises(PdfTooLargeError) as refused:
|
||||
await knowledge_manager.add_pdf_urls([signed])
|
||||
|
||||
assert "25 MB" in str(refused.value)
|
||||
assert "amazonaws" not in str(refused.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refused_upload_is_reported_and_its_copy_removed(knowledge_manager, tmp_path):
|
||||
path = tmp_path / "0e9f_catalogue.pdf"
|
||||
path.write_bytes(pdf_with_pages(4))
|
||||
queue_item = MagicMock(id=123, source_type="pdf_file", source=str(path))
|
||||
|
||||
with patch('app.core.config.settings.KB_PDF_MAX_PAGES', 3), \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as queue_repo, \
|
||||
pytest.raises(PdfTooManyPages):
|
||||
await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
assert not path.exists()
|
||||
queue_repo.return_value.update_status.assert_not_called() # never reported as completed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("source_type,method", [("pdf_file", "add_pdf_files"), ("pdf_url", "add_pdf_urls")])
|
||||
async def test_a_pdf_that_fails_is_never_marked_completed_first(knowledge_manager, source_type, method):
|
||||
queue_item = MagicMock(id=123, source_type=source_type, source="https://example.com/x.pdf")
|
||||
|
||||
with patch.object(knowledge_manager, method, AsyncMock(return_value=False)), \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as queue_repo, \
|
||||
pytest.raises(Exception, match="Failed to process knowledge source"):
|
||||
await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
queue_repo.return_value.update_status.assert_not_called()
|
||||
Reference in new issue
Block a user