roboco/tests/unit/services/test_project_month_spend_usd_db.py

"""TaskService.project_month_spend_usd against a real Postgres DB.

The join (agent_spawn_sessions.task_id, a plain String(36), to tasks.id via a
cast) and the month-boundary filter are exactly the kind of thing a mocked
unit test can't prove actually executes as real SQL. This also pins the
open-session live-token pricing fix: a still-open session's
estimated_cost_usd is null until close, so it must be priced from its token
columns via calculate_cost, not silently summed as $0.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING
from uuid import UUID, uuid4

import pytest
from roboco.billing.pricing import calculate_cost
from roboco.db.tables import AgentSpawnSessionTable, AgentTable, ProjectTable, TaskTable
from roboco.models.base import (
    AgentRole,
    AgentStatus,
    Complexity,
    TaskNature,
    TaskStatus,
    TaskType,
    Team,
)
from roboco.services.task import TaskService

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

_MODEL = "claude-sonnet-5"


async def _seed_project(session: AsyncSession) -> tuple[UUID, UUID]:
    """Seed a system agent + project. Returns ``(project_id, system_agent_id)``
    — the latter doubles as the task-row FK target below."""
    system_agent = AgentTable(
        id=uuid4(),
        name="System",
        slug=f"system-{uuid4().hex[:8]}",
        role=AgentRole.SYSTEM,
        team=None,
        status=AgentStatus.ACTIVE,
        model_config={},
        system_prompt="system",
        capabilities=[],
        permissions={},
        metrics={},
    )
    session.add(system_agent)
    await session.flush()

    project = ProjectTable(
        id=uuid4(),
        name="Budget Spend Test Project",
        slug=f"budget-spend-{uuid4().hex[:8]}",
        git_url="https://github.com/example/budget-spend.git",
        default_branch="main",
        protected_branches=["main"],
        assigned_cell=Team.BACKEND,
        created_by=system_agent.id,
        is_active=True,
    )
    session.add(project)
    await session.flush()
    return UUID(str(project.id)), UUID(str(system_agent.id))


async def _seed_task(session: AsyncSession, project_id: UUID, created_by: UUID) -> UUID:
    task = TaskTable(
        id=uuid4(),
        title="Budget spend fixture task",
        description="A description long enough to satisfy any length floor.",
        acceptance_criteria=["it exists"],
        status=TaskStatus.IN_PROGRESS,
        priority=2,
        task_type=TaskType.CODE,
        nature=TaskNature.TECHNICAL,
        estimated_complexity=Complexity.MEDIUM,
        created_by=created_by,
        team=Team.BACKEND,
        project_id=project_id,
    )
    session.add(task)
    await session.flush()
    return UUID(str(task.id))


def _spawn_session(
    task_id: UUID,
    *,
    started_at: datetime,
    estimated_cost_usd: float | None,
    ended_at: datetime | None,
    tokens: tuple[int, int] = (0, 0),
) -> AgentSpawnSessionTable:
    tokens_input, tokens_output = tokens
    return AgentSpawnSessionTable(
        id=uuid4(),
        agent_slug="be-dev-1",
        team="backend",
        role="developer",
        model=_MODEL,
        task_id=str(task_id),
        started_at=started_at,
        ended_at=ended_at,
        estimated_cost_usd=estimated_cost_usd,
        tokens_input=tokens_input,
        tokens_output=tokens_output,
    )


@pytest.mark.asyncio
async def test_sums_closed_and_prices_open_session(db_session: AsyncSession) -> None:
    """A closed session's estimated_cost_usd + an open session's live-token
    price (calculate_cost) — not the open session silently counted as $0."""
    project_id, agent_id = await _seed_project(db_session)
    task_id = await _seed_task(db_session, project_id, agent_id)
    now = datetime.now(UTC)

    db_session.add(
        _spawn_session(
            task_id,
            started_at=now - timedelta(hours=2),
            estimated_cost_usd=2.5,
            ended_at=now - timedelta(hours=1),
        )
    )
    db_session.add(
        _spawn_session(
            task_id,
            started_at=now - timedelta(minutes=30),
            estimated_cost_usd=None,
            ended_at=None,
            tokens=(100_000, 50_000),
        )
    )
    await db_session.flush()

    expected_open_cost = calculate_cost(
        model=_MODEL, tokens_input=100_000, tokens_output=50_000
    )
    assert expected_open_cost > 0, (
        "fixture model must be priced for this to be a real test"
    )

    svc = TaskService(db_session)
    total = await svc.project_month_spend_usd(project_id)
    assert total == pytest.approx(2.5 + expected_open_cost)


@pytest.mark.asyncio
async def test_excludes_last_months_session(db_session: AsyncSession) -> None:
    """A session that started before this calendar month's boundary must not
    count, even though its cost is closed and non-zero."""
    project_id, agent_id = await _seed_project(db_session)
    task_id = await _seed_task(db_session, project_id, agent_id)
    month_start = datetime.now(UTC).replace(
        day=1, hour=0, minute=0, second=0, microsecond=0
    )
    last_month = month_start - timedelta(days=1)

    db_session.add(
        _spawn_session(
            task_id,
            started_at=last_month,
            estimated_cost_usd=50.0,
            ended_at=last_month + timedelta(hours=1),
        )
    )
    await db_session.flush()

    svc = TaskService(db_session)
    total = await svc.project_month_spend_usd(project_id)
    assert total == 0.0


@pytest.mark.asyncio
async def test_join_excludes_other_projects_tasks(db_session: AsyncSession) -> None:
    """A session on ANOTHER project's task must never bleed into this
    project's sum — proves the join filters by project_id, not just presence
    in agent_spawn_sessions."""
    project_id, agent_id = await _seed_project(db_session)
    other_project_id, other_agent_id = await _seed_project(db_session)
    my_task_id = await _seed_task(db_session, project_id, agent_id)
    other_task_id = await _seed_task(db_session, other_project_id, other_agent_id)
    now = datetime.now(UTC)

    db_session.add(
        _spawn_session(
            my_task_id,
            started_at=now - timedelta(hours=1),
            estimated_cost_usd=1.0,
            ended_at=now,
        )
    )
    db_session.add(
        _spawn_session(
            other_task_id,
            started_at=now - timedelta(hours=1),
            estimated_cost_usd=999.0,
            ended_at=now,
        )
    )
    await db_session.flush()

    svc = TaskService(db_session)
    total = await svc.project_month_spend_usd(project_id)
    assert total == pytest.approx(1.0)