catnab/tests/conftest.py

from datetime import date
from unittest.mock import MagicMock

import pytest

from catnab.ai import ItemClassification


@pytest.fixture
def mk_ynab_tx():
    def ynab_tx(
        id='1',
        payee_name='Amazon',
        amount=123450,
        var_date=date(2025, 1, 1),
        flag_color=None,
        approved=False,
        subtransactions=[],
    ):
        return MagicMock(
            id=id,
            payee_name=payee_name,
            amount=amount,
            var_date=var_date,
            flag_color=flag_color,
            approved=approved,
            subtransactions=subtransactions
        )

    return ynab_tx


@pytest.fixture
def mk_amz_tx():
    def amz_tx(
        grand_total=123.45,
        completed_date=date(2025, 1, 1),
        order_number='A',
        is_refund=False,
        payment_method='Visa 1234',
    ):
        return MagicMock(
            grand_total=grand_total,
            completed_date=completed_date,
            order_number=order_number,
            is_refund=is_refund,
            payment_method=payment_method,
        )

    return amz_tx


def order_item(
    title='An example item',
    price=120.34,
    quantity=1,
):
    return MagicMock(
        title=title,
        price=price,
        quantity=quantity,
    )


def amz_order(
    order_number='A',
    order_placed_date=date(2025, 1, 1),
    items=[order_item()],
    grand_total=123.45, # add a little bit for pretend tax
    full_details=True,
):
    return MagicMock(
        order_number=order_number,
        order_placed_date=order_placed_date,
        items=items,
        grand_total=grand_total,
        full_details=full_details,
    )


@pytest.fixture
def setup_mock_order(monkeypatch):
    def setup(orders, amz_hist, ynab_hist, item_categories):
        def __enter__(self): return self
        monkeypatch.setattr('catnab.amazon.Session.__enter__', __enter__)

        def list_transactions(self, days): return amz_hist
        monkeypatch.setattr('catnab.amazon.Session.list_transactions', list_transactions)

        orders_map = {o.order_number: o for o in orders}
        def get_order(self, order_number): return orders_map[order_number]
        monkeypatch.setattr('catnab.amazon.Session.get_order', get_order)

        monkeypatch.setattr('catnab.ynab_client.YnabClient.history', ynab_hist)
        monkeypatch.setattr('catnab.ynab_client.YnabClient.update_transactions', MagicMock())

        # item_categories is mapping of item title -> category
        category_names = set(item_categories.values())
        category_ids = {cat: str(i) for i, cat in enumerate(category_names)}
        monkeypatch.setattr('catnab.ynab_client.YnabClient.category_ids', category_ids)

        def classify(item_names, _categories):
            result = {}
            for n in item_names:
                result[n] = ItemClassification(
                    name=n,
                    category=item_categories[n],
                    description='',
                )
            return result
        monkeypatch.setattr('catnab.ai.classify', classify)

    return setup


@pytest.fixture
def basic_order(mk_amz_tx, mk_ynab_tx, setup_mock_order):
    items = [order_item(title='test', price=120.34, quantity=1)]
    item_categories = {'test': 'example category'}

    order = amz_order(
        order_number='A',
        items=items,
        grand_total=123.45,
    )
    amz_tx = mk_amz_tx(order_number='A', grand_total=123.45)
    ynab_tx = mk_ynab_tx(id='1', amount=123450)

    setup_mock_order(
        orders=[order],
        amz_hist=[amz_tx],
        ynab_hist=[ynab_tx],
        item_categories=item_categories,
    )


@pytest.fixture
def multi_item_order(mk_amz_tx, mk_ynab_tx, setup_mock_order):
    items = [
        order_item(title='item 1', price=59.36, quantity=1),
        order_item(title='item 2', price=22.18, quantity=1),
    ]
    item_categories={'item 1': 'example category', 'item 2': 'example category'}

    order = amz_order(
        order_number='A',
        items=items,
        grand_total=81.54,
    )
    amz_tx = mk_amz_tx(order_number='A', grand_total=81.54)
    ynab_tx = mk_ynab_tx(id='1', amount=81540)

    setup_mock_order(
        orders=[order],
        amz_hist=[amz_tx],
        ynab_hist=[ynab_tx],
        item_categories=item_categories
    )


@pytest.fixture
def multi_category_order(mk_amz_tx, mk_ynab_tx, setup_mock_order):
    items = [
        order_item(title='item 1', price=59.36, quantity=1),
        order_item(title='item 2', price=22.18, quantity=1),
    ]
    # first category will have id 0, second category will have id 1
    item_categories={'item 1': 'example category', 'item 2': 'another category'}

    order = amz_order(
        order_number='A',
        items=items,
        grand_total=81.54,
    )
    amz_tx = mk_amz_tx(order_number='A', grand_total=81.54)
    ynab_tx = mk_ynab_tx(id='1', amount=81540)

    setup_mock_order(
        orders=[order],
        amz_hist=[amz_tx],
        ynab_hist=[ynab_tx],
        item_categories=item_categories
    )


@pytest.fixture
def ynab_creds(monkeypatch):
    monkeypatch.setenv('YNAB_BUDGET_ID', '00000000-0000-0000-0000-000000000000')
    monkeypatch.setenv('YNAB_TOKEN', 'very_secure_token')