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

"""Integration: the legacy one-shot run drain never picks up an engine task.

Engine tasks are created with the coarse status ``new``, exactly like a plain
task. ``AiRuntimeWorker._drain_runs`` reads ``TaskRepository.list_runnable`` and
hands every row to the legacy ``start_run``, which marks the task ``done``. If
that query returned an engine task, a manager order would run once through the
drain and never through ``Scheduler.tick()``, and the task would be billed on the
wrong path. A task that owns a ``task_pipeline_state`` row belongs to the engine,
so the drain has to skip it while still picking up the plain task.
"""

from __future__ import annotations

from app.application.ports.delivery_hooks import NullDeliveryHooks
from app.application.ports.repo_workspace import FakeRepoWorkspace
from app.domain.base import new_id
from app.domain.workforce.pipeline import PipelineKey, Stage, TaskKind, TaskPipelineState
from app.infrastructure.db.models import OrganizationRow, TaskRow, UserRow
from app.infrastructure.repositories.engine import PipelineStateRepository
from app.infrastructure.repositories.workforce import TaskRepository
from app.workers.ai_runtime import AiRuntimeWorker
from tests.engine_fakes import FakeSandbox, RecordingMemory, ScriptedModel, ScriptedReply
from tests.test_engine_flow import _scheduler, _task, _team
from tests.test_engine_flow import world as flow_world

#: The flow tests' organization fixture. Binding it under its own name registers it
#: for this module's tests (a bare import would be flagged as unused and shadowed).
world = flow_world


async def test_the_drain_skips_a_task_the_engine_owns(session_factory, seeded) -> None:
    async with session_factory() as session:
        user = UserRow(email=f"drain-{new_id()[:8]}@example.com", password_hash="x")
        session.add(user)
        await session.flush()
        org = OrganizationRow(
            name="سازمان صف اجرا", slug=f"drain-{user.id[:6]}", owner_user_id=user.id
        )
        session.add(org)
        await session.flush()
        plain = TaskRow(
            organization_id=org.id, title="کار ساده", created_by_user_id=user.id
        )
        engine_owned = TaskRow(
            organization_id=org.id, title="کار موتور", created_by_user_id=user.id
        )
        session.add_all([plain, engine_owned])
        await session.flush()
        await PipelineStateRepository(session).create(
            TaskPipelineState(
                task_id=engine_owned.id,
                organization_id=org.id,
                pipeline=PipelineKey.PLAN_CODE_TEST,
                kind=TaskKind.FEATURE,
                stage=Stage.PENDING_MANAGER_REVIEW,
            )
        )
        await session.commit()

        runnable = await TaskRepository(session).list_runnable()

    assert [task.id for task in runnable] == [plain.id]


async def _stuck_org(session) -> tuple[str, str]:
    """An organization with no employees, so none of its tasks can be staffed."""
    user = UserRow(email=f"stuck-{new_id()[:8]}@example.com", password_hash="x")
    session.add(user)
    await session.flush()
    org = OrganizationRow(
        name="سازمان بدون کارمند", slug=f"stuck-{user.id[:6]}", owner_user_id=user.id
    )
    session.add(org)
    await session.flush()
    return org.id, user.id


async def _stuck_task(session, organization_id: str, user_id: str, title: str) -> None:
    row = TaskRow(organization_id=organization_id, title=title, created_by_user_id=user_id)
    session.add(row)
    await session.flush()
    await PipelineStateRepository(session).create(
        TaskPipelineState(
            task_id=row.id,
            organization_id=organization_id,
            pipeline=PipelineKey.PLAN_CODE_TEST,
            kind=TaskKind.FEATURE,
            stage=Stage.PENDING_MANAGER_REVIEW,
        )
    )
    await session.flush()


async def test_twelve_stuck_tasks_do_not_starve_another_org_in_one_pass(
    world, session_factory, redirect_error_log, monkeypatch
) -> None:
    """The tick reads the oldest states first, so twelve unstaffed tasks come first.

    With a read window of ten they would fill it, and the staffed organization's task
    would never be read. The tick reads a wide window, and the per-organization cap in
    ``Scheduler._select`` still lets the staffed task take its turn.
    """
    async with session_factory() as session:
        stuck_org_id, stuck_user_id = await _stuck_org(session)
        for number in range(12):
            await _stuck_task(session, stuck_org_id, stuck_user_id, f"کار گیر {number}")
        await session.commit()

    scripted = ScriptedModel([ScriptedReply(content='[[ASSIGN employee=""/]]')])
    async with session_factory() as session:
        await _team(session, world)
        staffed = await _task(session, world, stage=Stage.PENDING_MANAGER_REVIEW)
        staffed_id = staffed.id
        await session.commit()

    monkeypatch.setattr(
        "app.workers.ai_runtime.build_scheduler",
        lambda session: _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=FakeRepoWorkspace(),
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        ),
    )
    await AiRuntimeWorker().drain()

    async with session_factory() as session:
        state = await PipelineStateRepository(session).get(staffed_id)

    assert state is not None and state.stage is Stage.PENDING_ANALYSIS
    assert scripted.calls == 1