require 'test/unit'
require 'flatten'

# This test may be run in three modes:
#  1) With no arguments, tests the built-in Array method flatten replaced with Flatten module implementation
#  2) With a single argument 'newclass', tests a subclass of Array with Flatten module implementation
#  2) With a single argument 'original', tests the built-in Array flatten method
if ARGV[0] == 'newclass'
  class MyArray < Array
    include Flatten
  end
  CLASS = MyArray
elsif ARGV[0] == 'original'
  CLASS = Array
else
  class Array
    include Flatten
  end
  CLASS = Array
end

puts "Working with array like class: #{CLASS.name}"

class FlattenTest < Test::Unit::TestCase

  def test_one_level
    assert_equal [1,2,3], CLASS.new([1,[2,3]]).flatten
  end

  def test_flatten_keeps_class
    my_array = CLASS.new([1,[2,3],4])
    flattened = my_array.flatten
    assert_equal my_array.class, flattened.class
    assert_equal CLASS.name, flattened.class.name
    assert_equal flattened, [1,2,3,4]
  end

  def test_flatten_many_levels
    assert_equal [0,0,0,1,2,3,4,5,6,7,8,9,10,11,12], CLASS.new([0,0,0,[1,[2,3],4],5,[[[[[[6,7]],8]],9],10,11],12]).flatten
  end

  def test_flatten_self_referencing_fails
    a = CLASS.new([1,2])
    b = CLASS.new([3,4,a])
    a << b
    assert_raises(ArgumentError) { a.flatten }
  end

end