247 lines
8.5 KiB
Python
247 lines
8.5 KiB
Python
|
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",
|
||
|
|
}
|
||
|
|
)
|