From eccf81a5ed4fb892841b6fd5ac9ff07d5043cc95 Mon Sep 17 00:00:00 2001 From: Mathias Schneider Date: Thu, 29 Oct 2026 14:43:18 +0100 Subject: [PATCH] more tests and bugfixes --- src/models/category_tree.py | 79 ++++++++++++++++--------------------- tests/test_category_tree.py | 73 +++++++++++++++++++++++++++++++++- 2 files changed, 105 insertions(+), 47 deletions(-) diff --git a/src/models/category_tree.py b/src/models/category_tree.py index 048c4ea..c64db81 100644 --- a/src/models/category_tree.py +++ b/src/models/category_tree.py @@ -1,6 +1,6 @@ from datetime import date -from pydantic import BaseModel +from pydantic import BaseModel, computed_field from src.models.category import Category from src.models.transactions import Transaction @@ -8,32 +8,32 @@ from src.models.transactions import Transaction class CategoryTree(BaseModel): category: Category | None - total_amount: float | None = None - total_budget: float | None = None sub_categories: list[CategoryTree] | None = None transactions: list[Transaction] | None = None - def get_total_amount(self) -> float: + @computed_field # type: ignore[prop-decorator] + @property + def total_amount(self) -> float: result = 0.0 if self.transactions: - sum(transaction.amount for transaction in self.transactions) + result += sum( + float(transaction.amount) for transaction in self.transactions + ) if self.sub_categories: - sum(sub_category.get_total_amount() for sub_category in self.sub_categories) + result += sum( + float(sub_category.total_amount) for sub_category in self.sub_categories + ) return result - def get_total_budget(self, date_valid: date) -> float: + def total_budget(self, date_valid: date) -> float: result = 0.0 if self.category is None: return result result += self.category.get_budget(date_valid) if self.sub_categories: result += sum( - ( - category.category.get_budget(date_valid) - if category.category is not None - else 0.0 - ) - for category in self.sub_categories + sub_category.total_budget(date_valid) + for sub_category in self.sub_categories ) return result @@ -41,48 +41,37 @@ class CategoryTree(BaseModel): if self.category != other.category: return - total_amount = ( - (self.total_amount or 0.0) + (other.total_amount or 0.0) - ) or None - total_budget = ( - (self.total_budget or 0.0) + (other.total_budget or 0.0) - ) or None - - sub_categories = [] - for category1 in self.sub_categories or []: - for category2 in other.sub_categories or []: - if category1.category == category2.category: - sub_categories.append(category1 + category2) - else: - sub_categories.append(category1) - sub_categories.append(category2) + sub_trees = [] + seen = [] + for sub_tree_1 in self.sub_categories or []: + new = sub_tree_1 + for sub_tree_2 in other.sub_categories or []: + if sub_tree_1.category == sub_tree_2.category: + new += sub_tree_2 + seen.append(sub_tree_2) + break + sub_trees.append(new) + sub_trees += [ + sub_tree + for sub_tree in (other.sub_categories or []) + if sub_tree not in seen + ] transactions = ((self.transactions or []) + (other.transactions or [])) or None return CategoryTree( category=self.category, - total_amount=total_amount, - total_budget=total_budget, - sub_categories=sub_categories, + sub_categories=list(sub_trees), transactions=transactions, ) def build_category_tree(transactions: list[Transaction]) -> list[CategoryTree]: - category_branches = map(_transaction_to_category_branch, transactions) - results: list = [] - for category_branch in category_branches: - if not results: - results += [category_branch] - continue - if category_branch.category not in [result.category for result in results]: - results += [category_branch] - else: - for result in results: - if (category_branch.category or "") == (result.category or ""): - result += category_branch - break - return results + results: dict[str, CategoryTree] = {} + for branch in map(_transaction_to_category_branch, transactions): + key = str(branch.category) if branch.category else "" + results[key] = (results[key] + branch) if key in results else branch + return list(results.values()) def _transaction_to_category_branch(transaction: Transaction) -> CategoryTree: diff --git a/tests/test_category_tree.py b/tests/test_category_tree.py index 0d4fff4..e16c9f1 100644 --- a/tests/test_category_tree.py +++ b/tests/test_category_tree.py @@ -1,5 +1,9 @@ from src.models.category import Category -from src.models.category_tree import CategoryTree, _transaction_to_category_branch +from src.models.category_tree import ( + CategoryTree, + _transaction_to_category_branch, + build_category_tree, +) from src.models.transactions import Transaction from tests.factories import TransactionFactory @@ -14,7 +18,6 @@ def test_transaction_to_category_branch_simple(): assert isinstance(result.transactions, list) assert len(result.transactions) == 1 assert isinstance(result.transactions[0], Transaction) - assert result.transactions[0] == transaction def test_transaction_to_category_branch_two_category_levels(): @@ -45,3 +48,69 @@ def test_transaction_to_category_branch_no_category(): assert isinstance(result.transactions, list) assert len(result.transactions) == 1 assert result.transactions[0] == transaction + + +def test_category_tree_add(): + category_tree_1 = CategoryTree( + category=Category.from_str("super"), + sub_categories=[ + CategoryTree( + category=Category.from_str("super:sub1"), + transactions=[TransactionFactory.build(amount=10)], + ), + CategoryTree( + category=Category.from_str("super:sub2"), + transactions=[TransactionFactory.build(amount=30)], + ), + ], + ) + category_tree_2 = CategoryTree( + category=Category.from_str("super"), + sub_categories=[ + CategoryTree( + category=Category.from_str("super:sub2"), + transactions=[TransactionFactory.build(amount=20)], + ), + CategoryTree( + category=Category.from_str("super:sub3"), + transactions=[TransactionFactory.build(amount=100)], + ), + ], + ) + result = category_tree_1 + category_tree_2 + assert result + assert result.sub_categories + assert isinstance(result.sub_categories, list) + assert len(result.sub_categories) == 3 + assert all(s.category for s in result.sub_categories) + assert sorted([s.category.name for s in result.sub_categories]) == [ + "sub1", + "sub2", + "sub3", + ] + assert result.total_amount == 160.0 + + +def test_build_category_tree(): + transactions = [ + TransactionFactory.build(category=Category.from_str("super1:sub1"), amount=10), + TransactionFactory.build(category=Category.from_str("super1:sub2"), amount=13), + TransactionFactory.build( + category=Category.from_str("super1:sub2:sub2_sub1"), amount=104 + ), + TransactionFactory.build( + category=Category.from_str("super1:sub2:sub2_sub1"), amount=109 + ), + TransactionFactory.build(category=Category.from_str("super2"), amount=5000), + ] + result = build_category_tree(transactions) + assert result + assert isinstance(result, list) + assert len(result) == 2 + assert isinstance(result[0], CategoryTree) + assert result[0].category + assert result[0].category.name == "super1" + assert result[0].transactions is None + assert result[0].sub_categories + assert isinstance(result[0].sub_categories, list) + assert len(result[0].sub_categories) == 2