catnab/tests/test_main.py

import datetime
import json
import pickle
from types import SimpleNamespace

from amazonorders.entity.order import Order
from amazonorders.entity.item import Item
from amazonorders.entity.shipment import Shipment
from amazonorders.entity.transaction import Transaction
from amazonorders.entity.recipient import Recipient
import pytest

import main


KEYS = {
    Item: [
        'title', 'link', 'price', 'seller', 'condition',
        'return_eligible_date', 'image_link', 'quantity'
    ],
    Order: [
        'full_details', 'index', 'shipments', 'items', 'order_number', 'order_details_link',
        'grand_total', 'order_placed_date', 'recipient', 'payment_method', 'payment_method_last_4',
        'subtotal', 'shipping_total', 'free_shipping', 'promotion_applied', 'coupon_savings',
        'reward_points', 'subscription_discount', 'total_before_tax', 'estimated_tax',
        'refund_total', 'multibuy_discount', 'amazon_discount', 'gift_card', 'gift_wrap',
    ],
    Transaction: [
        'completed_date', 'payment_method', 'grand_total', 'is_refund',
        'order_number', 'order_details_link', 'seller',
    ],
    Shipment: [
        'items', 'delivery_status', 'tracking_link',
    ],
    Recipient: [
        'name', 'address',
    ],
}


def to_dict(entity):
    entity_keys = KEYS[type(entity)]
    result = {}
    for k, v in entity.__dict__.items():
        if k in entity_keys:
            if type(v) in KEYS:
                result[k] = to_dict(v)
            elif type(v) == list and len(v) > 0 and type(v[0]) in KEYS:
                result[k] = [to_dict(i) for i in v]
            elif type(v) in (datetime.date, datetime.datetime):
                result[k] = v.isoformat()
            else:
                result[k] = v
    return result


# @pytest.fixture
def mock_amz_hist():
    with open('output/transactions_2025.pkl', 'rb') as f:
        return pickle.load(f)


# @pytest.fixture
def mock_ynab_hist():
    with open('output/ynab_history.json') as f:
        return json.load(f)


def test_get_tx_categories(monkeypatch):
    order = SimpleNamespace(
        items=[
            SimpleNamespace(
                title='An example item',
                price=17.49
            )
        ]
    )

    item_details={
        'An example item': SimpleNamespace(
            category='Example category',
            description='example',
        )
    }

    def classify(_names, _categories):
        return item_details
    monkeypatch.setattr('ai.classify', classify)

    res = main.get_tx_categories(order, item_details)
    assert len(res) == 1
    assert res[0].category_name == 'Example category'
    assert res[0].description == 'example'
    assert res[0].ratio == 1.0


def test_get_categories_multi_items(monkeypatch):
    order = SimpleNamespace(
        items=[
            SimpleNamespace(
                title='An example item',
                price=17.49
            ),
            SimpleNamespace(
                title='Another example',
                price=12.34,
            )
        ]
    )

    item_details={
        'An example item': SimpleNamespace(
            category='Example category',
            description='example',
        ),
        'Another example': SimpleNamespace(
            category='Example category',
            description='second example',
        )
    }

    def classify(_names, _categories):
            return item_details
    monkeypatch.setattr('ai.classify', classify)

    res = main.get_tx_categories(order, item_details)
    assert len(res) == 1
    assert res[0].category_name == 'Example category'
    assert res[0].description == 'example, second example'
    assert res[0].ratio == 1.0


def test_get_categories_multi_categories(monkeypatch):
    order = SimpleNamespace(
        items=[
            SimpleNamespace(
                title='An example item',
                price=17.49
            ),
            SimpleNamespace(
                title='Another example',
                price=12.34,
            )
        ]
    )

    item_details={
        'An example item': SimpleNamespace(
            category='Example category',
            description='example',
        ),
        'Another example': SimpleNamespace(
            category='Other category',
            description='second example',
        )
    }

    def classify(_names, _categories):
            return item_details
    monkeypatch.setattr('ai.classify', classify)

    res = main.get_tx_categories(order, item_details)
    assert len(res) == 2
    assert res[0].category_name == 'Example category'
    assert res[0].description == 'example'
    assert res[0].ratio == pytest.approx(0.59, abs=0.01)

    assert res[1].category_name == 'Other category'
    assert res[1].description == 'second example'
    assert res[1].ratio == pytest.approx(0.41, abs=0.01)


def test_build_tx_update_multi():
    tx_categories = [
        SimpleNamespace(
            category_name='Example category',
            description='example',
            ratio=0.68,
        ),
        SimpleNamespace(
            category_name='Other category',
            description='second example',
            ratio=0.32,
        )
    ]

    category_ids = {
        'Example category': '123',
        'Other category': '456',
    }

    # ynab uses tenths of a cent as its unit
    tx_total = 23760

    res = main.build_tx_update(tx_categories, category_ids, tx_total)
    assert len(res.subtransactions) == 2

    sub0 = res.subtransactions[0]
    assert sub0.category_id == '123'
    assert sub0.memo == 'example'
    assert sub0.amount == 16157

    sub1 = res.subtransactions[1]
    assert sub1.category_id == '456'
    assert sub1.memo == 'second example'
    assert sub1.amount == 7603