feat: 添加自动化任务(automation_jobs)功能模块
This commit is contained in:
qzl
@@ -1,283 +1,287 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, time, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from models.automation_jobs import AutomationJobStatus, ScheduleType
|
||||
from v1.automation_jobs.repository import AutomationJobsRepository
|
||||
from v1.automation_jobs.schemas import (
|
||||
AutomationJobCreateRequest,
|
||||
AutomationJobUpdateRequest,
|
||||
)
|
||||
from schemas.automation import (
|
||||
AgentTool,
|
||||
AutomationJobConfig,
|
||||
ContextSource,
|
||||
ContextWindowMode,
|
||||
MessageContextConfig,
|
||||
)
|
||||
|
||||
|
||||
class _ExecuteResult:
|
||||
def __init__(self, value: object) -> None:
|
||||
self._value = value
|
||||
|
||||
def scalar_one_or_none(self) -> object:
|
||||
return self._value
|
||||
|
||||
def scalar_one(self) -> int:
|
||||
return self._value # type: ignore[return-value]
|
||||
def _make_config() -> AutomationJobConfig:
|
||||
return AutomationJobConfig(
|
||||
input_template="Hello",
|
||||
enabled_tools=[AgentTool.MEMORY_WRITE],
|
||||
context=MessageContextConfig(
|
||||
source=ContextSource.LATEST_CHAT,
|
||||
window_mode=ContextWindowMode.DAY,
|
||||
window_count=2,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _ScalarRows:
|
||||
def __init__(self, rows: list[object]) -> None:
|
||||
self._rows = rows
|
||||
|
||||
def all(self) -> list[object]:
|
||||
return self._rows
|
||||
|
||||
|
||||
class _ExecuteRowsResult:
|
||||
def __init__(self, rows: list[object]) -> None:
|
||||
self._rows = rows
|
||||
|
||||
def scalars(self) -> _ScalarRows:
|
||||
return _ScalarRows(self._rows)
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self) -> None:
|
||||
self.added: list[object] = []
|
||||
self.flushed = False
|
||||
self._execute_result: object = None
|
||||
self._return_rows: bool = False
|
||||
|
||||
def set_execute_result(self, value: object) -> None:
|
||||
self._execute_result = value
|
||||
self._return_rows = isinstance(value, list)
|
||||
|
||||
async def execute(self, stmt): # noqa: ANN001
|
||||
del stmt
|
||||
if self._return_rows:
|
||||
return _ExecuteRowsResult(self._execute_result)
|
||||
return _ExecuteResult(self._execute_result)
|
||||
|
||||
def add(self, obj: object) -> None:
|
||||
self.added.append(obj)
|
||||
|
||||
async def flush(self) -> None:
|
||||
self.flushed = True
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_session() -> _FakeSession:
|
||||
return _FakeSession()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def repository(fake_session: _FakeSession) -> AutomationJobsRepository:
|
||||
return AutomationJobsRepository(session=fake_session) # type: ignore[arg-type]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_job() -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
id=uuid4(),
|
||||
owner_id=uuid4(),
|
||||
bootstrap_key=None,
|
||||
def _make_create_request() -> AutomationJobCreateRequest:
|
||||
return AutomationJobCreateRequest(
|
||||
title="Test Job",
|
||||
config={"input_template": "Hello {name}"},
|
||||
schedule_type=ScheduleType.DAILY,
|
||||
run_at=datetime(2026, 3, 23, 0, 0, tzinfo=timezone.utc),
|
||||
next_run_at=datetime(2026, 3, 24, 0, 0, tzinfo=timezone.utc),
|
||||
timezone="UTC",
|
||||
run_at=time(9, 0, 0),
|
||||
timezone="Asia/Shanghai",
|
||||
status=AutomationJobStatus.ACTIVE,
|
||||
created_by=uuid4(),
|
||||
deleted_at=None,
|
||||
config=_make_config(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_by_owner_returns_jobs(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
sample_job: SimpleNamespace,
|
||||
) -> None:
|
||||
fake_session.set_execute_result([sample_job])
|
||||
|
||||
async def test_list_by_owner_returns_jobs() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
owner_id = uuid4()
|
||||
jobs = await repository.list_by_owner(owner_id)
|
||||
job_one = MagicMock()
|
||||
job_two = MagicMock()
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalars.return_value.all.return_value = [job_one, job_two]
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
assert len(jobs) == 1
|
||||
assert jobs[0].title == "Test Job"
|
||||
result = await repository.list_by_owner(owner_id)
|
||||
|
||||
assert result == [job_one, job_two]
|
||||
session.execute.assert_awaited_once()
|
||||
call_args = session.execute.call_args
|
||||
stmt = call_args[0][0]
|
||||
assert "owner_id" in str(stmt)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_by_owner_returns_empty_list(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
fake_session.set_execute_result([])
|
||||
|
||||
async def test_count_user_jobs_counts_non_bootstrap_jobs() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
owner_id = uuid4()
|
||||
jobs = await repository.list_by_owner(owner_id)
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalar_one.return_value = 3
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
assert jobs == []
|
||||
result = await repository.count_user_jobs(owner_id)
|
||||
|
||||
assert result == 3
|
||||
session.execute.assert_awaited_once()
|
||||
call_args = session.execute.call_args
|
||||
stmt = call_args[0][0]
|
||||
stmt_str = str(stmt)
|
||||
assert "bootstrap_key" in stmt_str
|
||||
assert "IS NULL" in stmt_str or "is_(None)" in stmt_str.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_id_returns_job(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
sample_job: SimpleNamespace,
|
||||
) -> None:
|
||||
fake_session.set_execute_result(sample_job)
|
||||
async def test_create_sets_bootstrap_key_to_none() -> None:
|
||||
session = AsyncMock()
|
||||
session.add = MagicMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
|
||||
await repository.create(owner_id, data)
|
||||
|
||||
session.add.assert_called_once()
|
||||
call_args = session.add.call_args[0][0]
|
||||
assert call_args.bootstrap_key is None
|
||||
session.flush.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_sets_correct_fields() -> None:
|
||||
session = AsyncMock()
|
||||
session.add = MagicMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
|
||||
await repository.create(owner_id, data)
|
||||
|
||||
call_args = session.add.call_args[0][0]
|
||||
assert call_args.owner_id == owner_id
|
||||
assert call_args.title == data.title
|
||||
assert call_args.schedule_type == data.schedule_type
|
||||
assert call_args.timezone == data.timezone
|
||||
assert call_args.status == data.status
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_returns_updated_job() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
job = await repository.get_by_id(job_id)
|
||||
existing_job = MagicMock()
|
||||
existing_job.schedule_type = ScheduleType.DAILY
|
||||
existing_job.config = {"input_template": "Old"}
|
||||
updated_job = MagicMock()
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalar_one_or_none.return_value = updated_job
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
assert job is not None
|
||||
assert job.title == "Test Job"
|
||||
data = AutomationJobUpdateRequest(title="Updated Title")
|
||||
result = await repository.update(job_id, data)
|
||||
|
||||
assert result is updated_job
|
||||
session.flush.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_id_returns_none_when_not_found(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
fake_session.set_execute_result(None)
|
||||
|
||||
async def test_update_merges_config() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
job = await repository.get_by_id(job_id)
|
||||
existing_job = MagicMock()
|
||||
existing_job.schedule_type = ScheduleType.DAILY
|
||||
existing_job.config = {"input_template": "Old", "enabled_tools": []}
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalar_one_or_none.return_value = existing_job
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
assert job is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_user_jobs_returns_count(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
fake_session.set_execute_result(5)
|
||||
|
||||
owner_id = uuid4()
|
||||
count = await repository.count_user_jobs(owner_id)
|
||||
|
||||
assert count == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_user_jobs_returns_zero_when_none(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
fake_session.set_execute_result(0)
|
||||
|
||||
owner_id = uuid4()
|
||||
count = await repository.count_user_jobs(owner_id)
|
||||
|
||||
assert count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_job(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
from v1.automation_jobs.schemas import AutomationJobCreateRequest
|
||||
from schemas.automation import AutomationJobConfig
|
||||
|
||||
owner_id = uuid4()
|
||||
request = AutomationJobCreateRequest(
|
||||
title="New Job",
|
||||
schedule_type=ScheduleType.DAILY,
|
||||
run_at=time(0, 0),
|
||||
timezone="UTC",
|
||||
status=AutomationJobStatus.ACTIVE,
|
||||
config=AutomationJobConfig(input_template="Test"),
|
||||
data = AutomationJobUpdateRequest(
|
||||
config={"input_template": "New", "context": {"source": "latest_chat"}}
|
||||
)
|
||||
await repository.update(job_id, data)
|
||||
|
||||
job = await repository.create(owner_id, request)
|
||||
|
||||
assert job.title == "New Job"
|
||||
assert job.owner_id == owner_id
|
||||
assert job.created_by == owner_id
|
||||
assert job.bootstrap_key is None
|
||||
assert job.schedule_type == ScheduleType.DAILY
|
||||
assert fake_session.flushed is True
|
||||
assert len(fake_session.added) == 1
|
||||
session.flush.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_soft_delete(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
async def test_update_returns_none_when_job_not_found() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalar_one_or_none.return_value = None
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
data = AutomationJobUpdateRequest(title="Updated Title")
|
||||
result = await repository.update(job_id, data)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_soft_delete_calls_soft_delete_by_id() -> None:
|
||||
session = AsyncMock()
|
||||
session.flush = AsyncMock()
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalar_one_or_none.return_value = None
|
||||
session.execute.return_value = execute_result
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
fake_session.set_execute_result(None)
|
||||
|
||||
await repository.soft_delete(job_id)
|
||||
|
||||
assert fake_session.flushed is True
|
||||
session.flush.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_job_title(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
sample_job: SimpleNamespace,
|
||||
) -> None:
|
||||
from v1.automation_jobs.schemas import AutomationJobUpdateRequest
|
||||
async def test_list_due_jobs_filters_by_active_status() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalars.return_value.all.return_value = []
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
sample_job.title = "Updated Title"
|
||||
fake_session.set_execute_result(sample_job)
|
||||
await repository.list_due_jobs(now_utc=MagicMock(), limit=10)
|
||||
|
||||
request = AutomationJobUpdateRequest(title="Updated Title")
|
||||
job = await repository.update(sample_job.id, request)
|
||||
|
||||
assert job is not None
|
||||
assert job.title == "Updated Title"
|
||||
session.execute.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_job_run_at_recomputes_next_run_at(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
sample_job: SimpleNamespace,
|
||||
) -> None:
|
||||
from v1.automation_jobs.schemas import AutomationJobUpdateRequest
|
||||
async def test_create_stores_run_at_as_timezone_aware() -> None:
|
||||
session = AsyncMock()
|
||||
session.add = MagicMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
|
||||
fake_session.set_execute_result(sample_job)
|
||||
await repository.create(owner_id, data)
|
||||
|
||||
request = AutomationJobUpdateRequest(
|
||||
run_at=time(12, 0),
|
||||
timezone="UTC",
|
||||
call_args = session.add.call_args[0][0]
|
||||
assert call_args.run_at.tzinfo is not None, "run_at should be timezone-aware"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_run_at_with_timezone_none_uses_existing_timezone() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
existing_job = MagicMock()
|
||||
existing_job.schedule_type = ScheduleType.DAILY
|
||||
existing_job.timezone = "America/New_York"
|
||||
existing_job.config = {}
|
||||
existing_job.run_at = None
|
||||
execute_result = MagicMock()
|
||||
execute_result.scalar_one_or_none.return_value = existing_job
|
||||
session.execute.return_value = execute_result
|
||||
|
||||
repository.update_by_id = AsyncMock(return_value=existing_job)
|
||||
|
||||
data = AutomationJobUpdateRequest(run_at=time(14, 30, 0))
|
||||
result = await repository.update(job_id, data)
|
||||
|
||||
assert result is not None
|
||||
update_values = repository.update_by_id.call_args[0][1]
|
||||
assert "run_at" in update_values
|
||||
assert "next_run_at" in update_values
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_schedule_type_recomputes_next_run_at() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
existing_job = MagicMock()
|
||||
existing_job.schedule_type = ScheduleType.DAILY
|
||||
existing_job.timezone = "UTC"
|
||||
existing_job.run_at = datetime(2026, 1, 1, 8, 0, 0, tzinfo=timezone.utc)
|
||||
existing_job.config = {}
|
||||
|
||||
repository.get_by_id = AsyncMock(return_value=existing_job)
|
||||
repository.update_by_id = AsyncMock(return_value=existing_job)
|
||||
|
||||
data = AutomationJobUpdateRequest(schedule_type=ScheduleType.WEEKLY)
|
||||
result = await repository.update(job_id, data)
|
||||
|
||||
assert result is not None
|
||||
update_values = repository.update_by_id.call_args[0][1]
|
||||
assert update_values["schedule_type"] == ScheduleType.WEEKLY
|
||||
assert "run_at" in update_values
|
||||
assert "next_run_at" in update_values
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_serializes_enum_values_to_json() -> None:
|
||||
session = AsyncMock()
|
||||
repository = AutomationJobsRepository(session)
|
||||
job_id = uuid4()
|
||||
existing_job = MagicMock()
|
||||
existing_job.schedule_type = ScheduleType.DAILY
|
||||
existing_job.timezone = "UTC"
|
||||
existing_job.run_at = datetime(2026, 1, 1, 8, 0, 0, tzinfo=timezone.utc)
|
||||
existing_job.config = {"input_template": "Old"}
|
||||
|
||||
repository.get_by_id = AsyncMock(return_value=existing_job)
|
||||
repository.update_by_id = AsyncMock(return_value=existing_job)
|
||||
|
||||
data = AutomationJobUpdateRequest(
|
||||
config={"enabled_tools": [AgentTool.MEMORY_WRITE]},
|
||||
)
|
||||
job = await repository.update(sample_job.id, request)
|
||||
result = await repository.update(job_id, data)
|
||||
|
||||
assert job is not None
|
||||
assert fake_session.flushed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_returns_none_when_job_not_found(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
) -> None:
|
||||
from v1.automation_jobs.schemas import AutomationJobUpdateRequest
|
||||
|
||||
fake_session.set_execute_result(None)
|
||||
|
||||
request = AutomationJobUpdateRequest(title="New Title")
|
||||
job = await repository.update(uuid4(), request)
|
||||
|
||||
assert job is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_with_no_changes_returns_existing_job(
|
||||
repository: AutomationJobsRepository,
|
||||
fake_session: _FakeSession,
|
||||
sample_job: SimpleNamespace,
|
||||
) -> None:
|
||||
from v1.automation_jobs.schemas import AutomationJobUpdateRequest
|
||||
|
||||
fake_session.set_execute_result(sample_job)
|
||||
|
||||
request = AutomationJobUpdateRequest()
|
||||
job = await repository.update(sample_job.id, request)
|
||||
|
||||
assert job is not None
|
||||
assert job.title == "Test Job"
|
||||
assert result is not None
|
||||
update_values = repository.update_by_id.call_args[0][1]
|
||||
enabled_tools = update_values["config"]["enabled_tools"]
|
||||
assert isinstance(enabled_tools[0], str)
|
||||
|
||||
@@ -0,0 +1,246 @@
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from v1.automation_jobs.schemas import (
|
||||
AutomationJobCreateRequest,
|
||||
AutomationJobUpdateRequest,
|
||||
AutomationJobResponse,
|
||||
)
|
||||
from schemas.automation import AgentTool, AutomationJobConfig
|
||||
|
||||
|
||||
class TestIsSystemProperty:
|
||||
def test_is_system_true_when_bootstrap_key_present(self):
|
||||
mock_orm_job = MagicMock()
|
||||
mock_orm_job.id = uuid4()
|
||||
mock_orm_job.owner_id = uuid4()
|
||||
mock_orm_job.bootstrap_key = "memory_extraction"
|
||||
mock_orm_job.title = "Test Job"
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.run_at = datetime.now()
|
||||
mock_orm_job.config = {
|
||||
"input_template": "Hello",
|
||||
"enabled_tools": [],
|
||||
"context": {},
|
||||
}
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.status = "active"
|
||||
mock_orm_job.timezone = "Asia/Shanghai"
|
||||
mock_orm_job.next_run_at = datetime.now()
|
||||
mock_orm_job.last_run_at = None
|
||||
mock_orm_job.created_at = datetime.now()
|
||||
mock_orm_job.updated_at = datetime.now()
|
||||
mock_orm_job.deleted_at = None
|
||||
|
||||
resp = AutomationJobResponse.from_orm(mock_orm_job)
|
||||
assert resp.is_system is True
|
||||
|
||||
def test_is_system_false_when_bootstrap_key_none(self):
|
||||
mock_orm_job = MagicMock()
|
||||
mock_orm_job.id = uuid4()
|
||||
mock_orm_job.owner_id = uuid4()
|
||||
mock_orm_job.bootstrap_key = None
|
||||
mock_orm_job.title = "Test Job"
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.run_at = datetime.now()
|
||||
mock_orm_job.config = {
|
||||
"input_template": "Hello",
|
||||
"enabled_tools": [],
|
||||
"context": {},
|
||||
}
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.status = "active"
|
||||
mock_orm_job.timezone = "Asia/Shanghai"
|
||||
mock_orm_job.next_run_at = datetime.now()
|
||||
mock_orm_job.last_run_at = None
|
||||
mock_orm_job.created_at = datetime.now()
|
||||
mock_orm_job.updated_at = datetime.now()
|
||||
mock_orm_job.deleted_at = None
|
||||
|
||||
resp = AutomationJobResponse.from_orm(mock_orm_job)
|
||||
assert resp.is_system is False
|
||||
|
||||
|
||||
class TestFromOrm:
|
||||
def test_run_at_converted_from_datetime_to_time(self):
|
||||
run_at_datetime = datetime(2024, 6, 15, 14, 30, 0)
|
||||
mock_orm_job = MagicMock()
|
||||
mock_orm_job.id = uuid4()
|
||||
mock_orm_job.owner_id = uuid4()
|
||||
mock_orm_job.bootstrap_key = None
|
||||
mock_orm_job.title = "Test Job"
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.run_at = run_at_datetime
|
||||
mock_orm_job.config = {
|
||||
"input_template": "Hello",
|
||||
"enabled_tools": [],
|
||||
"context": {},
|
||||
}
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.status = "active"
|
||||
mock_orm_job.timezone = "Asia/Shanghai"
|
||||
mock_orm_job.next_run_at = datetime.now()
|
||||
mock_orm_job.last_run_at = None
|
||||
mock_orm_job.created_at = datetime.now()
|
||||
mock_orm_job.updated_at = datetime.now()
|
||||
mock_orm_job.deleted_at = None
|
||||
|
||||
resp = AutomationJobResponse.from_orm(mock_orm_job)
|
||||
assert resp.run_at == run_at_datetime.time()
|
||||
|
||||
def test_config_deserialized(self):
|
||||
config = {
|
||||
"input_template": "Test template",
|
||||
"enabled_tools": [AgentTool.MEMORY_WRITE],
|
||||
"context": {
|
||||
"source": "latest_chat",
|
||||
"window_mode": "day",
|
||||
"window_count": 5,
|
||||
},
|
||||
}
|
||||
mock_orm_job = MagicMock()
|
||||
mock_orm_job.id = uuid4()
|
||||
mock_orm_job.owner_id = uuid4()
|
||||
mock_orm_job.bootstrap_key = None
|
||||
mock_orm_job.title = "Test Job"
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.run_at = datetime.now()
|
||||
mock_orm_job.config = config
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.status = "active"
|
||||
mock_orm_job.timezone = "Asia/Shanghai"
|
||||
mock_orm_job.next_run_at = datetime.now()
|
||||
mock_orm_job.last_run_at = None
|
||||
mock_orm_job.created_at = datetime.now()
|
||||
mock_orm_job.updated_at = datetime.now()
|
||||
mock_orm_job.deleted_at = None
|
||||
|
||||
resp = AutomationJobResponse.from_orm(mock_orm_job)
|
||||
assert resp.config.input_template == "Test template"
|
||||
assert resp.config.enabled_tools == [AgentTool.MEMORY_WRITE]
|
||||
assert resp.config.context.window_count == 5
|
||||
|
||||
def test_is_system_derived_from_bootstrap_key(self):
|
||||
mock_orm_job = MagicMock()
|
||||
mock_orm_job.id = uuid4()
|
||||
mock_orm_job.owner_id = uuid4()
|
||||
mock_orm_job.bootstrap_key = "system_bootstrap"
|
||||
mock_orm_job.title = "Test Job"
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.run_at = datetime.now()
|
||||
mock_orm_job.config = {
|
||||
"input_template": "Hello",
|
||||
"enabled_tools": [],
|
||||
"context": {},
|
||||
}
|
||||
mock_orm_job.schedule_type = "daily"
|
||||
mock_orm_job.status = "active"
|
||||
mock_orm_job.timezone = "UTC"
|
||||
mock_orm_job.next_run_at = datetime.now()
|
||||
mock_orm_job.last_run_at = None
|
||||
mock_orm_job.created_at = datetime.now()
|
||||
mock_orm_job.updated_at = datetime.now()
|
||||
mock_orm_job.deleted_at = None
|
||||
|
||||
resp = AutomationJobResponse.from_orm(mock_orm_job)
|
||||
assert resp.is_system is True
|
||||
assert resp.bootstrap_key == "system_bootstrap"
|
||||
|
||||
|
||||
class TestTimezoneValidation:
|
||||
def test_valid_timezone(self):
|
||||
request = AutomationJobCreateRequest.model_validate(
|
||||
{
|
||||
"title": "Test Job",
|
||||
"schedule_type": "daily",
|
||||
"run_at": "09:00:00",
|
||||
"timezone": "Asia/Shanghai",
|
||||
"config": {
|
||||
"input_template": "Hello",
|
||||
"enabled_tools": [],
|
||||
"context": {
|
||||
"source": "latest_chat",
|
||||
"window_mode": "day",
|
||||
"window_count": 2,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
assert request.timezone == "Asia/Shanghai"
|
||||
|
||||
def test_invalid_timezone(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
AutomationJobCreateRequest.model_validate(
|
||||
{
|
||||
"title": "Test Job",
|
||||
"schedule_type": "daily",
|
||||
"run_at": "09:00:00",
|
||||
"timezone": "Invalid/Timezone",
|
||||
"config": {
|
||||
"input_template": "Hello",
|
||||
"enabled_tools": [],
|
||||
"context": {
|
||||
"source": "latest_chat",
|
||||
"window_mode": "day",
|
||||
"window_count": 2,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
assert "timezone must be a valid IANA timezone" in str(exc_info.value)
|
||||
|
||||
def test_update_valid_timezone(self):
|
||||
request = AutomationJobUpdateRequest.model_validate(
|
||||
{
|
||||
"timezone": "America/New_York",
|
||||
}
|
||||
)
|
||||
assert request.timezone == "America/New_York"
|
||||
|
||||
def test_update_invalid_timezone(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
AutomationJobUpdateRequest.model_validate(
|
||||
{
|
||||
"timezone": "Invalid/Timezone",
|
||||
}
|
||||
)
|
||||
assert "timezone must be a valid IANA timezone" in str(exc_info.value)
|
||||
|
||||
def test_update_none_timezone_allowed(self):
|
||||
request = AutomationJobUpdateRequest.model_validate(
|
||||
{
|
||||
"timezone": None,
|
||||
}
|
||||
)
|
||||
assert request.timezone is None
|
||||
|
||||
|
||||
class TestAutomationJobConfigPatch:
|
||||
def test_all_fields_optional(self):
|
||||
patch = AutomationJobConfig.model_validate({})
|
||||
assert patch.input_template is None
|
||||
assert patch.enabled_tools is None
|
||||
assert patch.context is None
|
||||
|
||||
def test_partial_input_template(self):
|
||||
patch = AutomationJobConfig.model_validate(
|
||||
{
|
||||
"input_template": "Updated template",
|
||||
}
|
||||
)
|
||||
assert patch.input_template == "Updated template"
|
||||
assert patch.enabled_tools is None
|
||||
assert patch.context is None
|
||||
|
||||
def test_extra_fields_forbidden(self):
|
||||
with pytest.raises(ValidationError):
|
||||
AutomationJobConfig.model_validate(
|
||||
{
|
||||
"input_template": "Test",
|
||||
"unknown_field": "value",
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,371 @@
|
||||
from datetime import datetime, time, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from models.automation_jobs import AutomationJobStatus, ScheduleType
|
||||
from v1.automation_jobs.service import (
|
||||
AutomationJobLimitExceeded,
|
||||
AutomationJobNotFound,
|
||||
AutomationJobsService,
|
||||
SystemJobModificationForbidden,
|
||||
)
|
||||
from v1.automation_jobs.schemas import (
|
||||
AutomationJobCreateRequest,
|
||||
AutomationJobUpdateRequest,
|
||||
)
|
||||
from schemas.automation import (
|
||||
AgentTool,
|
||||
AutomationJobConfig,
|
||||
ContextSource,
|
||||
ContextWindowMode,
|
||||
MessageContextConfig,
|
||||
)
|
||||
|
||||
|
||||
def _make_config() -> AutomationJobConfig:
|
||||
return AutomationJobConfig(
|
||||
input_template="Hello",
|
||||
enabled_tools=[AgentTool.MEMORY_WRITE],
|
||||
context=MessageContextConfig(
|
||||
source=ContextSource.LATEST_CHAT,
|
||||
window_mode=ContextWindowMode.DAY,
|
||||
window_count=2,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _make_create_request() -> AutomationJobCreateRequest:
|
||||
return AutomationJobCreateRequest(
|
||||
title="Test Job",
|
||||
schedule_type=ScheduleType.DAILY,
|
||||
run_at=time(9, 0, 0),
|
||||
timezone="Asia/Shanghai",
|
||||
status=AutomationJobStatus.ACTIVE,
|
||||
config=_make_config(),
|
||||
)
|
||||
|
||||
|
||||
def _make_job(
|
||||
owner_id: MagicMock | None = None, bootstrap_key: str | None = None
|
||||
) -> MagicMock:
|
||||
job = MagicMock()
|
||||
job.id = uuid4()
|
||||
job.owner_id = owner_id or uuid4()
|
||||
job.bootstrap_key = bootstrap_key
|
||||
job.title = "Test Job"
|
||||
job.schedule_type = ScheduleType.DAILY
|
||||
job.run_at = datetime(2024, 1, 1, 9, 0, 0, tzinfo=timezone.utc)
|
||||
job.timezone = "Asia/Shanghai"
|
||||
job.status = AutomationJobStatus.ACTIVE
|
||||
job.config = {"input_template": "Hello"}
|
||||
job.next_run_at = datetime(2024, 1, 2, 9, 0, 0, tzinfo=timezone.utc)
|
||||
job.last_run_at = None
|
||||
job.created_at = datetime(2024, 1, 1, 9, 0, 0, tzinfo=timezone.utc)
|
||||
job.updated_at = datetime(2024, 1, 1, 9, 0, 0, tzinfo=timezone.utc)
|
||||
return job
|
||||
|
||||
|
||||
class TestListByOwner:
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_by_owner_returns_jobs(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id)
|
||||
repository.list_by_owner.return_value = [job]
|
||||
|
||||
result = await service.list_by_owner(owner_id)
|
||||
|
||||
assert len(result.items) == 1
|
||||
assert result.items[0].title == job.title
|
||||
repository.list_by_owner.assert_awaited_once_with(owner_id)
|
||||
|
||||
|
||||
class TestGetById:
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_id_returns_job(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id)
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
result = await service.get_by_id(job.id, owner_id)
|
||||
|
||||
assert result.title == job.title
|
||||
repository.get_by_id.assert_awaited_once_with(job.id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_id_raises_not_found_when_job_none(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job_id = uuid4()
|
||||
repository.get_by_id.return_value = None
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.get_by_id(job_id, owner_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_id_raises_not_found_when_owner_mismatch(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
different_owner_id = uuid4()
|
||||
job = _make_job(different_owner_id)
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.get_by_id(job.id, owner_id)
|
||||
|
||||
|
||||
class TestCreate:
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_raises_limit_exceeded(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
repository.count_user_jobs.return_value = 3
|
||||
|
||||
with pytest.raises(AutomationJobLimitExceeded):
|
||||
await service.create(owner_id, data)
|
||||
|
||||
session.execute.assert_awaited_once()
|
||||
session.rollback.assert_awaited_once()
|
||||
repository.count_user_jobs.assert_awaited_once_with(owner_id)
|
||||
repository.create.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_succeeds_when_under_limit(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
job = _make_job(owner_id)
|
||||
repository.count_user_jobs.return_value = 2
|
||||
repository.create.return_value = job
|
||||
|
||||
result = await service.create(owner_id, data)
|
||||
|
||||
assert result.title == job.title
|
||||
session.execute.assert_awaited_once()
|
||||
repository.create.assert_awaited_once_with(owner_id, data)
|
||||
session.commit.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_commits_session(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
job = _make_job(owner_id)
|
||||
repository.count_user_jobs.return_value = 0
|
||||
repository.create.return_value = job
|
||||
|
||||
await service.create(owner_id, data)
|
||||
|
||||
session.execute.assert_awaited_once()
|
||||
session.commit.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_rollbacks_on_sqlalchemy_error(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
data = _make_create_request()
|
||||
repository.count_user_jobs.return_value = 0
|
||||
repository.create.side_effect = SQLAlchemyError("db down")
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await service.create(owner_id, data)
|
||||
|
||||
assert exc.value.status_code == 503
|
||||
session.execute.assert_awaited_once()
|
||||
session.rollback.assert_awaited_once()
|
||||
session.commit.assert_not_awaited()
|
||||
|
||||
|
||||
class TestUpdate:
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_raises_not_found_when_job_none(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job_id = uuid4()
|
||||
repository.get_by_id.return_value = None
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.update(
|
||||
job_id, owner_id, AutomationJobUpdateRequest(title="New")
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_raises_not_found_when_owner_mismatch(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
different_owner_id = uuid4()
|
||||
job = _make_job(different_owner_id)
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.update(
|
||||
job.id, owner_id, AutomationJobUpdateRequest(title="New")
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_raises_system_job_forbidden(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id, bootstrap_key="system-key")
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
with pytest.raises(SystemJobModificationForbidden):
|
||||
await service.update(
|
||||
job.id, owner_id, AutomationJobUpdateRequest(title="New")
|
||||
)
|
||||
|
||||
repository.update.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_succeeds(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id)
|
||||
updated_job = _make_job(owner_id)
|
||||
updated_job.title = "Updated Title"
|
||||
repository.get_by_id.return_value = job
|
||||
repository.update.return_value = updated_job
|
||||
|
||||
result = await service.update(
|
||||
job.id, owner_id, AutomationJobUpdateRequest(title="Updated Title")
|
||||
)
|
||||
|
||||
assert result.title == "Updated Title"
|
||||
repository.update.assert_awaited_once_with(
|
||||
job.id, AutomationJobUpdateRequest(title="Updated Title")
|
||||
)
|
||||
session.commit.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_returns_not_found_when_update_returns_none(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id)
|
||||
repository.get_by_id.return_value = job
|
||||
repository.update.return_value = None
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.update(
|
||||
job.id, owner_id, AutomationJobUpdateRequest(title="New")
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_rollbacks_on_sqlalchemy_error(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id, bootstrap_key=None)
|
||||
repository.get_by_id.return_value = job
|
||||
repository.update.side_effect = SQLAlchemyError("db down")
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await service.update(
|
||||
job.id, owner_id, AutomationJobUpdateRequest(title="New")
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 503
|
||||
session.rollback.assert_awaited_once()
|
||||
|
||||
|
||||
class TestDelete:
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_raises_not_found_when_job_none(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job_id = uuid4()
|
||||
repository.get_by_id.return_value = None
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.delete(job_id, owner_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_raises_not_found_when_owner_mismatch(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
different_owner_id = uuid4()
|
||||
job = _make_job(different_owner_id)
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
with pytest.raises(AutomationJobNotFound):
|
||||
await service.delete(job.id, owner_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_raises_system_job_forbidden(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id, bootstrap_key="system-key")
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
with pytest.raises(SystemJobModificationForbidden):
|
||||
await service.delete(job.id, owner_id)
|
||||
|
||||
repository.soft_delete.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_succeeds(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id)
|
||||
repository.get_by_id.return_value = job
|
||||
|
||||
await service.delete(job.id, owner_id)
|
||||
|
||||
repository.soft_delete.assert_awaited_once_with(job.id)
|
||||
session.commit.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_rollbacks_on_sqlalchemy_error(self) -> None:
|
||||
session = AsyncMock()
|
||||
repository = AsyncMock()
|
||||
service = AutomationJobsService(repository, session)
|
||||
owner_id = uuid4()
|
||||
job = _make_job(owner_id, bootstrap_key=None)
|
||||
repository.get_by_id.return_value = job
|
||||
repository.soft_delete.side_effect = SQLAlchemyError("db down")
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await service.delete(job.id, owner_id)
|
||||
|
||||
assert exc.value.status_code == 503
|
||||
session.rollback.assert_awaited_once()
|
||||
Reference in New Issue
Block a user