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')