270 lines
9.6 KiB
Python
270 lines
9.6 KiB
Python
"""Tests for TaskDispatcher - 任务分发器"""
|
|
|
|
import json
|
|
import uuid
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from agentkit.core.dispatcher import TaskDispatcher
|
|
from agentkit.core.exceptions import TaskDispatchError, TaskNotFoundError
|
|
from agentkit.core.protocol import AgentStatus, TaskResult, TaskStatus
|
|
|
|
|
|
class _ColumnMock:
|
|
"""Mock for SQLAlchemy column attributes that supports comparison operators."""
|
|
|
|
def __init__(self, name):
|
|
self._name = name
|
|
|
|
def __eq__(self, other):
|
|
return MagicMock()
|
|
|
|
def __ne__(self, other):
|
|
return MagicMock()
|
|
|
|
def __lt__(self, other):
|
|
return MagicMock()
|
|
|
|
def __le__(self, other):
|
|
return MagicMock()
|
|
|
|
def __gt__(self, other):
|
|
return MagicMock()
|
|
|
|
def __ge__(self, other):
|
|
return MagicMock()
|
|
|
|
def like(self, pattern):
|
|
return MagicMock()
|
|
|
|
def desc(self):
|
|
return MagicMock()
|
|
|
|
|
|
class MockAgentModel:
|
|
"""Mock Agent ORM model with class-level column mocks."""
|
|
name = _ColumnMock("name")
|
|
status = _ColumnMock("status")
|
|
agent_type = _ColumnMock("agent_type")
|
|
id = _ColumnMock("id")
|
|
|
|
def __init__(self, **kwargs):
|
|
self.id = kwargs.get("id", uuid.uuid4())
|
|
self.name = kwargs.get("name", "test_agent")
|
|
self.agent_type = kwargs.get("agent_type", "test")
|
|
self.status = kwargs.get("status", AgentStatus.ONLINE)
|
|
self.version = kwargs.get("version", "1.0")
|
|
self.endpoint = kwargs.get("endpoint", "http://localhost:8000")
|
|
self.description = kwargs.get("description", "Test agent")
|
|
|
|
|
|
class MockTaskModel:
|
|
"""Mock Task ORM model with class-level column mocks."""
|
|
id = _ColumnMock("id")
|
|
agent_id = _ColumnMock("agent_id")
|
|
task_type = _ColumnMock("task_type")
|
|
status = _ColumnMock("status")
|
|
priority = _ColumnMock("priority")
|
|
input_data = _ColumnMock("input_data")
|
|
output_data = _ColumnMock("output_data")
|
|
error_message = _ColumnMock("error_message")
|
|
started_at = _ColumnMock("started_at")
|
|
completed_at = _ColumnMock("completed_at")
|
|
organization_id = _ColumnMock("organization_id")
|
|
created_by = _ColumnMock("created_by")
|
|
project_id = _ColumnMock("project_id")
|
|
scheduled_at = _ColumnMock("scheduled_at")
|
|
created_at = _ColumnMock("created_at")
|
|
|
|
def __init__(self, **kwargs):
|
|
self.id = kwargs.get("id", uuid.uuid4())
|
|
self.agent_id = kwargs.get("agent_id", uuid.uuid4())
|
|
self.task_type = kwargs.get("task_type", "test_task")
|
|
self.status = kwargs.get("status", TaskStatus.PENDING)
|
|
self.priority = kwargs.get("priority", 1)
|
|
self.input_data = kwargs.get("input_data", {})
|
|
self.output_data = kwargs.get("output_data", None)
|
|
self.error_message = kwargs.get("error_message", None)
|
|
self.started_at = kwargs.get("started_at", None)
|
|
self.completed_at = kwargs.get("completed_at", None)
|
|
self.organization_id = kwargs.get("organization_id", uuid.uuid4())
|
|
self.created_by = kwargs.get("created_by", None)
|
|
self.project_id = kwargs.get("project_id", None)
|
|
self.scheduled_at = kwargs.get("scheduled_at", None)
|
|
self.created_at = kwargs.get("created_at", None)
|
|
|
|
|
|
class MockTaskLogModel:
|
|
"""Mock TaskLog ORM model with class-level column mocks."""
|
|
id = _ColumnMock("id")
|
|
task_id = _ColumnMock("task_id")
|
|
agent_id = _ColumnMock("agent_id")
|
|
log_level = _ColumnMock("log_level")
|
|
message = _ColumnMock("message")
|
|
|
|
def __init__(self, **kwargs):
|
|
self.id = kwargs.get("id", uuid.uuid4())
|
|
self.task_id = kwargs.get("task_id", uuid.uuid4())
|
|
self.agent_id = kwargs.get("agent_id", uuid.uuid4())
|
|
self.log_level = kwargs.get("log_level", "info")
|
|
self.message = kwargs.get("message", "")
|
|
|
|
|
|
def _make_mock_session(agent=None, task=None, log_entries=None):
|
|
"""Create a mock async session that simulates SQLAlchemy queries."""
|
|
session = AsyncMock()
|
|
|
|
async def mock_execute(stmt):
|
|
result = MagicMock()
|
|
|
|
if agent is not None:
|
|
result.scalar_one_or_none.return_value = agent
|
|
elif task is not None:
|
|
result.scalar_one_or_none.return_value = task
|
|
result.scalars.return_value.all.return_value = [task] if task else []
|
|
else:
|
|
result.scalar_one_or_none.return_value = None
|
|
result.scalars.return_value.all.return_value = log_entries or []
|
|
|
|
if log_entries is not None:
|
|
result.scalars.return_value.all.return_value = log_entries
|
|
|
|
return result
|
|
|
|
session.execute = mock_execute
|
|
session.add = MagicMock()
|
|
session.commit = AsyncMock()
|
|
session.rollback = AsyncMock()
|
|
session.refresh = AsyncMock()
|
|
|
|
return session
|
|
|
|
|
|
def _make_dispatcher(agent=None, task=None, log_entries=None):
|
|
"""Create a TaskDispatcher with mocked dependencies."""
|
|
mock_session = _make_mock_session(agent=agent, task=task, log_entries=log_entries)
|
|
|
|
session_factory = MagicMock()
|
|
session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_session)
|
|
session_factory.return_value.__aexit__ = AsyncMock(return_value=False)
|
|
|
|
mock_redis = AsyncMock()
|
|
mock_redis.lpush = AsyncMock()
|
|
redis_factory = AsyncMock(return_value=mock_redis)
|
|
|
|
dispatcher = TaskDispatcher(
|
|
redis_factory=redis_factory,
|
|
session_factory=session_factory,
|
|
agent_model=MockAgentModel,
|
|
task_model=MockTaskModel,
|
|
task_log_model=MockTaskLogModel,
|
|
)
|
|
|
|
return dispatcher, mock_session, mock_redis
|
|
|
|
|
|
_mock_select = MagicMock()
|
|
|
|
|
|
class TestTaskDispatcherDispatch:
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_dispatch_to_online_agent(self, make_task):
|
|
"""分发任务到在线 Agent"""
|
|
agent = MockAgentModel(name="test_agent", status=AgentStatus.ONLINE)
|
|
dispatcher, session, redis = _make_dispatcher(agent=agent)
|
|
task_id = str(uuid.uuid4())
|
|
task = make_task(task_id=task_id, agent_name="test_agent")
|
|
|
|
result_task_id = await dispatcher.dispatch(task)
|
|
assert result_task_id == task_id
|
|
redis.lpush.assert_called_once()
|
|
|
|
# Verify the queue key format
|
|
call_args = redis.lpush.call_args
|
|
assert call_args[0][0] == "agent:test_agent:tasks"
|
|
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_dispatch_agent_not_found(self, make_task):
|
|
"""分发到不存在的 Agent 抛出异常"""
|
|
dispatcher, session, redis = _make_dispatcher(agent=None)
|
|
task_id = str(uuid.uuid4())
|
|
task = make_task(task_id=task_id, agent_name="nonexistent")
|
|
|
|
with pytest.raises(TaskDispatchError):
|
|
await dispatcher.dispatch(task)
|
|
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_dispatch_agent_offline(self, make_task):
|
|
"""分发到离线 Agent 抛出异常"""
|
|
agent = MockAgentModel(name="offline_agent", status=AgentStatus.OFFLINE)
|
|
dispatcher, session, redis = _make_dispatcher(agent=agent)
|
|
task_id = str(uuid.uuid4())
|
|
task = make_task(task_id=task_id, agent_name="offline_agent")
|
|
|
|
with pytest.raises(TaskDispatchError):
|
|
await dispatcher.dispatch(task)
|
|
|
|
|
|
class TestTaskDispatcherCancel:
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_cancel_pending_task(self, make_task):
|
|
"""取消待执行的任务"""
|
|
task_uuid = uuid.uuid4()
|
|
task = MockTaskModel(id=task_uuid, status=TaskStatus.PENDING)
|
|
dispatcher, session, redis = _make_dispatcher(task=task)
|
|
|
|
await dispatcher.cancel_task(str(task_uuid))
|
|
assert task.status == TaskStatus.CANCELLED
|
|
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_cancel_completed_task(self, make_task):
|
|
"""取消已完成的任务不改变状态"""
|
|
task_uuid = uuid.uuid4()
|
|
task = MockTaskModel(id=task_uuid, status=TaskStatus.COMPLETED)
|
|
dispatcher, session, redis = _make_dispatcher(task=task)
|
|
|
|
await dispatcher.cancel_task(str(task_uuid))
|
|
# Status should remain COMPLETED (not changed to CANCELLED)
|
|
assert task.status == TaskStatus.COMPLETED
|
|
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_cancel_nonexistent_task(self):
|
|
"""取消不存在的任务抛出异常"""
|
|
dispatcher, session, redis = _make_dispatcher(task=None)
|
|
|
|
with pytest.raises(TaskNotFoundError):
|
|
await dispatcher.cancel_task(str(uuid.uuid4()))
|
|
|
|
|
|
class TestTaskDispatcherHandleResult:
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_handle_completed_result(self, make_task, make_result):
|
|
"""处理成功结果"""
|
|
task_uuid = uuid.uuid4()
|
|
task = MockTaskModel(id=task_uuid, status=TaskStatus.RUNNING)
|
|
dispatcher, session, redis = _make_dispatcher(task=task)
|
|
|
|
result = make_result(task_id=str(task_uuid), status=TaskStatus.COMPLETED)
|
|
await dispatcher.handle_result(result)
|
|
|
|
assert task.status == TaskStatus.COMPLETED
|
|
assert task.output_data == result.output_data
|
|
|
|
@patch("sqlalchemy.select", _mock_select)
|
|
async def test_handle_failed_result(self, make_task, make_result):
|
|
"""处理失败结果"""
|
|
task_uuid = uuid.uuid4()
|
|
task = MockTaskModel(id=task_uuid, status=TaskStatus.RUNNING)
|
|
dispatcher, session, redis = _make_dispatcher(task=task)
|
|
|
|
result = make_result(
|
|
task_id=str(task_uuid),
|
|
status=TaskStatus.FAILED,
|
|
error_message="Something went wrong",
|
|
)
|
|
await dispatcher.handle_result(result)
|
|
|
|
assert task.status == TaskStatus.FAILED
|
|
assert task.error_message == "Something went wrong"
|