move everything into src/catnab and add more tests

This commit is contained in:
2025-12-09 08:39:04 -05:00
parent 2509c4dd70
commit 761c7eeda4
7 changed files with 351 additions and 26 deletions

View File

@@ -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

View File

@@ -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:

189
tests/conftest.py Normal file
View File

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

View File

@@ -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'}