test: 补充知识处理、检索与 Agent 回归用例
覆盖切块、向量检索、工具调用和历史预算;模型相关用例使用桩。
This commit is contained in:
1 parent
2dd062787d
commit
d9fe727c3f
123 files changed
+29005
No files matched your search
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
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.workers.knowledge_processor import process_queue_item, run_processor
|
||||
from app.models.knowledge_queue import QueueStatus
|
||||
from app.models.notification import NotificationType
|
||||
from uuid import uuid4
|
||||
import pytest_asyncio
|
||||
from datetime import datetime
|
||||
|
||||
@pytest.fixture
|
||||
def mock_db():
|
||||
"""Create a mock database session"""
|
||||
db = MagicMock()
|
||||
return db
|
||||
|
||||
@pytest.fixture
|
||||
def mock_queue_item():
|
||||
"""Create a mock queue item"""
|
||||
queue_item = MagicMock()
|
||||
queue_item.id = 1
|
||||
queue_item.organization_id = uuid4()
|
||||
queue_item.agent_id = str(uuid4())
|
||||
queue_item.user_id = uuid4()
|
||||
queue_item.source = "test_document.pdf"
|
||||
queue_item.status = QueueStatus.PENDING
|
||||
return queue_item
|
||||
|
||||
@pytest.fixture
|
||||
def mock_queue_repo(mock_queue_item):
|
||||
"""Create a mock queue repository"""
|
||||
repo = MagicMock()
|
||||
repo.get_by_id.return_value = mock_queue_item
|
||||
repo.get_pending.return_value = [mock_queue_item]
|
||||
return repo
|
||||
|
||||
@pytest.fixture
|
||||
def mock_knowledge_manager():
|
||||
"""Create a mock knowledge manager"""
|
||||
manager = MagicMock()
|
||||
manager.process_knowledge = AsyncMock()
|
||||
return manager
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def mock_dependencies(mock_db, mock_queue_repo, mock_knowledge_manager):
|
||||
"""Set up all mock dependencies"""
|
||||
# Create a context manager mock that returns mock_db
|
||||
mock_session_local = MagicMock()
|
||||
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
||||
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
with patch('app.workers.knowledge_processor.SessionLocal', mock_session_local), \
|
||||
patch('app.workers.knowledge_processor.KnowledgeQueueRepository', return_value=mock_queue_repo), \
|
||||
patch('app.workers.knowledge_processor.KnowledgeManager', return_value=mock_knowledge_manager), \
|
||||
patch('app.workers.knowledge_processor.send_fcm_notification', new_callable=AsyncMock) as mock_fcm:
|
||||
yield {
|
||||
'db': mock_db,
|
||||
'queue_repo': mock_queue_repo,
|
||||
'knowledge_manager': mock_knowledge_manager,
|
||||
'fcm': mock_fcm
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_queue_item_success(mock_dependencies, mock_queue_item):
|
||||
"""Test successful processing of a queue item"""
|
||||
# Execute
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
# Assert
|
||||
assert mock_queue_item.status == QueueStatus.COMPLETED
|
||||
mock_dependencies['knowledge_manager'].process_knowledge.assert_awaited_once_with(mock_queue_item)
|
||||
|
||||
# Verify notification was created
|
||||
mock_dependencies['db'].add.assert_called_once()
|
||||
notification_call = mock_dependencies['db'].add.call_args[0][0]
|
||||
assert notification_call.type == NotificationType.KNOWLEDGE_PROCESSED
|
||||
assert notification_call.user_id == mock_queue_item.user_id
|
||||
assert mock_queue_item.source in notification_call.message
|
||||
|
||||
# Verify FCM notification was sent
|
||||
mock_dependencies['fcm'].assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_queue_item_not_found(mock_dependencies):
|
||||
"""Test processing when queue item is not found"""
|
||||
# Setup
|
||||
mock_dependencies['queue_repo'].get_by_id.return_value = None
|
||||
|
||||
# Execute
|
||||
await process_queue_item(999)
|
||||
|
||||
# Assert
|
||||
mock_dependencies['knowledge_manager'].process_knowledge.assert_not_awaited()
|
||||
mock_dependencies['fcm'].assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_queue_item_error(mock_dependencies, mock_queue_item):
|
||||
"""Test error handling during queue item processing"""
|
||||
# Setup
|
||||
error_message = "Processing failed"
|
||||
mock_dependencies['knowledge_manager'].process_knowledge.side_effect = Exception(error_message)
|
||||
|
||||
# Execute and assert exception is raised
|
||||
with pytest.raises(Exception, match=error_message):
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
# Assert
|
||||
assert mock_queue_item.status == QueueStatus.FAILED
|
||||
|
||||
# Verify error notification was created
|
||||
mock_dependencies['db'].add.assert_called_once()
|
||||
notification_call = mock_dependencies['db'].add.call_args[0][0]
|
||||
assert notification_call.type == NotificationType.KNOWLEDGE_FAILED
|
||||
assert notification_call.user_id == mock_queue_item.user_id
|
||||
assert error_message in notification_call.message
|
||||
|
||||
# Verify FCM notification was sent
|
||||
mock_dependencies['fcm'].assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_processor_success(mock_dependencies):
|
||||
"""Test successful run of the processor"""
|
||||
from app.api.knowledge import PROCESSOR_STATUS
|
||||
|
||||
# Execute
|
||||
await run_processor()
|
||||
|
||||
# Assert
|
||||
assert not PROCESSOR_STATUS["is_running"]
|
||||
assert PROCESSOR_STATUS["error"] is None
|
||||
assert isinstance(PROCESSOR_STATUS["last_run"], str)
|
||||
mock_dependencies['queue_repo'].get_pending.assert_called_once()
|
||||
mock_dependencies['knowledge_manager'].process_knowledge.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_processor_no_pending_items(mock_dependencies):
|
||||
"""Test processor run with no pending items"""
|
||||
# Setup
|
||||
mock_dependencies['queue_repo'].get_pending.return_value = []
|
||||
|
||||
# Execute
|
||||
await run_processor()
|
||||
|
||||
# Assert
|
||||
mock_dependencies['knowledge_manager'].process_knowledge.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_processor_error(mock_dependencies):
|
||||
"""Test error handling in processor run"""
|
||||
# Setup
|
||||
error_message = "Processor error"
|
||||
mock_dependencies['queue_repo'].get_pending.side_effect = Exception(error_message)
|
||||
|
||||
# Execute and assert exception is raised
|
||||
with pytest.raises(Exception, match=error_message):
|
||||
await run_processor()
|
||||
|
||||
# Assert
|
||||
from app.api.knowledge import PROCESSOR_STATUS
|
||||
assert not PROCESSOR_STATUS["is_running"]
|
||||
assert PROCESSOR_STATUS["error"] == error_message
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Vector agent-filter re-assert
|
||||
#
|
||||
# Ingestion writes filters.agent_id from the queue row's single optional
|
||||
# agent_id. A run queued without one leaves every chunk with an empty list, and
|
||||
# retrieval matches on that filter — so the source is invisible to the agent
|
||||
# while the dashboard shows it as linked. Two production sources were left
|
||||
# unreadable this way. The relational links are the authority, so the worker
|
||||
# re-derives the filter from them after every successful run.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _linked_knowledge(*agent_ids):
|
||||
"""A knowledge row whose agent_links resolve to the given agent ids."""
|
||||
knowledge = MagicMock()
|
||||
knowledge.agent_links = [MagicMock(agent_id=a) for a in agent_ids]
|
||||
return knowledge
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_queue_item_reasserts_agent_links(mock_dependencies, mock_queue_item):
|
||||
"""After a successful run the chunks point at every linked agent."""
|
||||
agent_a, agent_b = str(uuid4()), str(uuid4())
|
||||
knowledge = _linked_knowledge(agent_a, agent_b)
|
||||
repo = MagicMock()
|
||||
repo.get_by_sources.return_value = [knowledge]
|
||||
|
||||
with patch('app.workers.knowledge_processor.KnowledgeRepository', return_value=repo), \
|
||||
patch('app.workers.knowledge_processor.knowledge_vector_links.sync_agent_ids') as sync:
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
sync.assert_called_once()
|
||||
_db, synced_knowledge, agent_ids = sync.call_args[0]
|
||||
assert synced_knowledge is knowledge
|
||||
# Both linked agents, not just the one the queue row happened to carry.
|
||||
assert sorted(agent_ids) == sorted([agent_a, agent_b])
|
||||
assert mock_queue_item.status == QueueStatus.COMPLETED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reasserts_agent_links_when_queue_row_has_no_agent(mock_dependencies, mock_queue_item):
|
||||
"""The production incident: a crawl queued with agent_id=None.
|
||||
|
||||
Ingestion writes an empty filter, so without this the source stays
|
||||
unreadable until someone unlinks and relinks it by hand.
|
||||
"""
|
||||
mock_queue_item.agent_id = None
|
||||
linked_agent = str(uuid4())
|
||||
knowledge = _linked_knowledge(linked_agent)
|
||||
repo = MagicMock()
|
||||
repo.get_by_sources.return_value = [knowledge]
|
||||
|
||||
with patch('app.workers.knowledge_processor.KnowledgeRepository', return_value=repo), \
|
||||
patch('app.workers.knowledge_processor.knowledge_vector_links.sync_agent_ids') as sync:
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
assert sync.call_args[0][2] == [linked_agent]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_link_sync_failure_does_not_fail_the_run(mock_dependencies, mock_queue_item):
|
||||
"""The content is already ingested — a filter refresh must not lose the run."""
|
||||
repo = MagicMock()
|
||||
repo.get_by_sources.return_value = [_linked_knowledge(str(uuid4()))]
|
||||
|
||||
with patch('app.workers.knowledge_processor.KnowledgeRepository', return_value=repo), \
|
||||
patch('app.workers.knowledge_processor.knowledge_vector_links.sync_agent_ids',
|
||||
side_effect=Exception("vector store unavailable")):
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
assert mock_queue_item.status == QueueStatus.COMPLETED
|
||||
mock_dependencies['fcm'].assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_source_skips_sync_without_error(mock_dependencies, mock_queue_item):
|
||||
"""No knowledge row yet (e.g. the run failed before creating one)."""
|
||||
repo = MagicMock()
|
||||
repo.get_by_sources.return_value = []
|
||||
|
||||
with patch('app.workers.knowledge_processor.KnowledgeRepository', return_value=repo), \
|
||||
patch('app.workers.knowledge_processor.knowledge_vector_links.sync_agent_ids') as sync:
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
sync.assert_not_called()
|
||||
assert mock_queue_item.status == QueueStatus.COMPLETED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_crawl_fails_the_queue_item(mock_dependencies, mock_queue_item):
|
||||
"""Regression: a crawl that stored nothing used to be marked COMPLETED.
|
||||
|
||||
Production logged "Crawling completed - 0 URLs, 0 documents, 0.00s" followed
|
||||
by "Marked queue item as completed", so a broken source was indistinguishable
|
||||
from a healthy one in the dashboard.
|
||||
"""
|
||||
from app.knowledge.enhanced_website_reader import EmptyCrawlError
|
||||
|
||||
mock_dependencies['knowledge_manager'].process_knowledge.side_effect = EmptyCrawlError(
|
||||
"No content could be read from https://example.com."
|
||||
)
|
||||
|
||||
with pytest.raises(EmptyCrawlError):
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
assert mock_queue_item.status == QueueStatus.FAILED
|
||||
notification = mock_dependencies['db'].add.call_args[0][0]
|
||||
assert notification.type == NotificationType.KNOWLEDGE_FAILED
|
||||
assert "No content could be read" in notification.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_knowledge_rows_sync_the_union_of_their_agents(
|
||||
mock_dependencies, mock_queue_item
|
||||
):
|
||||
"""Two knowledge rows can share one URL — and one set of vector chunks.
|
||||
|
||||
Syncing them one after another would let the last write drop the other
|
||||
row's agents, so the union is applied once.
|
||||
"""
|
||||
agent_a, agent_b = str(uuid4()), str(uuid4())
|
||||
repo = MagicMock()
|
||||
repo.get_by_sources.return_value = [
|
||||
_linked_knowledge(agent_a),
|
||||
_linked_knowledge(agent_b),
|
||||
]
|
||||
|
||||
with patch('app.workers.knowledge_processor.KnowledgeRepository', return_value=repo), \
|
||||
patch('app.workers.knowledge_processor.knowledge_vector_links.sync_agent_ids') as sync:
|
||||
await process_queue_item(mock_queue_item.id)
|
||||
|
||||
sync.assert_called_once()
|
||||
assert sync.call_args[0][2] == sorted([agent_a, agent_b])
|
||||
@@ -0,0 +1,166 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
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.
|
||||
|
||||
A source whose worker died partway does not spin forever. Seen in
|
||||
production: the processor was killed for memory mid-source, restarted, and
|
||||
only ever picks up pending items - so a new signup's site stayed
|
||||
"processing" for good (and one from July with it).
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from app.knowledge.failure_advice import FAILURE_ADVICE, FailureCause
|
||||
from app.models.knowledge_queue import KnowledgeQueue, QueueStatus
|
||||
from app.models.notification import Notification, NotificationType
|
||||
from app.workers import knowledge_processor
|
||||
from app.workers.knowledge_processor import ABANDONED_AFTER
|
||||
|
||||
SITE = "https://ian-vpn.site"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sessions(db):
|
||||
return sessionmaker(bind=db.get_bind())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def push(monkeypatch, sessions):
|
||||
monkeypatch.setattr(knowledge_processor, "SessionLocal", sessions)
|
||||
sent = AsyncMock()
|
||||
monkeypatch.setattr(knowledge_processor, "send_fcm_notification", sent)
|
||||
return sent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queued(db, test_organization):
|
||||
def make(status, last_progress_ago, user_id=None, source=SITE, source_type="website"):
|
||||
when = datetime.now(timezone.utc) - last_progress_ago
|
||||
item = KnowledgeQueue(organization_id=test_organization.id, user_id=user_id, source=source,
|
||||
source_type=source_type, status=status, created_at=when)
|
||||
db.add(item)
|
||||
db.commit()
|
||||
item.updated_at = when # set explicitly: onupdate would stamp now
|
||||
db.commit()
|
||||
return item.id
|
||||
return make
|
||||
|
||||
|
||||
def state(sessions, item_id):
|
||||
with sessions() as db:
|
||||
item = db.get(KnowledgeQueue, item_id)
|
||||
notes = [n for n in db.query(Notification).all()
|
||||
if (n.notification_metadata or {}).get("queue_id") == item_id]
|
||||
return item.status, item.error, [(n.user_id, n.type) for n in notes]
|
||||
|
||||
|
||||
LONG_AGO = ABANDONED_AFTER + timedelta(minutes=1)
|
||||
JUST_NOW = timedelta(minutes=2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_abandoned_source_fails_and_its_owner_is_told(push, sessions, queued, test_user):
|
||||
item_id = queued(QueueStatus.PROCESSING, LONG_AGO, user_id=test_user.id)
|
||||
|
||||
await knowledge_processor.fail_abandoned_items()
|
||||
|
||||
status, error, notes = state(sessions, item_id)
|
||||
assert status == QueueStatus.FAILED
|
||||
assert error == FAILURE_ADVICE[FailureCause.INTERRUPTED].format(url=SITE)
|
||||
assert notes == [(test_user.id, NotificationType.KNOWLEDGE_FAILED)]
|
||||
push.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_with_no_owner_fails_quietly(push, sessions, queued):
|
||||
item_id = queued(QueueStatus.PROCESSING, LONG_AGO) # a Shopify install's site
|
||||
|
||||
await knowledge_processor.fail_abandoned_items()
|
||||
|
||||
assert state(sessions, item_id)[:2] == (QueueStatus.FAILED,
|
||||
FAILURE_ADVICE[FailureCause.INTERRUPTED].format(url=SITE))
|
||||
push.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_source_still_being_read_is_left_alone(push, sessions, queued, test_user):
|
||||
item_id = queued(QueueStatus.PROCESSING, JUST_NOW, user_id=test_user.id)
|
||||
|
||||
await knowledge_processor.fail_abandoned_items()
|
||||
|
||||
assert state(sessions, item_id) == (QueueStatus.PROCESSING, None, [])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_at_the_edge_of_its_time_limit_is_left_alone(push, sessions, queued):
|
||||
"""A live run stops at its time limit; until well past it, it may still be
|
||||
another process's."""
|
||||
item_id = queued(QueueStatus.PROCESSING, ABANDONED_AFTER - timedelta(minutes=1))
|
||||
|
||||
await knowledge_processor.fail_abandoned_items()
|
||||
|
||||
assert state(sessions, item_id)[0] == QueueStatus.PROCESSING
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status", [QueueStatus.PENDING, QueueStatus.COMPLETED, QueueStatus.FAILED])
|
||||
async def test_only_processing_sources_are_touched(push, sessions, queued, status):
|
||||
item_id = queued(status, LONG_AGO)
|
||||
|
||||
await knowledge_processor.fail_abandoned_items()
|
||||
|
||||
assert state(sessions, item_id) == (status, None, [])
|
||||
|
||||
|
||||
def test_abandoned_means_well_past_the_time_limit():
|
||||
from app.core.config import settings
|
||||
assert ABANDONED_AFTER > timedelta(seconds=settings.KB_SOURCE_TIME_LIMIT_SECONDS)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_every_processor_run_checks_for_them(push, sessions, queued, monkeypatch):
|
||||
item_id = queued(QueueStatus.PROCESSING, LONG_AGO)
|
||||
monkeypatch.setattr(knowledge_processor, "process_queue_item", AsyncMock())
|
||||
|
||||
await knowledge_processor.run_processor()
|
||||
|
||||
assert state(sessions, item_id)[0] == QueueStatus.FAILED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_check_never_holds_up_pending_sources(push, sessions, queued, monkeypatch):
|
||||
pending = queued(QueueStatus.PENDING, JUST_NOW)
|
||||
monkeypatch.setattr(knowledge_processor, "fail_abandoned_items", AsyncMock(side_effect=RuntimeError("db hiccup")))
|
||||
processed = AsyncMock()
|
||||
monkeypatch.setattr(knowledge_processor, "process_queue_item", processed)
|
||||
|
||||
await knowledge_processor.run_processor()
|
||||
|
||||
processed.assert_awaited_once_with(pending)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_uploaded_file_is_named_by_its_name_not_our_server_path(push, sessions, queued):
|
||||
item_id = queued(QueueStatus.PROCESSING, LONG_AGO, source="/app/temp/0e9f1d922da1467f_menu.pdf",
|
||||
source_type="pdf_file")
|
||||
|
||||
await knowledge_processor.fail_abandoned_items()
|
||||
|
||||
error = state(sessions, item_id)[1]
|
||||
assert "/app/temp" not in error
|
||||
assert error.startswith("0e9f1d922da1467f_menu.pdf could not be read")
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
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.
|
||||
|
||||
A source's own status never depends on telling someone about it. Seen in
|
||||
production: a Shopify install queues the shop's site with no user, so the
|
||||
"processed" notification hit the NOT NULL user_id constraint, and the failure
|
||||
handler then failed the same way and lost the FAILED status too.
|
||||
|
||||
Run against a real database: the bug was a database constraint.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from app.models.knowledge_queue import KnowledgeQueue, QueueStatus
|
||||
from app.models.notification import Notification, NotificationType
|
||||
from app.workers import knowledge_processor
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue_item(db, test_organization):
|
||||
def make(user_id=None, source="https://shop.example"):
|
||||
item = KnowledgeQueue(organization_id=test_organization.id, user_id=user_id, source=source,
|
||||
source_type="website", status=QueueStatus.PENDING)
|
||||
db.add(item)
|
||||
db.commit()
|
||||
return item.id
|
||||
return make
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sessions(db):
|
||||
"""New sessions on the fixture's database, as the worker opens its own."""
|
||||
return sessionmaker(bind=db.get_bind())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def worker(monkeypatch, sessions):
|
||||
"""The worker on the test database, with crawling and the side jobs stubbed."""
|
||||
manager = AsyncMock()
|
||||
monkeypatch.setattr(knowledge_processor, "SessionLocal", sessions)
|
||||
monkeypatch.setattr(knowledge_processor, "KnowledgeManager", lambda **kw: manager)
|
||||
monkeypatch.setattr(knowledge_processor, "_sync_vector_agent_links", lambda db, item: None)
|
||||
push = AsyncMock()
|
||||
monkeypatch.setattr(knowledge_processor, "send_fcm_notification", push)
|
||||
return manager, push
|
||||
|
||||
|
||||
def state(sessions, item_id):
|
||||
with sessions() as db:
|
||||
item = db.get(KnowledgeQueue, item_id)
|
||||
notifications = db.query(Notification).filter(
|
||||
Notification.notification_metadata.isnot(None)).all()
|
||||
return item.status, item.error, [n for n in notifications if n.notification_metadata.get("queue_id") == item_id]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_source_with_no_user_completes_without_a_notification(worker, sessions, queue_item):
|
||||
manager, push = worker
|
||||
item_id = queue_item(user_id=None) # what a Shopify install queues
|
||||
|
||||
await knowledge_processor.process_queue_item(item_id)
|
||||
|
||||
status, error, notifications = state(sessions, item_id)
|
||||
assert (status, error, notifications) == (QueueStatus.COMPLETED, None, [])
|
||||
push.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_source_with_no_user_that_fails_keeps_its_failed_status(worker, sessions, queue_item):
|
||||
manager, push = worker
|
||||
manager.process_knowledge.side_effect = RuntimeError("https://shop.example did not respond in time.")
|
||||
item_id = queue_item(user_id=None)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await knowledge_processor.process_queue_item(item_id)
|
||||
|
||||
status, error, notifications = state(sessions, item_id)
|
||||
assert status == QueueStatus.FAILED
|
||||
assert "did not respond in time" in error
|
||||
assert notifications == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_person_who_added_it_is_told_it_is_ready(worker, sessions, queue_item, test_user):
|
||||
manager, push = worker
|
||||
item_id = queue_item(user_id=test_user.id)
|
||||
|
||||
await knowledge_processor.process_queue_item(item_id)
|
||||
|
||||
status, _, notifications = state(sessions, item_id)
|
||||
assert status == QueueStatus.COMPLETED
|
||||
assert [(n.user_id, n.type) for n in notifications] == [(test_user.id, NotificationType.KNOWLEDGE_PROCESSED)]
|
||||
push.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_person_who_added_it_is_told_it_failed(worker, sessions, queue_item, test_user):
|
||||
manager, push = worker
|
||||
manager.process_knowledge.side_effect = RuntimeError("could not be found")
|
||||
item_id = queue_item(user_id=test_user.id)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await knowledge_processor.process_queue_item(item_id)
|
||||
|
||||
status, error, notifications = state(sessions, item_id)
|
||||
assert (status, error) == (QueueStatus.FAILED, "could not be found")
|
||||
assert [n.type for n in notifications] == [NotificationType.KNOWLEDGE_FAILED]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_notification_that_cannot_be_saved_does_not_fail_the_source(worker, sessions, queue_item, test_user, monkeypatch):
|
||||
manager, push = worker
|
||||
item_id = queue_item(user_id=test_user.id)
|
||||
real_init = Notification.__init__
|
||||
|
||||
def broken(self, **kwargs):
|
||||
real_init(self, **{**kwargs, "user_id": None}) # violates NOT NULL on flush
|
||||
monkeypatch.setattr(Notification, "__init__", broken)
|
||||
|
||||
await knowledge_processor.process_queue_item(item_id)
|
||||
|
||||
status, error, notifications = state(sessions, item_id)
|
||||
assert (status, error, notifications) == (QueueStatus.COMPLETED, None, [])
|
||||
push.assert_not_called()
|
||||
|
||||
|
||||
def test_the_service_starts_firebase_before_its_loops():
|
||||
order = []
|
||||
with patch("app.services.firebase.initialize_firebase", side_effect=lambda: order.append("firebase")), \
|
||||
patch.object(knowledge_processor.asyncio, "run", side_effect=lambda coro: (coro.close(), order.append("loops"))):
|
||||
knowledge_processor.run_service()
|
||||
assert order == ["firebase", "loops"]
|
||||
|
||||
|
||||
def test_the_service_still_runs_without_firebase():
|
||||
started = []
|
||||
with patch("app.services.firebase.initialize_firebase", side_effect=ValueError("no credentials")), \
|
||||
patch.object(knowledge_processor.asyncio, "run", side_effect=lambda coro: (coro.close(), started.append(True))):
|
||||
knowledge_processor.run_service()
|
||||
assert started == [True]
|
||||
Reference in new issue
Block a user