diff --git a/pyproject.toml b/pyproject.toml index 09cd3f9..67c75ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ ] [project.scripts] -catnab = "main:main" +catnab = "catnab:main" [dependency-groups] dev = [ @@ -25,8 +25,5 @@ dev = [ "pytest>=8.4.2", ] -[tool.pytest.ini_options] -pythonpath = ["src"] - [tool.uv] package = true diff --git a/src/main.py b/src/catnab/__init__.py similarity index 93% rename from src/main.py rename to src/catnab/__init__.py index b242cdf..a7a72c2 100644 --- a/src/main.py +++ b/src/catnab/__init__.py @@ -10,10 +10,11 @@ import click from ynab.models.transaction_detail import TransactionDetail import ynab -import amazon -from amazon import AmazonAccount -import ai -import ynab_client +from . import amazon +from .amazon import AmazonAccount +from . import ai +from .ai import ItemClassification +from . import ynab_client @dataclass @@ -81,9 +82,9 @@ def tx_matches(amz_hist: list[Transaction], ynab_hist: list[TransactionDetail]) matched = [] for amz_tx in filter(is_matchable_amz, amz_hist): # ynab uses tenths of a cent for whatever reason - total_cents = round(amz_tx.grand_total * 1000) + amount = round(amz_tx.grand_total * 1000) found = None - for ynab_tx in ynab_by_amt[total_cents]: + for ynab_tx in ynab_by_amt[amount]: # valid transactions have the same amount, and occur within 7 days of each other. # Ties are broken by date, with the nearest taking precedence. date_diff = tx_date_diff(amz_tx, ynab_tx) @@ -102,11 +103,11 @@ def tx_matches(amz_hist: list[Transaction], ynab_hist: list[TransactionDetail]) def get_tx_categories( order: Order, - item_details: dict[str, ai.ItemClassification], + classified: dict[str, ItemClassification], ) -> list[TxCategory]: by_category = collections.defaultdict(list) for item in order.items: - details = item_details[item.title] + details = classified[item.title] by_category[details.category].append(item) # we have already ensured that order details are fully populated, so i.price should not be None @@ -117,7 +118,7 @@ def get_tx_categories( item_descriptions = [] for i in items: category_total += i.price - details = item_details[i.title] + details = classified[i.title] item_descriptions.append(details.description) txc = TxCategory( @@ -192,7 +193,7 @@ def sync_account(ctx: Context, acct: AmazonAccount): print('Classifying order items...') item_names = [i.title for o in amz_orders.values() for i in o.items] - item_details = ai.classify(item_names, ctx.ynab_client.text_categories) + classified = ai.classify(item_names, ctx.ynab_client.text_categories) updates = [] print('Building transaction updates...') @@ -200,7 +201,7 @@ def sync_account(ctx: Context, acct: AmazonAccount): try: print(f' ${abs(amz_tx.grand_total)} on {amz_tx.completed_date}') order = amz_orders[amz_tx.order_number] - tx_categories = get_tx_categories(order, item_details) + tx_categories = get_tx_categories(order, classified) tx_update = build_tx_update(tx_categories, ctx.ynab_client.category_ids, ynab_tx) updates.append(tx_update) except Exception: diff --git a/src/ai.py b/src/catnab/ai.py similarity index 100% rename from src/ai.py rename to src/catnab/ai.py diff --git a/src/amazon.py b/src/catnab/amazon.py similarity index 100% rename from src/amazon.py rename to src/catnab/amazon.py diff --git a/src/ynab_client.py b/src/catnab/ynab_client.py similarity index 100% rename from src/ynab_client.py rename to src/catnab/ynab_client.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..743ded9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,189 @@ +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') diff --git a/tests/test_main.py b/tests/test_main.py index 6e82c5a..bbbf3bb 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,10 +1,102 @@ +from datetime import date from unittest.mock import MagicMock import pytest +from ynab import TransactionFlagColor -from ai import ItemClassification -import main -from main import TxCategory +from catnab.ai import ItemClassification +import catnab +from catnab import Context, TxCategory +from catnab.amazon import AmazonAccount + + +def test_is_matchable_ynab(mk_ynab_tx): + matchable = mk_ynab_tx() + assert catnab.is_matchable_ynab(matchable) + + wrong_payee = mk_ynab_tx(payee_name='Walmart') + assert not catnab.is_matchable_ynab(wrong_payee) + + already_matched = mk_ynab_tx(flag_color=TransactionFlagColor.PURPLE) + assert not catnab.is_matchable_ynab(already_matched) + + already_approved = mk_ynab_tx(approved=True) + assert not catnab.is_matchable_ynab(already_approved) + + # fortunately we only need to check the length of the subtransactions list + already_split = mk_ynab_tx(subtransactions=[1, 2, 3]) + assert not catnab.is_matchable_ynab(already_split) + + +def test_is_matchable_amz(mk_amz_tx): + matchable = mk_amz_tx() + assert catnab.is_matchable_amz(matchable) + + refund = mk_amz_tx(is_refund=True) + assert not catnab.is_matchable_amz(refund) + + empty_order_number = mk_amz_tx(order_number='') + assert not catnab.is_matchable_amz(empty_order_number) + + gift_card_order=mk_amz_tx(payment_method='Amazon Gift Card') + assert not catnab.is_matchable_amz(gift_card_order) + + points_order = mk_amz_tx(payment_method='Amazon Visa points') + assert not catnab.is_matchable_amz(points_order) + + +def test_tx_matches(mk_amz_tx, mk_ynab_tx): + amz_tx = mk_amz_tx( + grand_total=123.45, + completed_date=date(2025, 1, 1), + ) + ynab_tx=mk_ynab_tx( + amount=123450, + var_date=date(2025, 1, 1), + ) + + matches = catnab.tx_matches([amz_tx], [ynab_tx]) + assert len(matches) == 1 + + +def test_tx_matches_max_datediff(mk_amz_tx, mk_ynab_tx): + # with a 7-day difference, transactions match + amz_tx = mk_amz_tx( + grand_total=123.45, + completed_date=date(2025, 1, 1), + ) + ynab_tx=mk_ynab_tx( + amount=123450, + var_date=date(2025, 1, 8), + ) + + matches = catnab.tx_matches([amz_tx], [ynab_tx]) + assert len(matches) == 1 + + # but with an 8-day difference, they don't + ynab_tx.var_date = date(2026, 1, 9) + matches = catnab.tx_matches([amz_tx], [ynab_tx]) + assert len(matches) == 0 + + +def test_tx_matches_tiebreak(mk_amz_tx, mk_ynab_tx): + amz_tx = mk_amz_tx( + grand_total=123.45, + completed_date=date(2025, 1, 1), + ) + ynab_1 = mk_ynab_tx( + amount=123450, + var_date=date(2025, 1, 5), + ) + ynab_2 = mk_ynab_tx( + amount=123450, + var_date=date(2025, 1, 3), + ) + + matches = catnab.tx_matches([amz_tx], [ynab_1, ynab_2]) + assert len(matches) == 1 + matched_amz, matched_ynab = matches[0] + assert matched_ynab is ynab_2 def test_get_tx_categories(monkeypatch): @@ -27,9 +119,9 @@ def test_get_tx_categories(monkeypatch): def classify(_names, _categories): return item_details - monkeypatch.setattr('ai.classify', classify) + monkeypatch.setattr('catnab.ai.classify', classify) - res = main.get_tx_categories(order, item_details) + res = catnab.get_tx_categories(order, item_details) assert len(res) == 1 assert res[0].category_name == 'Example category' assert res[0].description == 'example' @@ -65,9 +157,9 @@ def test_get_categories_multi_items(monkeypatch): def classify(_names, _categories): return item_details - monkeypatch.setattr('ai.classify', classify) + monkeypatch.setattr('catnab.ai.classify', classify) - res = main.get_tx_categories(order, item_details) + res = catnab.get_tx_categories(order, item_details) assert len(res) == 1 assert res[0].category_name == 'Example category' assert res[0].description == 'example, second example' @@ -103,9 +195,9 @@ def test_get_categories_multi_categories(monkeypatch): def classify(_names, _categories): return item_details - monkeypatch.setattr('ai.classify', classify) + monkeypatch.setattr('catnab.ai.classify', classify) - res = main.get_tx_categories(order, item_details) + res = catnab.get_tx_categories(order, item_details) assert len(res) == 2 assert res[0].category_name == 'Example category' assert res[0].description == 'example' @@ -141,8 +233,8 @@ def test_build_tx_update_multi(): amount=23760, ) - res = main.build_tx_update(tx_categories, category_ids, tx) - assert len(res.subtransactions) == 2 + res = catnab.build_tx_update(tx_categories, category_ids, tx) + assert res.subtransactions is not None and len(res.subtransactions) == 2 sub0 = res.subtransactions[0] assert sub0.category_id == '123' @@ -153,3 +245,49 @@ def test_build_tx_update_multi(): assert sub1.category_id == '456' assert sub1.memo == 'second example' assert sub1.amount == 7603 + + +def test_sync_acct_basic(basic_order, ynab_creds): + ctx = Context(sync_days=30) + acct = AmazonAccount('test@example.com', 'password123') + catnab.sync_account(ctx, acct) + + mock_update = ctx.ynab_client.update_transactions + mock_update.assert_called_once() + + tx_update = mock_update.call_args.args[0][0] + assert tx_update.id == '1' + assert tx_update.category_id == '0' + assert tx_update.flag_color == 'purple' + + +def test_sync_acct_multi_items(multi_item_order, ynab_creds): + ctx = Context(sync_days=30) + acct = AmazonAccount('test@example.com', 'password123') + catnab.sync_account(ctx, acct) + + mock_update = ctx.ynab_client.update_transactions + mock_update.assert_called_once() + + tx_update = mock_update.call_args.args[0][0] + assert tx_update.id == '1' + assert tx_update.category_id == '0' + assert tx_update.flag_color == 'purple' + + +def test_sync_acct_multi_categories(multi_category_order, ynab_creds): + ctx = Context(sync_days=30) + acct = AmazonAccount('test@example.com', 'password123') + catnab.sync_account(ctx, acct) + + mock_update = ctx.ynab_client.update_transactions + mock_update.assert_called_once() + + tx_update = mock_update.call_args.args[0][0] + assert tx_update.id == '1' + assert tx_update.category_id == None + assert tx_update.flag_color == 'purple' + + assert len(tx_update.subtransactions) == 2 + # we don't care about order, but both categories need to be present + assert set(st.category_id for st in tx_update.subtransactions) == {'0', '1'}