code: def test_sum_tree():
      # test power of 2 sum tree
      sum_tree = SegmentTree(
          size=4, operation=SegmentTree.Operation.SUM
      )
      sum_tree.add(10, "10")
      assert sum_tree.total_value() == 10
      sum_tree.add(20, "20")
      assert sum_tree.total_value() == 30
      sum_tree.add(5, "5")
      assert sum_tree.total_value() == 35
      sum_tree.add(7.5, "7.5")
      assert sum_tree.total_value() == 42.5
      sum_tree.add(2.5, "2.5")
      assert sum_tree.total_value() == 35
      sum_tree.add(5, "5")
      assert sum_tree.total_value() == 20
  
      assert sum_tree.get_element_by_partial_sum(2) == (0, 2.5, "2.5")
      assert sum_tree.get_element_by_partial_sum(3) == (1, 5.0, "5")
      assert sum_tree.get_element_by_partial_sum(10) == (2, 5.0, "5")
      assert sum_tree.get_element_by_partial_sum(13) == (3, 7.5, "7.5")
  
      sum_tree.update(2, 10)
      assert (
          sum_tree.__str__()
          == "[25.]\n[ 7.5 17.5]\n[ 2.5  5.  10.   7.5]\n"
      )
  
      # test non power of 2 sum tree
      with pytest.raises(ValueError):
          sum_tree = SegmentTree(
              size=5, operation=SegmentTree.Operation.SUM
          )
  
imports: import os
  import sys
  import pytest
  from rl_coach.memories.non_episodic.prioritized_experience_replay import SegmentTree
label_desc: Use image node_id 1 for calling a global function with example usage: test_sum_tree() without return types