smart-org/apps/api/tests/test_engine_control.py

"""The person's controls on an engine task: pause, resume, cancel and answer.

``Scheduler.user_action`` is the one place these moves are made for a person.
What it defends:

* pause moves a running task to ``paused``; resume brings a paused task back to
  triage, and is refused while the wallet is empty (the task stays paused)
* cancel ends the task and releases its worktree
* an answer is accepted only from ``pending_user_decision``; it reaches the task
  description and the thread, and the task goes back to triage
* a task of another organization is a 404, never a conflict (tenant isolation)
"""

# ruff: noqa: F811 - the `world` fixture is imported from test_engine_flow; the test parameters
# named `world` are that fixture, which ruff would otherwise read as a redefinition.

from __future__ import annotations

import pytest

from app.application.engine.scheduler import BUDGET_PAUSED_MESSAGE
from app.application.ports.delivery_hooks import NullDeliveryHooks
from app.application.ports.repo_workspace import FakeRepoWorkspace
from app.core.errors import ConflictError, NotFoundError, ValidationError
from app.domain.base import new_id
from app.domain.workforce.models import MessageRole, TaskStatus, Thread
from app.domain.workforce.pipeline import Stage
from app.infrastructure.db.models import OrganizationRow, TaskRow
from app.infrastructure.repositories.engine import PipelineStateRepository
from app.infrastructure.repositories.workforce import (
    MessageRepository,
    TaskRepository,
    ThreadRepository,
)
from tests.engine_fakes import FakeSandbox, RecordingMemory, ScriptedModel
from tests.test_engine_flow import (
    _calls,
    _FundedWallet,
    _scheduler,
    _task,
    _Wallets,
    _ZeroWallet,
    world,  # noqa: F401 - pytest injects this fixture into the tests below
)


def _scheduler_for(session, org, *, workspace=None, wallets=None):
    return _scheduler(
        session,
        org,
        scripted=ScriptedModel([]),
        sandbox=FakeSandbox(),
        workspace=workspace if workspace is not None else FakeRepoWorkspace(),
        memory=RecordingMemory(),
        delivery=NullDeliveryHooks(),
        wallets=wallets,
    )


async def test_pause_moves_a_running_task_to_paused(world, session_factory):
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        await session.commit()
        scheduler = _scheduler_for(session, world)
        state = await scheduler.user_action(
            organization_id=world.organization_id,
            task_id=row.id,
            action="pause",
            user_id=world.user_id,
        )
        task = await TaskRepository(session).get(row.id)

    assert state.stage is Stage.PAUSED
    assert task is not None and task.status is TaskStatus.WAITING_APPROVAL


async def test_resume_brings_a_paused_task_back_to_triage_when_funded(world, session_factory):
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PAUSED)
        await session.commit()
        scheduler = _scheduler_for(session, world, wallets=_Wallets(_FundedWallet()))
        state = await scheduler.user_action(
            organization_id=world.organization_id,
            task_id=row.id,
            action="resume",
            user_id=world.user_id,
        )

    assert state.stage is Stage.PENDING_MANAGER_REVIEW


async def test_resume_with_a_zero_wallet_is_refused_and_the_task_stays_paused(
    world, session_factory
):
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PAUSED)
        await session.commit()
        scheduler = _scheduler_for(session, world, wallets=_Wallets(_ZeroWallet()))
        with pytest.raises(ConflictError) as refused:
            await scheduler.user_action(
                organization_id=world.organization_id,
                task_id=row.id,
                action="resume",
                user_id=world.user_id,
            )
        state = await PipelineStateRepository(session).get(row.id)

    assert refused.value.message == BUDGET_PAUSED_MESSAGE
    assert state is not None and state.stage is Stage.PAUSED


async def test_resume_is_refused_for_a_task_that_is_not_paused(world, session_factory):
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        await session.commit()
        scheduler = _scheduler_for(session, world, wallets=_Wallets(_FundedWallet()))
        with pytest.raises(ConflictError, match="فقط کار"):
            await scheduler.user_action(
                organization_id=world.organization_id,
                task_id=row.id,
                action="resume",
                user_id=world.user_id,
            )
        state = await PipelineStateRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION


async def test_cancel_ends_the_task_and_releases_its_worktree(world, session_factory):
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        await session.commit()
        scheduler = _scheduler_for(session, world, workspace=workspace)
        state = await scheduler.user_action(
            organization_id=world.organization_id,
            task_id=row.id,
            action="cancel",
            user_id=world.user_id,
        )
        task = await TaskRepository(session).get(row.id)

    assert state.stage is Stage.CANCELLED
    assert task is not None and task.status is TaskStatus.CANCELLED
    assert _calls(workspace, "remove_worktree")


async def test_an_answer_reaches_the_description_and_the_thread_then_triage(
    world, session_factory
):
    async with session_factory() as session:
        thread = await ThreadRepository(session).create(
            Thread(
                organization_id=world.organization_id,
                title="گفتگو",
                project_id=world.project_id,
            )
        )
        row = await _task(session, world, stage=Stage.PENDING_USER_DECISION)
        row.thread_id = thread.id
        await session.commit()
        scheduler = _scheduler_for(session, world)
        state = await scheduler.user_action(
            organization_id=world.organization_id,
            task_id=row.id,
            action="answer",
            answer="  نسخهٔ دوم را انجام بده  ",
            user_id=world.user_id,
        )
        task = await TaskRepository(session).get(row.id)
        messages = await MessageRepository(session).list_for_thread(thread.id)

    assert state.stage is Stage.PENDING_MANAGER_REVIEW
    assert task is not None
    assert task.description.endswith("پاسخ کاربر: نسخهٔ دوم را انجام بده")
    assert [(m.role, m.content) for m in messages] == [
        (MessageRole.USER, "نسخهٔ دوم را انجام بده")
    ]


async def test_an_answer_is_refused_unless_the_task_waits_for_a_decision(
    world, session_factory
):
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        await session.commit()
        scheduler = _scheduler_for(session, world)
        with pytest.raises(ConflictError):
            await scheduler.user_action(
                organization_id=world.organization_id,
                task_id=row.id,
                action="answer",
                answer="پاسخ",
                user_id=world.user_id,
            )
        state = await PipelineStateRepository(session).get(row.id)
        task = await TaskRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION
    assert task is not None and "پاسخ کاربر" not in task.description


async def test_a_blank_answer_is_refused(world, session_factory):
    async with session_factory() as session:
        row = await _task(session, world, stage=Stage.PENDING_USER_DECISION)
        await session.commit()
        scheduler = _scheduler_for(session, world)
        with pytest.raises(ValidationError):
            await scheduler.user_action(
                organization_id=world.organization_id,
                task_id=row.id,
                action="answer",
                answer="   ",
                user_id=world.user_id,
            )
        state = await PipelineStateRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.PENDING_USER_DECISION


async def test_a_legacy_task_without_an_engine_state_is_a_conflict(world, session_factory):
    async with session_factory() as session:
        row = TaskRow(
            organization_id=world.organization_id,
            title="کار قدیمی",
            project_id=world.project_id,
            created_by_user_id=world.user_id,
        )
        session.add(row)
        await session.commit()
        scheduler = _scheduler_for(session, world)
        with pytest.raises(ConflictError, match="موتور خودکار"):
            await scheduler.user_action(
                organization_id=world.organization_id,
                task_id=row.id,
                action="pause",
                user_id=world.user_id,
            )


async def test_another_organization_gets_not_found_for_every_control(world, session_factory):
    async with session_factory() as session:
        other = OrganizationRow(
            name="سازمان دیگر",
            slug=f"other-{new_id()[:6]}",
            owner_user_id=world.user_id,
        )
        session.add(other)
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        await session.commit()
        scheduler = _scheduler_for(session, world)
        for action in ("pause", "resume", "cancel", "answer"):
            with pytest.raises(NotFoundError):
                await scheduler.user_action(
                    organization_id=other.id,
                    task_id=row.id,
                    action=action,
                    answer="پاسخ",
                    user_id=world.user_id,
                )
        state = await PipelineStateRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION