fin-tracker/backend/tests/factories.py

"""Row factories: insert core data directly, without going through a source sync.

Everything returns the new id. Amounts accept str/int/Decimal and are normalised to Decimal,
so tests can write `outcome="1234.56"` and still exercise the NUMERIC(24,10) path.
"""

from __future__ import annotations

from datetime import UTC, date, datetime
from decimal import Decimal
from itertools import count
from typing import Any

from fintracker.db import get_sessionmaker
from fintracker.models import (
    Account,
    AccountKind,
    AccountRole,
    AssetClass,
    CashTxn,
    CashTxnTag,
    Category,
    Event,
    EventKind,
    EventStatus,
    Instrument,
    PriceDaily,
    RawCbrRate,
    Rule,
    RuleKind,
    RuleMatchType,
    Trip,
)

_seq = count(1)


def month_back(months: int) -> date:
    """First day of the month `months` before the current one (in the deployment tz)."""
    from fintracker.analytics import today_local

    t = today_local()
    total = (t.year * 12 + t.month - 1) - months
    return date(total // 12, total % 12 + 1, 1)


Amount = str | int | float | Decimal


def money(value: Amount | None) -> Decimal | None:
    if value is None:
        return None
    return Decimal(str(value))


async def _add(obj: Any) -> Any:
    async with get_sessionmaker()() as session:
        session.add(obj)
        await session.commit()
        await session.refresh(obj)
        return obj


async def make_account(
    *,
    name: str = "Карта",
    currency: str = "RUB",
    role: AccountRole = AccountRole.liquid,
    kind: AccountKind = AccountKind.zm_card,
    balance: Amount | None = 0,
    balance_as_of: datetime | None = None,
    include_in_net_worth: bool = True,
    archived: bool = False,
    disabled: bool = False,
    mirror_of_account_id: int | None = None,
    source: str = "zenmoney",
    source_id: str | None = None,
) -> int:
    acc = await _add(
        Account(
            kind=kind,
            source=source,
            source_id=source_id or f"acc-{next(_seq)}",
            name=name,
            currency=currency,
            role=role,
            balance=money(balance),
            balance_as_of=balance_as_of or datetime.now(UTC),
            include_in_net_worth=include_in_net_worth,
            archived=archived,
            disabled=disabled,
            mirror_of_account_id=mirror_of_account_id,
        )
    )
    return acc.id


async def make_category(
    name: str, *, parent_id: int | None = None, source_id: str | None = None
) -> int:
    cat = await _add(
        Category(
            source="zenmoney",
            source_id=source_id or f"cat-{next(_seq)}",
            name=name,
            parent_id=parent_id,
        )
    )
    return cat.id


async def make_txn(
    d: date,
    *,
    income: Amount = 0,
    income_account_id: int | None = None,
    income_currency: str | None = None,
    outcome: Amount = 0,
    outcome_account_id: int | None = None,
    outcome_currency: str | None = None,
    payee: str | None = None,
    comment: str | None = None,
    mcc: int | None = None,
    hold: bool = False,
    deleted: bool = False,
    primary_category_id: int | None = None,
    tag_ids: list[int] | None = None,
    source_id: str | None = None,
) -> int:
    income_d = money(income) or Decimal(0)
    outcome_d = money(outcome) or Decimal(0)
    txn = await _add(
        CashTxn(
            source="zenmoney",
            source_id=source_id or f"txn-{next(_seq)}",
            ts=datetime.combine(d, datetime.min.time(), tzinfo=UTC),
            date=d,
            income=income_d,
            income_account_id=income_account_id,
            income_currency=income_currency or ("RUB" if income_d else None),
            outcome=outcome_d,
            outcome_account_id=outcome_account_id,
            outcome_currency=outcome_currency or ("RUB" if outcome_d else None),
            payee=payee,
            comment=comment,
            mcc=mcc,
            hold=hold,
            deleted=deleted,
            primary_category_id=primary_category_id,
        )
    )
    tags = (
        tag_ids if tag_ids is not None else ([primary_category_id] if primary_category_id else [])
    )
    if tags:
        async with get_sessionmaker()() as session:
            for ord_, category_id in enumerate(tags):
                session.add(CashTxnTag(txn_id=txn.id, ord=ord_, category_id=category_id))
            await session.commit()
    return txn.id


async def make_cbr_rate(rate_date: date, ccy: str, value: Amount, *, nominal: int = 1) -> None:
    async with get_sessionmaker()() as session:
        session.add(
            RawCbrRate(
                rate_date=rate_date, ccy=ccy, nominal=nominal, value=money(value) or Decimal(0)
            )
        )
        await session.commit()


async def make_rule(
    *,
    kind: RuleKind,
    match_type: RuleMatchType,
    pattern: str,
    value: str | None = None,
    enabled: bool = True,
    priority: int = 100,
) -> int:
    rule = await _add(
        Rule(
            kind=kind,
            match_type=match_type,
            pattern=pattern,
            value=value,
            enabled=enabled,
            priority=priority,
        )
    )
    return rule.id


async def make_trip(
    name: str, date_from: date, date_to: date, *, country: str | None = None
) -> int:
    trip = await _add(Trip(name=name, date_from=date_from, date_to=date_to, country=country))
    return trip.id


async def refresh(trigger: str = "test") -> Any:
    """Run the whole metric refresh the way the worker does."""
    from fintracker.metrics.refresh import refresh_all

    async with get_sessionmaker()() as session:
        entry = await refresh_all(session, trigger=trigger)
        assert entry.error is None, entry.error
        return entry


async def make_instrument(
    *,
    ticker: str = "GAZP",
    name: str | None = None,
    asset_class: AssetClass = AssetClass.share,
    currency: str = "RUB",
    board: str | None = "TQBR",
) -> int:
    instrument = await _add(
        Instrument(
            asset_class=asset_class,
            ticker=ticker,
            board=board,
            name=name or ticker,
            currency=currency,
        )
    )
    return instrument.id


async def make_event(
    d: date,
    *,
    account_id: int,
    kind: EventKind,
    instrument_id: int | None = None,
    quantity: Amount | None = None,
    price: Amount | None = None,
    amount: Amount = 0,
    currency: str = "RUB",
    fee: Amount | None = None,
    accrued_interest: Amount | None = None,
    status: EventStatus = EventStatus.confirmed,
    meta: dict[str, Any] | None = None,
    source_id: str | None = None,
) -> int:
    key = source_id or f"ev-{next(_seq)}"
    event = await _add(
        Event(
            account_id=account_id,
            instrument_id=instrument_id,
            kind=kind,
            status=status,
            ts=datetime.combine(d, datetime.min.time(), tzinfo=UTC),
            trade_date=d,
            quantity=money(quantity),
            price=money(price),
            price_currency=currency if price is not None else None,
            amount=money(amount) or Decimal(0),
            currency=currency,
            fee=money(fee),
            fee_currency=currency if fee is not None else None,
            accrued_interest=money(accrued_interest),
            source="tinvest",
            source_id=key,
            dedupe_key=f"tinvest:{key}",
            meta=meta,
        )
    )
    return event.id


async def make_price(
    d: date,
    *,
    instrument_id: int,
    close: Amount,
    currency: str = "RUB",
    accrued_interest: Amount | None = None,
) -> None:
    async with get_sessionmaker()() as session:
        session.add(
            PriceDaily(
                instrument_id=instrument_id,
                d=d,
                close=money(close) or Decimal(0),
                currency=currency,
                source="moex",
                accrued_interest=money(accrued_interest),
            )
        )
        await session.commit()