require 'test/unit'
require 'mycode.rb'

class StringStream < String
  def write(message)
    self.<< message
  end
end

def wrap(&b)
  raise "Expected block!" unless block_given?
  s = StringStream.new
  old = $stdout.clone
  $stdout = s
  b.call
  $stdout = old
  s
end

class TC_MyCode < Test::Unit::TestCase
  def test_mycode
    assert_equal "First Option\n", wrap {foo(1)}, "First Option"
    assert_equal "Second Option\n", wrap {foo(2)}, "Second Option"
    assert_equal "Third Option\n", wrap {foo(3)}, "Third Option"
    assert_equal "Unknown Option\n", wrap {foo(4)}, "Unknown Option"
    assert_equal "Unknown Option\n", wrap {foo("hi!")}, "Unknown Option"
  end
end