From 6061b70635051b6b2438340a68f21f76cd5b3afc Mon Sep 17 00:00:00 2001
From: Ira Fesefeldt <fesefeldt@cs.rwth-aachen.de>
Date: Wed, 15 Nov 2023 16:59:33 +0100
Subject: [PATCH] fixed error and added test method

---
 pyalgotask/tasks/trees/avl_tree.py           | 11 ++++++++---
 pyalgotask/tasks/trees/binary_search_tree.py |  2 +-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/pyalgotask/tasks/trees/avl_tree.py b/pyalgotask/tasks/trees/avl_tree.py
index 919c4a4..bc3e680 100644
--- a/pyalgotask/tasks/trees/avl_tree.py
+++ b/pyalgotask/tasks/trees/avl_tree.py
@@ -74,7 +74,8 @@ def insert(self, value):
         self.adjust_depth(new_node.parent)
 
         # rebalance (A7-A9)
-        self.balance(new_node)
+        root = self.balance(new_node)
+        self.algorithm_root = root
 
     def adjust_depth(self, node):
         """adjusts the depth values starting from node
@@ -140,8 +141,9 @@ def balance(self, node):
         :return: the new root of the tree
         """
         if node.is_nil():
-            return
+            return node
 
+        child = node
         current = node.parent
         while not current.is_nil():
             if current.left.depth == current.right.depth + 2:
@@ -151,8 +153,11 @@ def balance(self, node):
                 # right imbalance
                 current = self.balance_right(current)
 
+            child = current
             current = current.parent
 
+        return child
+
     def delete(self, value):
         """Performs a delete of the value from the tree
         :param value: the value to delete
@@ -209,7 +214,7 @@ def algorithm(self):
                 raise NotImplementedError(
                     f"Operator {operation.type} not supported by task {self.__class__.__name__}"
                 )
-            yield (self.algorithm_root, str(operation))
+            yield (self.algorithm_root, operation)
 
 
 task_base.register_task("tree", AVLTree())
diff --git a/pyalgotask/tasks/trees/binary_search_tree.py b/pyalgotask/tasks/trees/binary_search_tree.py
index e3abe2f..79ab1f3 100644
--- a/pyalgotask/tasks/trees/binary_search_tree.py
+++ b/pyalgotask/tasks/trees/binary_search_tree.py
@@ -146,7 +146,7 @@ def algorithm(self):
                 raise NotImplementedError(
                     f"Operator {operation.type} not supported by task {self.__class__.__name__}"
                 )
-            yield (self.algorithm_root, str(operation))
+            yield (self.algorithm_root, operation)
 
 
 task_base.register_task("tree", BinarySearchTree())