Skip to content

Commit 085c3a4

Browse files
committed
Fix each_level enumerator
Return a level-wise enumerator when no block is given and add a test to assert the yielded level arrays match the tree structure.
1 parent bf547fe commit 085c3a4

2 files changed

Lines changed: 13 additions & 1 deletion

File tree

lib/tree.rb

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -782,7 +782,7 @@ def each_level
782782
end
783783
self
784784
else
785-
each
785+
to_enum(:each_level)
786786
end
787787
end
788788

test/test_tree.rb

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -846,6 +846,18 @@ def test_each_leaf
846846
assert(result_array.include?(@child4), 'Should have child 4')
847847
end
848848

849+
# Test the each_level method without a block (Enumerator).
850+
def test_each_level
851+
setup_test_tree
852+
853+
levels = @root.each_level.to_a
854+
855+
assert_equal(3, levels.length, 'Should have three levels')
856+
assert_equal([@root], levels[0])
857+
assert_equal([@child1, @child2, @child3], levels[1])
858+
assert_equal([@child4], levels[2])
859+
end
860+
849861
# Test the parent method.
850862
def test_parent
851863
setup_test_tree

0 commit comments

Comments
 (0)