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

"""Every usage row debits its organization's prepaid wallet once, in the same transaction.

What these tests defend:

* a model call debits its price: the balance falls by that price and exactly one
  ``usage`` history row names the usage row it paid for
* a zero-priced usage row is still recorded, but debits nothing
* a sandbox usage debits the wallet too (every unit goes through the one ledger)
* another organization's wallet is never touched
* a rolled-back session leaves neither the usage row nor the debit
* the live wallet event carries the new balance, so the sidebar badge follows it
"""

from __future__ import annotations

import pytest_asyncio
from sqlalchemy import select

from app.application.ports.audit import UsageRecord
from app.application.ports.realtime import RealtimeEvent, get_publisher, set_publisher
from app.application.services.workforce_service import record_model_usage
from app.domain.ai_control.models import RateCard, RateUnit
from app.domain.base import new_id
from app.domain.billing.models import CreditKind
from app.infrastructure.db.models import OrganizationRow, UsageEventRow, UserRow
from app.infrastructure.repositories.ai_control import RateCardRepository
from app.infrastructure.repositories.audit import UsageLedgerRepository
from app.infrastructure.repositories.billing import CreditWalletRepository
from app.infrastructure.repositories.credit import CreditLedgerRepository
from tests.billing_helpers import fund_org


class _Recorder:
    """Stands in for the realtime bus and keeps every event it is given."""

    def __init__(self) -> None:
        self.events: list[RealtimeEvent] = []

    async def publish(self, event: RealtimeEvent) -> None:
        self.events.append(event)


@pytest_asyncio.fixture
async def recorder():
    """Install a recording publisher for one test and put the previous one back."""
    previous = get_publisher()
    bus = _Recorder()
    set_publisher(bus)
    try:
        yield bus
    finally:
        set_publisher(previous)


@pytest_asyncio.fixture
async def two_orgs(session_factory):
    """Two organizations created directly, so neither starts with a wallet."""
    org_ids: list[str] = []
    async with session_factory() as session:
        for label in ("الف", "ب"):
            user = UserRow(email=f"debit-{new_id()[:8]}@example.com", password_hash="x")
            session.add(user)
            await session.flush()
            org = OrganizationRow(
                name=f"سازمان {label}", slug=f"debit-{new_id()[:8]}", owner_user_id=user.id
            )
            session.add(org)
            await session.flush()
            org_ids.append(org.id)
        await session.commit()
    return session_factory, org_ids


async def _balance(session_factory, organization_id: str) -> float:
    async with session_factory() as session:
        wallet = await CreditWalletRepository(session).get_for_org(organization_id)
    assert wallet is not None
    return wallet.balance


async def _usage_ids(session_factory, organization_id: str) -> list[str]:
    async with session_factory() as session:
        result = await session.execute(
            select(UsageEventRow.id).where(UsageEventRow.organization_id == organization_id)
        )
        return list(result.scalars())


async def test_a_model_call_debits_its_price_and_names_its_usage_row(two_orgs, recorder):
    session_factory, (org_a, _org_b) = two_orgs
    await fund_org(session_factory, org_a, 100.0)
    async with session_factory() as session:
        await RateCardRepository(session).upsert(
            RateCard(
                name="توکن ورودی",
                unit=RateUnit.MODEL_INPUT_TOKENS,
                customer_price=2.5,
                internal_cost_estimate=1.0,
            )
        )
        await session.commit()

    async with session_factory() as session:
        await record_model_usage(
            ledger=UsageLedgerRepository(session),
            rate_cards=RateCardRepository(session),
            organization_id=org_a,
            unit=RateUnit.MODEL_INPUT_TOKENS,
            quantity=4.0,
        )
        await session.commit()

    # Four units at 2.5 each: the wallet falls by 10.
    assert await _balance(session_factory, org_a) == 90.0
    usage_ids = await _usage_ids(session_factory, org_a)
    assert len(usage_ids) == 1
    async with session_factory() as session:
        debits = await CreditLedgerRepository(session).history(
            org_a, kinds=(CreditKind.USAGE,)
        )
    assert len(debits) == 1
    assert debits[0].amount == -10.0
    assert debits[0].balance_after == 90.0
    assert debits[0].usage_event_id == usage_ids[0]

    wallet_events = [e for e in recorder.events if e.type == "billing.wallet"]
    assert len(wallet_events) == 1
    assert wallet_events[0].data["balance"] == 90.0


async def test_a_zero_priced_usage_row_is_recorded_but_debits_nothing(two_orgs):
    session_factory, (org_a, _org_b) = two_orgs
    await fund_org(session_factory, org_a, 100.0)
    async with session_factory() as session:
        balance_after = await UsageLedgerRepository(session).record(
            usage=UsageRecord(
                organization_id=org_a,
                unit=RateUnit.TOOL_CALL,
                quantity=1.0,
                total_price=0.0,
            )
        )
        await session.commit()

    assert balance_after is None
    assert len(await _usage_ids(session_factory, org_a)) == 1
    assert await _balance(session_factory, org_a) == 100.0
    async with session_factory() as session:
        debits = await CreditLedgerRepository(session).history(
            org_a, kinds=(CreditKind.USAGE,)
        )
    assert debits == []


async def test_a_sandbox_usage_debits_the_wallet_too(two_orgs):
    session_factory, (org_a, _org_b) = two_orgs
    await fund_org(session_factory, org_a, 20.0)
    async with session_factory() as session:
        balance_after = await UsageLedgerRepository(session).record(
            usage=UsageRecord(
                organization_id=org_a,
                unit=RateUnit.SANDBOX_CPU_SECOND,
                quantity=10.0,
                unit_price=0.1,
                total_price=1.0,
            )
        )
        await session.commit()

    assert balance_after == 19.0
    assert await _balance(session_factory, org_a) == 19.0


async def test_another_organizations_wallet_is_untouched(two_orgs):
    session_factory, (org_a, org_b) = two_orgs
    await fund_org(session_factory, org_a, 50.0)
    await fund_org(session_factory, org_b, 50.0)
    async with session_factory() as session:
        await UsageLedgerRepository(session).record(
            usage=UsageRecord(
                organization_id=org_a,
                unit=RateUnit.TOOL_CALL,
                quantity=1.0,
                total_price=5.0,
            )
        )
        await session.commit()

    assert await _balance(session_factory, org_a) == 45.0
    assert await _balance(session_factory, org_b) == 50.0
    async with session_factory() as session:
        b_debits = await CreditLedgerRepository(session).history(
            org_b, kinds=(CreditKind.USAGE,)
        )
    assert b_debits == []


async def test_a_rolled_back_session_leaves_neither_the_usage_row_nor_the_debit(two_orgs):
    session_factory, (org_a, _org_b) = two_orgs
    await fund_org(session_factory, org_a, 100.0)
    async with session_factory() as session:
        await UsageLedgerRepository(session).record(
            usage=UsageRecord(
                organization_id=org_a,
                unit=RateUnit.TOOL_CALL,
                quantity=1.0,
                total_price=7.0,
            )
        )
        await session.rollback()

    assert await _usage_ids(session_factory, org_a) == []
    assert await _balance(session_factory, org_a) == 100.0
    async with session_factory() as session:
        debits = await CreditLedgerRepository(session).history(
            org_a, kinds=(CreditKind.USAGE,)
        )
    assert debits == []