bambuddy/backend/tests/integration/test_scheduler_budget_reservation.py

"""Budget-reservation lifecycle through the unified queue scheduler."""

from contextlib import ExitStack
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import HTTPException
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

import backend.app.models  # noqa: F401 - populate Base.metadata
import backend.app.services.print_scheduler as scheduler_module
from backend.app.core.database import Base
from backend.app.models.archive import PrintArchive
from backend.app.models.finance import BudgetReservation, CostCenter
from backend.app.models.print_queue import PrintQueueItem
from backend.app.models.printer import Printer
from backend.app.models.settings import Settings
from backend.app.models.user import User
from backend.app.services.finance_budget import validate_print_budget
from backend.app.services.print_scheduler import PrintScheduler

pytestmark = pytest.mark.integration


@pytest.fixture
async def billing_dispatch_case(tmp_path):
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    session_maker = async_sessionmaker(engine, expire_on_commit=False)

    base_dir = tmp_path / "billing-dispatch"
    archive_rel = Path("archives") / "job.3mf"
    archive_abs = base_dir / archive_rel
    archive_abs.parent.mkdir(parents=True)
    archive_abs.write_bytes(b"archive payload")

    async with session_maker() as db:
        db.add(Settings(key="billing_enabled", value="true"))
        user = User(username="scheduler-budget-admin", role="admin", is_active=True)
        cost_center = CostCenter(name="Scheduler Budget", is_active=True, monthly_budget=10.0)
        printer = Printer(
            name="Budget Printer",
            serial_number="BUDGET-SERIAL",
            ip_address="127.0.0.1",
            access_code="access-code",
            model="X1C",
        )
        db.add_all([user, cost_center, printer])
        await db.flush()

        archive = PrintArchive(
            printer_id=printer.id,
            filename="job.3mf",
            file_path=str(archive_rel),
            file_size=archive_abs.stat().st_size,
            status="completed",
            cost=4.0,
            created_by_id=user.id,
            cost_center_id=cost_center.id,
        )
        db.add(archive)
        await db.flush()

        item = PrintQueueItem(
            printer_id=printer.id,
            archive_id=archive.id,
            cost_center_id=cost_center.id,
            estimated_cost=4.0,
            created_by_id=user.id,
            status="pending",
        )
        db.add(item)
        await db.commit()

        ids = SimpleNamespace(
            user_id=user.id,
            cost_center_id=cost_center.id,
            printer_id=printer.id,
            archive_id=archive.id,
            item_id=item.id,
        )

    try:
        yield SimpleNamespace(session_maker=session_maker, base_dir=base_dir, ids=ids)
    finally:
        await engine.dispose()


async def _dispatch(ctx, *, uploaded: bool = True, cancel_during_upload: bool = False):
    scheduler = PrintScheduler()
    start_print = MagicMock(return_value=True)

    async def upload(*_args, **_kwargs):
        if cancel_during_upload:
            async with ctx.session_maker() as other_db:
                item = await other_db.get(PrintQueueItem, ctx.ids.item_id)
                item.status = "cancelled"
                await other_db.commit()
        return uploaded

    patches = [
        patch.object(scheduler_module, "async_session", ctx.session_maker),
        patch.object(scheduler_module.settings, "base_dir", ctx.base_dir),
        patch("backend.app.services.print_scheduler.printer_manager.is_connected", MagicMock(return_value=True)),
        patch("backend.app.services.print_scheduler.printer_manager.get_status", MagicMock(return_value=None)),
        patch("backend.app.services.print_scheduler.printer_manager.start_print", start_print),
        patch("backend.app.services.print_scheduler.printer_manager.set_awaiting_plate_clear", MagicMock()),
        patch(
            "backend.app.services.print_scheduler.get_ftp_retry_settings",
            AsyncMock(return_value=(False, 0, 0, 1.0)),
        ),
        patch("backend.app.services.print_scheduler.delete_file_async", AsyncMock(return_value=True)),
        patch("backend.app.services.print_scheduler.upload_file_async", upload),
        patch("backend.app.services.print_scheduler.cache_3mf_download", MagicMock()),
        patch("backend.app.services.print_scheduler.spawn_background_task", MagicMock()),
        patch("backend.app.services.notification_service.notification_service.on_queue_job_started", AsyncMock()),
        patch("backend.app.services.notification_service.notification_service.on_queue_job_failed", AsyncMock()),
        patch("backend.app.services.mqtt_relay.mqtt_relay.on_queue_job_started", AsyncMock()),
        patch.object(scheduler, "_propagate_owner_to_printer_manager", AsyncMock()),
        patch.object(scheduler, "_power_off_if_needed", AsyncMock()),
        patch.object(scheduler, "_preheat_and_soak", AsyncMock()),
    ]
    with ExitStack() as stack:
        for patcher in patches:
            stack.enter_context(patcher)
        await scheduler._dispatch_one(ctx.ids.item_id)

    return start_print


async def _reservation(ctx):
    async with ctx.session_maker() as db:
        return await db.scalar(
            select(BudgetReservation).where(
                BudgetReservation.source_type == "print_queue",
                BudgetReservation.source_id == ctx.ids.item_id,
            )
        )


@pytest.mark.asyncio
async def test_successful_scheduler_dispatch_keeps_one_active_reservation(billing_dispatch_case):
    start_print = await _dispatch(billing_dispatch_case)

    reservation = await _reservation(billing_dispatch_case)
    assert reservation is not None
    assert reservation.status == "active"
    assert reservation.amount == 4.0
    assert reservation.print_archive_id == billing_dispatch_case.ids.archive_id
    start_print.assert_called_once()

    # The printing queue row and its persisted reservation represent the same
    # €4 hold. A second €6 job must fit exactly; €6.01 must not.
    async with billing_dispatch_case.session_maker() as db:
        user = await db.get(User, billing_dispatch_case.ids.user_id)
        second = PrintQueueItem(
            printer_id=billing_dispatch_case.ids.printer_id,
            archive_id=billing_dispatch_case.ids.archive_id,
            cost_center_id=billing_dispatch_case.ids.cost_center_id,
            estimated_cost=6.0,
            created_by_id=user.id,
            status="pending",
        )
        db.add(second)
        await db.commit()
        await validate_print_budget(
            db,
            cost_center_id=second.cost_center_id,
            estimated_cost=6.0,
            current_user=user,
            exclude_queue_item_id=second.id,
        )
        with pytest.raises(HTTPException, match="exceeds available"):
            await validate_print_budget(
                db,
                cost_center_id=second.cost_center_id,
                estimated_cost=6.01,
                current_user=user,
                exclude_queue_item_id=second.id,
            )


@pytest.mark.asyncio
async def test_upload_failure_releases_scheduler_reservation(billing_dispatch_case):
    start_print = await _dispatch(billing_dispatch_case, uploaded=False)

    reservation = await _reservation(billing_dispatch_case)
    assert reservation is not None
    assert reservation.status == "released"
    assert reservation.released_at is not None
    start_print.assert_not_called()


@pytest.mark.asyncio
async def test_retried_dispatch_reuses_active_reservation(billing_dispatch_case):
    await _dispatch(billing_dispatch_case)

    # Simulate startup recovery after the process stopped with a persisted
    # reservation and the queue row was made dispatchable again.
    async with billing_dispatch_case.session_maker() as db:
        item = await db.get(PrintQueueItem, billing_dispatch_case.ids.item_id)
        item.status = "pending"
        item.started_at = None
        item.dispatching_at = None
        await db.commit()

    await _dispatch(billing_dispatch_case)

    async with billing_dispatch_case.session_maker() as db:
        reservations = (
            (
                await db.execute(
                    select(BudgetReservation).where(
                        BudgetReservation.source_type == "print_queue",
                        BudgetReservation.source_id == billing_dispatch_case.ids.item_id,
                    )
                )
            )
            .scalars()
            .all()
        )
    assert len(reservations) == 1
    assert reservations[0].status == "active"


@pytest.mark.asyncio
async def test_cancel_during_upload_releases_scheduler_reservation(billing_dispatch_case):
    start_print = await _dispatch(billing_dispatch_case, cancel_during_upload=True)

    reservation = await _reservation(billing_dispatch_case)
    assert reservation is not None
    assert reservation.status == "released"
    assert reservation.released_at is not None
    start_print.assert_not_called()

    async with billing_dispatch_case.session_maker() as db:
        item = await db.get(PrintQueueItem, billing_dispatch_case.ids.item_id)
        active_count = await db.scalar(
            select(func.count()).select_from(BudgetReservation).where(BudgetReservation.status == "active")
        )
    assert item.status == "cancelled"
    assert active_count == 0