more tests and bugfixes

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