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

"""§23 — the standalone workers do real work, and do it idempotently.

The workers are not HTTP requests, so there is no ``get_session`` dependency to
override; they open their own sessions from ``app.workers.base.SessionFactory``.
The autouse fixture below swaps that one name so the workers run against the same
in-memory schema the API fixtures build — which is also the point of the
``app.application.container`` extraction: worker and API compose services
identically, so a worker test is meaningful and not a re-implementation.

What is asserted here is the part that is easy to get wrong and expensive to get
wrong: a task that nobody clicked "run" on still runs, a billing period is not
billed twice, and an abandoned sandbox session is stopped *and* audited.
"""

from __future__ import annotations

from datetime import timedelta

import pytest
import pytest_asyncio
from httpx import AsyncClient

from app.domain.base import utcnow
from app.workers import base as worker_base
from app.workers.ai_runtime import AiRuntimeWorker
from app.workers.billing import BillingWorker
from app.workers.sandbox_gateway import SandboxGatewayWorker
from tests.conftest import auth_header, signup


@pytest_asyncio.fixture(autouse=True)
async def _platform_seed(seeded):
    return seeded


@pytest_asyncio.fixture(autouse=True)
async def _workers_use_the_test_database(monkeypatch, session_factory):
    """Point every worker *and the error recorder* at the in-memory database.

    The recorder opens its own session by design, so patching the worker session
    factory alone would leave it writing failures into the developer's real
    ``agy_web.db``.
    """
    import app.infrastructure.error_log as error_log

    monkeypatch.setattr(worker_base, "SessionFactory", session_factory)
    monkeypatch.setattr(error_log, "SessionFactory", session_factory)
    return session_factory


async def _make_org_with_employee(client: AsyncClient, email: str) -> tuple[str, dict, str, str]:
    """Return ``(org_id, headers, project_id, employee_id)`` for a fresh tenant."""
    token = (await signup(client, email))["access_token"]
    headers = auth_header(token)
    org_id = (
        await client.post("/api/organizations", json={"name": "سازمان کارگر"}, headers=headers)
    ).json()["id"]
    project_id = (
        await client.post(
            f"/api/organizations/{org_id}/projects",
            json={"name": "پروژهٔ کارگر"},
            headers=headers,
        )
    ).json()["id"]
    template = (
        await client.get(f"/api/organizations/{org_id}/agent-templates", headers=headers)
    ).json()[0]
    employee_id = (
        await client.post(
            f"/api/organizations/{org_id}/employees",
            json={
                "name": "کارمند کارگر",
                "role_key": template["role_key"],
                "template_id": template["id"],
            },
            headers=headers,
        )
    ).json()["id"]
    return org_id, headers, project_id, employee_id


# -- ai-runtime -------------------------------------------------------------- #


async def test_ai_runtime_runs_a_task_nobody_clicked(client: AsyncClient):
    """A task created without pressing "run" must still be picked up."""
    org_id, headers, project_id, employee_id = await _make_org_with_employee(
        client, "worker-runtime@example.com"
    )
    task_id = (
        await client.post(
            f"/api/organizations/{org_id}/tasks",
            json={
                "title": "کار بی‌صاحب",
                "project_id": project_id,
                "assigned_employee_id": employee_id,
            },
            headers=headers,
        )
    ).json()["id"]

    # Nothing has run it yet.
    detail = await client.get(
        f"/api/organizations/{org_id}/tasks/{task_id}", headers=headers
    )
    assert detail.json()["status"] in {"new", "triaged"}
    assert detail.json()["runs"] == []

    outcome = await AiRuntimeWorker().drain()

    assert outcome.processed == 1, outcome.summary()
    assert outcome.failed == 0, outcome.notes

    detail = await client.get(
        f"/api/organizations/{org_id}/tasks/{task_id}", headers=headers
    )
    assert detail.json()["status"] == "done"
    runs = detail.json()["runs"]
    assert len(runs) == 1
    assert runs[0]["status"] == "completed"


async def test_ai_runtime_does_not_run_a_task_twice(client: AsyncClient):
    """A closed task must not be re-run — re-running would double-bill."""
    org_id, headers, project_id, employee_id = await _make_org_with_employee(
        client, "worker-runtime-twice@example.com"
    )
    await client.post(
        f"/api/organizations/{org_id}/tasks",
        json={
            "title": "کار یک‌باره",
            "project_id": project_id,
            "assigned_employee_id": employee_id,
        },
        headers=headers,
    )

    first = await AiRuntimeWorker().drain()
    second = await AiRuntimeWorker().drain()

    assert first.processed == 1
    assert second.processed == 0, "a done task is not runnable"
    assert second.touched == 0, second.summary()


# -- billing-worker ---------------------------------------------------------- #


async def test_billing_roll_up_creates_a_draft_invoice(client: AsyncClient):
    """Usage in the period must become exactly one draft invoice."""
    org_id, headers, project_id, employee_id = await _make_org_with_employee(
        client, "worker-billing@example.com"
    )
    await client.post(
        f"/api/organizations/{org_id}/tasks",
        json={
            "title": "کار قابل صورتحساب",
            "project_id": project_id,
            "assigned_employee_id": employee_id,
        },
        headers=headers,
    )
    assert (await AiRuntimeWorker().drain()).processed == 1, "the run must be metered"

    end = utcnow()
    outcome = await BillingWorker().roll_up(
        period_start=end - timedelta(days=30), period_end=end
    )

    assert outcome.processed == 1, outcome.summary()
    assert outcome.failed == 0, outcome.notes

    invoices = await client.get(
        f"/api/organizations/{org_id}/billing/invoices", headers=headers
    )
    assert invoices.status_code == 200, invoices.text
    body = invoices.json()
    assert len(body) == 1, body
    assert body[0]["status"] == "draft"
    assert body[0]["amount"] > 0, "the invoice must carry the ledger's total"


async def test_billing_roll_up_is_idempotent(client: AsyncClient):
    """Two passes over the same period must not bill twice."""
    org_id, headers, project_id, employee_id = await _make_org_with_employee(
        client, "worker-billing-twice@example.com"
    )
    await client.post(
        f"/api/organizations/{org_id}/tasks",
        json={
            "title": "کار یک‌بار صورتحساب",
            "project_id": project_id,
            "assigned_employee_id": employee_id,
        },
        headers=headers,
    )
    await AiRuntimeWorker().drain()

    end = utcnow()
    start = end - timedelta(days=30)
    worker = BillingWorker()

    first = await worker.roll_up(period_start=start, period_end=end)
    second = await worker.roll_up(period_start=start, period_end=end)

    assert first.processed == 1
    assert second.processed == 0, "the period was already invoiced"
    assert second.skipped >= 1

    invoices = (
        await client.get(f"/api/organizations/{org_id}/billing/invoices", headers=headers)
    ).json()
    assert len(invoices) == 1, "exactly one invoice for one period"


async def test_billing_roll_up_skips_an_organization_with_no_usage(client: AsyncClient):
    """No usage means no $0 draft littering the invoice table."""
    org_id, headers, _, _ = await _make_org_with_employee(
        client, "worker-billing-empty@example.com"
    )

    end = utcnow()
    outcome = await BillingWorker().roll_up(
        period_start=end - timedelta(days=30), period_end=end
    )

    assert outcome.processed == 0
    assert outcome.skipped >= 1
    invoices = (
        await client.get(f"/api/organizations/{org_id}/billing/invoices", headers=headers)
    ).json()
    assert invoices == []


@pytest.mark.parametrize(
    "delta",
    [timedelta(0), timedelta(minutes=-5), timedelta(minutes=30)],
)
async def test_billing_roll_up_refuses_a_nonsense_period(delta: timedelta):
    """A period that is empty, backwards or absurdly short is a bug, not a bill."""
    end = utcnow()
    with pytest.raises(ValueError):
        await BillingWorker().roll_up(period_start=end - delta, period_end=end)


# -- sandbox-gateway --------------------------------------------------------- #


async def test_sandbox_gateway_stops_an_idle_session_and_audits_it(client: AsyncClient):
    """An abandoned session must be stopped, and the stop must be auditable."""
    org_id, headers, project_id, _ = await _make_org_with_employee(
        client, "worker-gateway@example.com"
    )
    template_id = (
        await client.get(f"/api/organizations/{org_id}/sandbox/templates", headers=headers)
    ).json()[0]["id"]
    workspace_id = (
        await client.post(
            f"/api/organizations/{org_id}/sandbox/workspaces",
            json={"project_id": project_id, "template_id": template_id},
            headers=headers,
        )
    ).json()["id"]
    session_id = (
        await client.post(
            f"/api/organizations/{org_id}/sandbox/sessions",
            json={"workspace_id": workspace_id},
            headers=headers,
        )
    ).json()["id"]

    # Backdate the session so it looks abandoned. The reaper will not touch a
    # session that was active a moment ago, which is the behaviour we want.
    from app.infrastructure.db.models import SandboxSessionRow

    async with worker_base.SessionFactory() as session:
        row = await session.get(SandboxSessionRow, session_id)
        assert row is not None
        row.last_active_at = utcnow() - timedelta(hours=6)
        await session.commit()

    outcome = await SandboxGatewayWorker().reap(idle_minutes=120)

    assert outcome.processed == 1, outcome.summary()
    assert outcome.failed == 0, outcome.notes

    # The session row itself must now read ``stopped``.
    async with worker_base.SessionFactory() as session:
        row = await session.get(SandboxSessionRow, session_id)
        assert row is not None
        assert row.status == "stopped"
        assert row.stopped_at is not None

    audit = await client.get(f"/api/organizations/{org_id}/audit", headers=headers)
    assert audit.status_code == 200, audit.text
    reaps = [e for e in audit.json() if e["action"] == "sandbox.reap"]
    assert reaps, "the platform stopping a sandbox must be audited"
    assert reaps[0]["target_id"] == session_id


async def test_sandbox_gateway_leaves_a_fresh_session_alone(client: AsyncClient):
    """A session that was just used must not be reaped."""
    org_id, headers, project_id, _ = await _make_org_with_employee(
        client, "worker-gateway-fresh@example.com"
    )
    template_id = (
        await client.get(f"/api/organizations/{org_id}/sandbox/templates", headers=headers)
    ).json()[0]["id"]
    workspace_id = (
        await client.post(
            f"/api/organizations/{org_id}/sandbox/workspaces",
            json={"project_id": project_id, "template_id": template_id},
            headers=headers,
        )
    ).json()["id"]
    await client.post(
        f"/api/organizations/{org_id}/sandbox/sessions",
        json={"workspace_id": workspace_id},
        headers=headers,
    )

    outcome = await SandboxGatewayWorker().reap(idle_minutes=120)

    assert outcome.processed == 0, outcome.summary()
    assert outcome.skipped == 1


def test_gateway_policy_report_names_no_secret():
    """The startup report must describe the envelope without leaking a key."""
    outcome = SandboxGatewayWorker().policy_report()
    assert outcome.notes, "the report must say something"

    lowered = " ".join(outcome.notes).lower()
    for forbidden in ("api_key", "token", "secret", "password", "daytona_api", "e2b_api"):
        assert forbidden not in lowered, f"{forbidden} leaked into the policy report"
    assert "deny_all" in lowered, "the default egress policy must be stated"