more tests and bugfixes

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:

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