156 lines
4.0 KiB
Python
156 lines
4.0 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from ai import ItemClassification
|
|
import main
|
|
from main import TxCategory
|
|
|
|
|
|
def test_get_tx_categories(monkeypatch):
|
|
order = MagicMock(
|
|
items=[
|
|
MagicMock(
|
|
title='An example item',
|
|
price=17.49
|
|
)
|
|
]
|
|
)
|
|
|
|
item_details={
|
|
'An example item': ItemClassification(
|
|
name='Example item',
|
|
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 = MagicMock(
|
|
items=[
|
|
MagicMock(
|
|
title='An example item',
|
|
price=17.49
|
|
),
|
|
MagicMock(
|
|
title='Another example',
|
|
price=12.34,
|
|
)
|
|
]
|
|
)
|
|
|
|
item_details={
|
|
'An example item': ItemClassification(
|
|
name='Example item 1',
|
|
category='Example category',
|
|
description='example',
|
|
),
|
|
'Another example': ItemClassification(
|
|
name='Example item 2',
|
|
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 = MagicMock(
|
|
items=[
|
|
MagicMock(
|
|
title='An example item',
|
|
price=17.49
|
|
),
|
|
MagicMock(
|
|
title='Another example',
|
|
price=12.34,
|
|
)
|
|
]
|
|
)
|
|
|
|
item_details={
|
|
'An example item': ItemClassification(
|
|
name='Example item 1',
|
|
category='Example category',
|
|
description='example',
|
|
),
|
|
'Another example': ItemClassification(
|
|
name='Example item 2',
|
|
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 = [
|
|
TxCategory(
|
|
category_name='Example category',
|
|
description='example',
|
|
ratio=0.68,
|
|
),
|
|
TxCategory(
|
|
category_name='Other category',
|
|
description='second example',
|
|
ratio=0.32,
|
|
)
|
|
]
|
|
|
|
category_ids = {
|
|
'Example category': '123',
|
|
'Other category': '456',
|
|
}
|
|
|
|
tx = MagicMock(
|
|
id='958e13e6-7bfd-465e-956a-887f15c8c456',
|
|
# ynab uses tenths of a cent as its unit
|
|
amount=23760,
|
|
)
|
|
|
|
res = main.build_tx_update(tx_categories, category_ids, tx)
|
|
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
|