social-app/backend/tests/unit/v1/automation_jobs/test_repository.py

from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4

import pytest

from models.automation_jobs import AutomationJobStatus, ScheduleType
from schemas.automation import (
    AgentTool,
    AutomationJobConfig,
    ContextSource,
    ContextWindowMode,
    MessageContextConfig,
    ScheduleConfig,
    ScheduleRunAt,
)
from v1.automation_jobs.repository import AutomationJobsRepository
from v1.automation_jobs.schemas import (
    AutomationJobCreateRequest,
    AutomationJobUpdateRequest,
)


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,
        ),
        schedule=ScheduleConfig(
            type=ScheduleType.DAILY,
            run_at=ScheduleRunAt(hour=9, minute=0),
        ),
    )


def _make_create_request() -> AutomationJobCreateRequest:
    return AutomationJobCreateRequest(
        title="Test Job",
        timezone="Asia/Shanghai",
        status=AutomationJobStatus.ACTIVE,
        config=_make_config(),
    )


@pytest.mark.asyncio
async def test_list_by_owner_returns_jobs() -> None:
    session = AsyncMock()
    repository = AutomationJobsRepository(session)
    owner_id = uuid4()
    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

    result = await repository.list_by_owner(owner_id)

    assert result == [job_one, job_two]
    session.execute.assert_awaited_once()


@pytest.mark.asyncio
async def test_count_user_jobs_counts_non_bootstrap_jobs() -> None:
    session = AsyncMock()
    repository = AutomationJobsRepository(session)
    owner_id = uuid4()
    execute_result = MagicMock()
    execute_result.scalar_one.return_value = 3
    session.execute.return_value = execute_result

    result = await repository.count_user_jobs(owner_id)

    assert result == 3


@pytest.mark.asyncio
async def test_create_sets_fields_and_next_run_at() -> 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.owner_id == owner_id
    assert call_args.bootstrap_key is None
    assert call_args.title == data.title
    assert call_args.timezone == data.timezone
    assert call_args.status == data.status
    assert call_args.next_run_at is not None


@pytest.mark.asyncio
async def test_update_returns_none_when_job_not_found() -> None:
    session = AsyncMock()
    repository = AutomationJobsRepository(session)
    job_id = uuid4()

    repository.get_by_id = AsyncMock(return_value=None)

    data = AutomationJobUpdateRequest(title="Updated Title")
    result = await repository.update(job_id, data)

    assert result is None


@pytest.mark.asyncio
async def test_update_merges_config_and_recomputes_next_run() -> None:
    session = AsyncMock()
    repository = AutomationJobsRepository(session)
    job_id = uuid4()
    existing_job = MagicMock()
    existing_job.timezone = "UTC"
    existing_job.config = {
        "input_template": "Old",
        "enabled_tools": ["memory.write"],
        "context": {
            "source": "latest_chat",
            "window_mode": "day",
            "window_count": 2,
        },
        "schedule": {
            "type": "daily",
            "run_at": {"hour": 8, "minute": 0},
        },
    }

    repository.get_by_id = AsyncMock(return_value=existing_job)
    repository.update_by_id = AsyncMock(return_value=existing_job)

    data = AutomationJobUpdateRequest(
        config=AutomationJobConfig(
            enabled_tools=[AgentTool.MEMORY_WRITE, AgentTool.MEMORY_FORGET],
            schedule=ScheduleConfig(
                type=ScheduleType.WEEKLY,
                run_at=ScheduleRunAt(hour=10, minute=30),
                weekdays=[2, 5],
            ),
        ),
    )
    result = await repository.update(job_id, data)

    assert result is not None
    update_values = repository.update_by_id.call_args[0][1]
    assert "config" in update_values
    assert "next_run_at" in update_values
    enabled_tools = update_values["config"]["enabled_tools"]
    assert isinstance(enabled_tools[0], str)


@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()

    await repository.soft_delete(job_id)

    session.flush.assert_awaited_once()


@pytest.mark.asyncio
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

    await repository.list_due_jobs(now_utc=MagicMock(), limit=10)

    session.execute.assert_awaited_once()