Custom flatten implementation

my_flatten.rb

require "test/unit"

class Array

  ##
  # Creates a new array, taking all elements into one or more dimensions, depending on the argument.
  # If no argument is set, the array dimension will be one.
  #
  # +examples:+
  #
  #   [[1], 2, [[[3]]], 4].my_flatten()
  #   #=> [1, 2, 3, 4]
  #   [[1], 2, [[3]], 4].my_flatten(1)
  #   #=> [1, 2, [3], 4]
  #
  def my_flatten(n = nil)
    n ? multiple_flatten(self, n) : recursive_flatten(self)
  end

  private  

  def recursive_flatten(array, results = [])
    array.each do |element|
      if element.class == Array
        recursive_flatten(element, results)
      else
        results << element
      end
    end
    results
  end  
  def multiple_flatten(array, n)
    count = 0
    arr = array
    while count < n do
      arr = single_flatten(arr)
      count += 1
    end
    arr
  end  

  def single_flatten(array)
    results = []
    array.each do |element|
      if element.class == Array
        element.each {|value| results << value}
      else
        results << element
      end
    end
    results
  end
end

class TestFlatten < Test::Unit::TestCase
  def test_simple_flatten
    assert_equal([1, 2, 3, 4], [[1], 2, [[[3]]], 4].my_flatten())
  end

  def test_flatten_with_arguments_one
    assert_equal([1, 2, [3], 4], [[1], 2, [[3]], 4].my_flatten(1))
  end

  def test_flatten_with_arguments_two
    assert_equal([1, 2, 3, [7], 4], [[1], 2, [[3, [7]]], 4].my_flatten(2))
  end

  def test_flatten_with_empty_array_and_arguments
    assert_equal([], [].my_flatten(1))
  end

  def test_flatten_with_empty_array_and_no_arguments
    assert_equal([], [].my_flatten())
  end
end