move everything into src/catnab and add more tests

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

src/catnab/__init__.py

@@ -10,10 +10,11 @@ import click
 from ynab.models.transaction_detail import TransactionDetail
 import ynab
 
-import amazon
+from . import amazon
-from amazon import AmazonAccount
+from .amazon import AmazonAccount
-import ai
+from . import ai
-import ynab_client
+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:

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

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
+from catnab.ai import ItemClassification
-import main
+import catnab
-from main import TxCategory
+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)
+    res = catnab.build_tx_update(tx_categories, category_ids, tx)
-    assert len(res.subtransactions) == 2
+    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'}